#!/usr/bin/env python
# coding: utf-8


import json
import csv
import PredysUtil8ans as pu
import numpy as np
import os
from sklearn.svm import SVC
from sklearn.svm import LinearSVC
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score
import nibabel as nib
from nibabel import load
import math    
import csv
import glob
import random
import pickle
import copy
import scipy
from scipy import signal
import scipy.stats
from scipy import ndimage
from statsmodels.stats.multitest import fdrcorrection
    


def IntraModalDec(TaskName, MethodName='Test', TargTime=[0,1,2,3,4], bdelay=3, Rad=5, RunFlag=2, svmType=1):

#Input:
# MethodName ... output label
# TargTime ... Target time points used
# bdelay ... Delay of bold signal from onset
# Rad ... Radius of searchlight sphere
# RunFlag ... 1: two runs are used (if available) / 2: Only the first run is used
# svmType ... 1: SVC with RBF kernel / 2: LinearSVC 

    SubList_trn = [1,2,3,4,5,6,7,9,10,11,12,13,14,15,17,18,19,20,24,25,27,28,29,30,31,32,37,38,40,41,43,44,45,46,48,50,54,55,58,59,61,64,67,68,70,72] 
    SubList_test = [2,4,8,13,16,19,22,27,28,29,30,35,40,49,54,56,57,67,77,83,84,87,90, 99,101,105,106,107,108,110,111,112,113,114,117,119,121,122,123,124,127,130,132]

    CurrDir = os.getcwd()
    RoiDir = CurrDir + '/WFU_AALAtlas'
    SaveDir = CurrDir + '/Result_8ans/Result_SearchLight_8to5_' + TaskName    
    PreprocDataDir_trn = CurrDir + '/DetrendData_8ans'  
    PreprocDataDir_test = CurrDir + '/DetrendData_5ans'               

    RoiName = 'GrayMatter'
    MinVoxN = 10 #Minimum voxel numbers for a sphere
    VolumeSize = [79, 95, 39]
    TotalVoxelSize = np.prod(VolumeSize)
    if RunFlag == 1:
        UseRun = [1, 2]
    else:
        UseRun = [1]


    #Load Whole brain data
    RoiImg = load(os.path.join(RoiDir,RoiName + '.nii'))
    RoiVol = RoiImg.get_fdata()
    Roi_x, Roi_y, Roi_z = np.nonzero(RoiVol)
    RoiLen = len(Roi_x)

    #Indices of voxels included in Sphere ROIs    
    VoxelIndex = [np.unravel_index(ii,VolumeSize) for ii in range(TotalVoxelSize)]
    TargVoxel = [ii for ii in range(TotalVoxelSize) if RoiVol[VoxelIndex[ii]] > 0]
    VoxelIndex_targ = [VoxelIndex[TargVoxel[tt]] for tt in range(len(TargVoxel))]

    #Load preprocessed brain response data
    tRespVal01_trn = []
    tRespVal02_trn = []
    for SubNum in SubList_trn:
        SubName = 'sub-{:02d}'.format(SubNum)
        print('Loading ' + SubName +' ...')
        fname = 'RespData_bd{:d}_{}_NoAdapt_{}_Run{:02d}.npy'
        if (TaskName == 'Words') and (RunFlag == 2) and (SubNum == 70): #For this subject, only 2nd run exists
            T01_List = [pu.CalcTrialRespVal(PreprocDataDir_trn, fname.format(bdelay, TaskName, SubName, RunNum), TargTime) for RunNum in [2]]
        else:
            T01_List = [pu.CalcTrialRespVal(PreprocDataDir_trn, fname.format(bdelay, TaskName, SubName, RunNum), TargTime) for RunNum in UseRun]
        T01_List = [T01_List[tt] for tt in range(len(T01_List)) if not len(T01_List[tt])==0]
        T01 = np.concatenate(T01_List, axis=0)
        tRespVal01_trn.append(T01)

        fname = 'RespData_bd{:d}_{}_Adapt_{}_Run{:02d}.npy' 
        if (TaskName == 'Words') and (RunFlag == 2) and (SubNum == 70): #For this subject, only 2nd run exists
            T02_List = [pu.CalcTrialRespVal(PreprocDataDir_trn, fname.format(bdelay, TaskName, SubName, RunNum), TargTime) for RunNum in [2]]
        else:        
            T02_List = [pu.CalcTrialRespVal(PreprocDataDir_trn, fname.format(bdelay, TaskName, SubName, RunNum), TargTime) for RunNum in UseRun]
        T02_List = [T02_List[tt] for tt in range(len(T02_List)) if not len(T02_List[tt])==0]
        T02 = np.concatenate(T02_List, axis=0)
        tRespVal02_trn.append(T02)

    tRespVal01_test = []
    tRespVal02_test = []
    for SubNum in SubList_test:
        SubName = 'sub-{:03d}'.format(SubNum)
        print('Loading ' + SubName +' ...')
        fname = 'RespData_bd{:d}_{}_NoAdapt_{}.npy'         
        T01 = pu.CalcTrialRespVal(PreprocDataDir_test, fname.format(bdelay, TaskName, SubName), TargTime)
        tRespVal01_test.append(T01)

        fname = 'RespData_bd{:d}_{}_Adapt_{}.npy' 
        T02 = pu.CalcTrialRespVal(PreprocDataDir_test, fname.format(bdelay, TaskName, SubName), TargTime)
        tRespVal02_test.append(T02)

    #Load searchlight sphere data
    with open('SearchSphereSet_Rad{}.pickle'.format(Rad), mode='rb') as f:
        SphereSet = pickle.load(f) 

    #Select svm type
    if svmType == 1:
        model = make_pipeline(StandardScaler(), SVC(gamma='auto')) 
    elif svmType == 2:
        model = make_pipeline(StandardScaler(), LinearSVC(random_state=0))

    #Prepare saving directory
    if not (os.path.exists(SaveDir)):
        os.mkdir(SaveDir)

    #Searchlight analysis
    tData = np.zeros([RoiLen,len(SubList_test)])
    for rr in range(RoiLen): #Loop for voxels in whole brain
        print('Calculating searchligth at voxel #' + str(rr) )
        #Get indices of target voxels included in Sphere ROI
        SpherePos = SphereSet[rr]
        vind = [VoxelIndex_targ.index(tuple(SpherePos[vv])) for vv in range(len(SpherePos))]
        if len(vind) < MinVoxN: #Skip if number of voxels in the shpere is very small
            tAcc = np.NaN
        else:
            tAcc = np.zeros(len(SubList_test))
            #Training with all subjects in the 1st group
            TrainingSub = list(range(len(SubList_trn)))

            #Concatenate training subjects data
            Val01_List = [tRespVal01_trn[ss][:,vind] for ss in TrainingSub]
            X01 = np.concatenate(Val01_List, axis=0)
            Val02_List = [tRespVal02_trn[ss][:,vind] for ss in TrainingSub]
            X02 = np.concatenate(Val02_List, axis=0)
            X = np.concatenate([X01,X02],axis=0)   #Concatenate data from two conditions
            X = np.nan_to_num(X) #Replace NaN to 0
            Y = np.concatenate([np.zeros(X01.shape[0]), np.ones(X02.shape[0])])

            #SVM model training
            model.fit(X, Y)

            for TestSub in range(len(SubList_test)):
                #Concatenate test subjects data in the 2nd group
                X01_test = tRespVal01_test[TestSub][:,vind]
                X02_test = tRespVal02_test[TestSub][:,vind]
                X_test = np.concatenate([X01_test,X02_test],axis=0)
                X_test = np.nan_to_num(X_test)
                Y_test = np.concatenate([np.zeros(X01_test.shape[0]), np.ones(X02_test.shape[0])])

                #Model testing
                Y_pred = model.predict(X_test)
                tAcc[TestSub]  = accuracy_score(Y_pred,Y_test)

        print('Mean Accuracy: {}'.format(round(np.mean(tAcc),3)))  
        tData[rr,:] = tAcc
        
    fname = 'SearchLight_8to5_{}_{}'.format(TaskName, MethodName)
    np.save(os.path.join(SaveDir,fname),tData)


