import os,shutil
from nipype.interfaces import fsl
import nibabel
import numpy
import nilearn.plotting

%matplotlib inline

import matplotlib.pyplot as plt
try:
    datadir=os.environ['FMRIDATADIR']
    assert not datadir==''
except:
    datadir='/Users/poldrack/data_unsynced/myconnectome/sub00001'

basedir=os.path.join(datadir,'ds009')
notebook_base= os.path.abspath(".")
cope='cope001_go'
results_dir = os.path.abspath("../results")

if not os.path.exists(results_dir):
    os.mkdir(results_dir)

from nipype.caching import Memory
mem = Memory(base_dir='.')

copefile=os.path.join(basedir,'task002_%s.nii.gz'%cope)
varcopefile=os.path.join(basedir,'task002_var%s.nii.gz'%cope)
nsubs=24


"""
Explanation: Run a group analysis using FSL.  The data are obtained from the Human Connectome Project, and include 32 subjects on a relational matching task (using the Relational - Match contrast).
End of explanation
"""


    
deshdr="""/NumWaves	1
/NumPoints	24
/PPheights		1.000000e+00

/Matrix"""

conhdr="""/ContrastName1	group mean
/NumWaves	1
/NumContrasts	1
/PPheights		1.000000e+00
/RequiredEffect		1.441

/Matrix"""

grouphdr="""/NumWaves	1
/NumPoints	24

/Matrix"""
desmtx=numpy.ones((nsubs,1))
numpy.savetxt('design.mat',desmtx,fmt='%1.0f',header=deshdr,comments='')
numpy.savetxt('covsplit.txt',desmtx,fmt='%1.0f',header=grouphdr,comments='')
conmtx=numpy.ones(1)
numpy.savetxt('design.con',conmtx,fmt='%1.0f',header=conhdr,comments='')

level2=mem.cache(fsl.FLAMEO)
flameo_results = level2(cope_file=copefile, 
                    var_cope_file=varcopefile,
                    design_file='design.mat',
                    cov_split_file='covsplit.txt',
                    t_con_file='design.con',
                    mask_file=os.path.join(os.getenv('FSLDIR'),'data/standard/MNI152_T1_2mm_brain_mask_dil.nii.gz'),
                    run_mode='ols')

flameo_results.outputs

shutil.move(flameo_results.outputs.stats_dir,os.path.join(results_dir,'stats_OLS'))


"""
Explanation: First run using OLS.
End of explanation
"""


est = mem.cache(fsl.SmoothEstimate)
smoothness=est(dof=23,
    residual_fit_file = os.path.join(results_dir,'stats_OLS/res4d.nii.gz'),
    mask_file = os.path.join(os.getenv('FSLDIR'),'data/standard/MNI152_T1_2mm_brain_mask_dil.nii.gz'))

print(smoothness.outputs)

cl = mem.cache(fsl.Cluster)
cluster_results=cl(threshold = 2.3,
   in_file = os.path.join(results_dir,'stats_OLS/zstat1.nii.gz'),
    dlh=smoothness.outputs.dlh,
    volume=smoothness.outputs.volume,
    pthreshold=0.05,
    out_localmax_txt_file=os.path.join(results_dir,'stats_OLS/zstat1_cluster_max.txt'),
    out_threshold_file=os.path.join(results_dir,'stats_OLS/zstat1_thresh.nii.gz'))

with open(os.path.join(results_dir,'stats_OLS/zstat1_cluster_list.txt'),'w') as f:
    f.write(cluster_results.runtime.stdout)

thresh_zstat_OLS=nibabel.load(os.path.join(results_dir,'stats_OLS/zstat1_thresh.nii.gz'))
bgimage=nibabel.load(os.path.join(os.getenv('FSLDIR'),'data/standard/MNI152_T1_2mm_brain.nii.gz'))
map_display=nilearn.plotting.plot_stat_map(thresh_zstat_OLS,bgimage,threshold=2.3,title='cluster corrected')

"""
Explanation: Estimate smoothness and compute cluster-corrected thresholded map
End of explanation
"""


level2=mem.cache(fsl.FLAMEO)
flameo_results = level2(cope_file=copefile, 
                    var_cope_file=varcopefile,
                    design_file='design.mat',
                    cov_split_file='covsplit.txt',
                    t_con_file='design.con',
                    mask_file=os.path.join(os.getenv('FSLDIR'),'data/standard/MNI152_T1_2mm_brain_mask_dil.nii.gz'),
                    run_mode='flame1')

flameo_results.outputs

shutil.move(flameo_results.outputs.stats_dir,os.path.join(results_dir,'stats_FLAME1'))



"""
Explanation: Run same model using mixed effects with FLAME 1
End of explanation
"""


level2=mem.cache(fsl.FLAMEO)
flameo_results = level2(cope_file=copefile, 
                    var_cope_file=varcopefile,
                    design_file='design.mat',
                    cov_split_file='covsplit.txt',
                    t_con_file='design.con',
                    mask_file=os.path.join(os.getenv('FSLDIR'),'data/standard/MNI152_T1_2mm_brain_mask_dil.nii.gz'),
                    run_mode='fe')

flameo_results.outputs

shutil.move(flameo_results.outputs.stats_dir,os.path.join(results_dir,'stats_fe'))



"""
Explanation: Run fixed effects model for comparison
End of explanation
"""


try:
    maskdata
except:
    zstat_fe=nibabel.load(os.path.join(results_dir,'stats_fe/zstat1.nii.gz')).get_data()
    zstat_FLAME1=nibabel.load(os.path.join(results_dir,'stats_FLAME1/zstat1.nii.gz')).get_data()
    zstat_OLS=nibabel.load(os.path.join(results_dir,'stats_OLS/zstat1.nii.gz')).get_data()
    maskdata=(zstat_fe!=0)*(zstat_FLAME1!=0)*(zstat_OLS!=0)

mindata=numpy.min([numpy.min(zstat_fe[maskdata>0]),numpy.min(zstat_OLS[maskdata>0])])
maxdata=numpy.max([numpy.max(zstat_fe[maskdata>0]),numpy.max(zstat_OLS[maskdata>0])])
plt.figure(figsize=(12,4))
plt.subplot(121)
plt.hist(zstat_fe[zstat_OLS>0] - zstat_OLS[zstat_OLS>0],100)
plt.xlabel('FE Z - OLS Z')
plt.subplot(122)

sc=plt.scatter(zstat_OLS[maskdata>0],zstat_fe[maskdata>0])
mindata=numpy.min([sc.get_axes().get_xlim()[0],sc.get_axes().get_ylim()[0]])
maxdata=numpy.max([sc.get_axes().get_xlim()[1],sc.get_axes().get_ylim()[1]])
plt.plot([mindata,maxdata],[mindata,maxdata],color='red')
plt.ylabel('Fixed effect Z')
plt.xlabel('OLS Z')


mindata=numpy.min([numpy.min(zstat_FLAME1[maskdata>0]),numpy.min(zstat_OLS[maskdata>0])])
maxdata=numpy.max([numpy.max(zstat_FLAME1[maskdata>0]),numpy.max(zstat_OLS[maskdata>0])])
plt.figure(figsize=(12,4))
plt.subplot(121)
plt.hist(zstat_OLS[zstat_OLS>0] - zstat_FLAME1[zstat_OLS>0],100)
plt.xlabel('OLS Z - FLAME1 Z')
plt.subplot(122)

sc=plt.scatter(zstat_FLAME1[maskdata>0],zstat_OLS[maskdata>0])
mindata=numpy.min([sc.get_axes().get_xlim()[0],sc.get_axes().get_ylim()[0]])
maxdata=numpy.max([sc.get_axes().get_xlim()[1],sc.get_axes().get_ylim()[1]])
plt.plot([mindata,maxdata],[mindata,maxdata],color='red')
plt.ylabel('OLS Z')
plt.xlabel('FLAME1 Z')

print ('mean absolute difference:')
print ('FE Z - OLS Z:',numpy.mean(numpy.abs(zstat_fe[maskdata>0] - zstat_OLS[maskdata>0])))
print ('OLS Z - FLAME1 Z:',numpy.mean(numpy.abs(zstat_OLS[maskdata>0] - zstat_FLAME1[maskdata>0])))

"""
Explanation: Plot fixed effects stats against OLS and OLS against FLAME1
End of explanation
"""


residfile=os.path.join(results_dir,'stats_FLAME1/res4d.nii.gz')
res4d=nibabel.load(residfile).get_data()
resdata=numpy.zeros((res4d.shape[3],numpy.sum(maskdata)))

for i in range(res4d.shape[3]):
    tmp=res4d[:,:,:,i]
    resdata[i,:]=tmp[maskdata>0]

import scipy.cluster.hierarchy
hclust=scipy.cluster.hierarchy.ward(resdata)
cl=scipy.cluster.hierarchy.dendrogram(hclust)

import sklearn.manifold
t=sklearn.manifold.TSNE(init='pca')
embedding=t.fit_transform(resdata)
plt.figure()
plt.axis([numpy.min(embedding[:,0])*1.2,numpy.max(embedding[:,0])*1.2,numpy.min(embedding[:,1])*1.2,numpy.max(embedding[:,1])*1.2])
for i in range(embedding.shape[0]):
    
    plt.text(embedding[i,0],embedding[i,1],'%d'%i)


"""
Explanation: Look for outlier subjects by examining residuals
End of explanation
"""