def SaveNifti_PermTest(TaskName, MethodName='Test', PeakThr=0.005, ClusterThr=0.05, N_Rand=50000):

    #N_Rand = 50000 #Number of permutation
    #PeakThr = 0.005 #Peak level threshold
    #ClusterThr = 0.05 #Cluster level threshold
    
    CurrDir = os.getcwd()
    SaveDir = CurrDir + '/Result_8ans/Result_SearchLight_8to5_' + TaskName + '/'
    fname = os.path.join(SaveDir, 'SearchLight_8to5_' + TaskName + '_' + MethodName + '.npy')
    SaveName = 'SearchLight_8to5_{}_{}_PermTest_LogPval.nii'.format(TaskName, MethodName)
    SaveName_FDR = 'SearchLight_8to5_{}_{}_PermTest_Peak{}_FDR.nii'.format(TaskName, MethodName, str(PeakThr)[2:])
    
    #Load decoding result
    tData = np.load(fname)
    
    #Sign permutation test
    pData, mData_rn = pu.SignPerm(tData, N_Rand) 

    #Transform 1d p value result into 3d MNI space
    RoiName = 'GrayMatter.nii'
    RoiDir = CurrDir + '/WFU_AALAtlas'
    RoiImg = nib.load(os.path.join(RoiDir, RoiName))
    RoiVol = RoiImg.get_fdata()

    #Null distribution of cluster size
    random.seed(1234)     
    N_Rand_Cluster = round(N_Rand / 100)
    tClusterSize = [pu.CalcClusterSize(tData, mData_rn, RoiVol, PeakThr, rr) for rr in range(N_Rand_Cluster)]
    ClusterSizeDist = np.concatenate(tClusterSize)

    #Calculate actual cluster size
    pData_vol = np.ones(np.shape(RoiVol))
    pData_vol[np.nonzero(RoiVol)] = pData
    cluster_map, n_clusters = ndimage.label(pData_vol < PeakThr)
    TargClusterSize = np.bincount(cluster_map.ravel())[1:]

    #Transform cluster size data to p values
    ClusterPval = np.zeros(np.shape(TargClusterSize)[0])
    for cc in range(np.shape(TargClusterSize)[0]):
        ClusterPval[cc] = np.count_nonzero(ClusterSizeDist >= TargClusterSize[cc]) / np.shape(ClusterSizeDist)[0]

    #FDR correction
    fdr_mask, p_fdr = fdrcorrection(ClusterPval, alpha=ClusterThr)

    #Masking p value with fdr clustere correction
    Mask = np.zeros(np.shape(pData_vol))
    for mm in np.where(fdr_mask)[0]:
        Cluster = np.where(cluster_map==mm+1)
        for vv in range(np.shape(Cluster)[1]):
            Mask[Cluster[0][vv],Cluster[1][vv],Cluster[2][vv]] = 1

    #Transform p values to -log(p)
    pData_vol[pData_vol==0] = 1 / N_Rand # Correction in case of pvalue = 0
    pData_vol = abs(np.log10(pData_vol))
    pData_vol_fdr = pData_vol*Mask

    #Save data
    SaveImg = nib.Nifti1Image(pData_vol, RoiImg.affine, RoiImg.header)
    nib.save(SaveImg, SaveDir + SaveName)
    SaveImg = nib.Nifti1Image(pData_vol_fdr, RoiImg.affine, RoiImg.header)
    nib.save(SaveImg, SaveDir + SaveName_FDR)
    
    return