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

# In[1]:


get_ipython().run_line_magic('load_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')
# %matplotlib inline
import os
from pathlib import Path
import warnings

import numpy as np
import nibabel as nb
import pandas as pd

import matplotlib as mpl
mpl.use('pgf')
from matplotlib import pyplot as plt
from matplotlib import gridspec
import seaborn as sn
import palettable

from nilearn.image import concat_imgs, mean_img
from nilearn import plotting

warnings.simplefilter('ignore')

DATA_HOME = Path(os.getenv('FMRIPREP_DATA_HOME', os.getcwd())).resolve()
DS030_HOME = DATA_HOME / 'ds000030' / '1.0.3'
DERIVS_HOME = DS030_HOME / 'derivatives'
ATLAS_HOME = DATA_HOME / 'templates' / 'mni_icbm152_nlin_asym_09c'
ANALYSIS_HOME = DERIVS_HOME / 'fmriprep_vs_feat_2.0-oe'

fprep_home = DERIVS_HOME / 'fmriprep_1.0.8' / 'fmriprep'
feat_home = DERIVS_HOME / 'fslfeat_5.0.10' / 'featbids'

out_folder = Path(os.getenv('FMRIPREP_OUTPUTS') or '').resolve()

# Load MNI152 nonlinear, asymmetric 2009c atlas
atlas = nb.load(str(ATLAS_HOME / '1mm_T1.nii.gz'))
mask1mm = nb.load(str(ATLAS_HOME / '1mm_brainmask.nii.gz')).get_data() > 0
mask2mm = nb.load(str(ATLAS_HOME / '2mm_brainmask.nii.gz')).get_data() > 0
data = atlas.get_data()
data[~mask1mm] = data[~mask1mm].max()
atlas = nb.Nifti1Image(data, atlas.affine, atlas.header)


# In[2]:


pgf_with_custom_preamble = {
    'text.usetex': True,    # use inline math for ticks
    'pgf.rcfonts': False,   # don't setup fonts from rc parameters
    'pgf.texsystem': 'xelatex',
    'verbose.level': 'debug-annoying',
    "pgf.preamble": [
        r"""\usepackage{fontspec}
\setsansfont{HelveticaLTStd-Light}[
Extension=.otf,
BoldFont=HelveticaLTStd-Bold,
ItalicFont=HelveticaLTStd-LightObl,
BoldItalicFont=HelveticaLTStd-BoldObl,
]
\setmainfont{HelveticaLTStd-Light}[
Extension=.otf,
BoldFont=HelveticaLTStd-Bold,
ItalicFont=HelveticaLTStd-LightObl,
BoldItalicFont=HelveticaLTStd-BoldObl,
]
\setmonofont{Inconsolata-dz}
""",
        r'\renewcommand\familydefault{\sfdefault}',
#         r'\setsansfont[Extension=.otf]{Helvetica-LightOblique}',
#         r'\setmainfont[Extension=.ttf]{DejaVuSansCondensed}',
#         r'\setmainfont[Extension=.otf]{FiraSans-Light}',
#         r'\setsansfont[Extension=.otf]{FiraSans-Light}',
    ]
}
mpl.rcParams.update(pgf_with_custom_preamble)


# In[3]:


def mean_std_map(pipe_home, meanmask, force=False, lazy=False, maskval=1000):
    pipe_std = pipe_home / 'summary_stdev.nii.gz'
    pipe_mean = pipe_home / 'summary_means.nii.gz'

    if force or not pipe_mean.is_file():
        all_mus = []
        if lazy:
            all_mus = [nb.load(str(f)) for f in pipe_home.glob(
                'sub-*/func/sub-*_task-stopsignal_bold_space-MNI152NLin2009cAsym_avgpreproc.nii.gz')]
        
        if not all_mus:
            print('Generating means file')
            pipe_files = list(pipe_home.glob(
                'sub-*/func/sub-*_task-stopsignal_bold_space-MNI152NLin2009cAsym_preproc.nii.gz'))
            all_mus = []
            for f in pipe_files:
                mean = mean_img(str(f))
                data = mean.get_data()
                sigma = np.percentile(data[meanmask], 50) / maskval
                data /= sigma
                all_mus.append(nb.Nifti1Image(data, mean.affine, mean.header))
                
        meannii = concat_imgs(all_mus, auto_resample=False)
        meannii.to_filename(str(pipe_mean))
        force = True

    if force or not pipe_std.is_file():
        print('Generating standard deviation map')
        meannii = nb.load(str(pipe_mean))
        nb.Nifti1Image(meannii.get_data().std(3), meannii.affine, meannii.header).to_filename(str(pipe_std))
        
    return pipe_mean, pipe_std

# Use the WM mask to normalize intensities of EPI means
meanmask = nb.load(str(ATLAS_HOME / '2mm_tpm_wm.nii.gz')).get_data() > 0.9


# In[4]:


# Calculate average and std
fprep_mean, fprep_std = mean_std_map(fprep_home, meanmask)
feat_mean, feat_std = mean_std_map(feat_home, meanmask)

# Trick to avoid nilearn zooming in
feat_nii = nb.load(str(feat_std))

fd = feat_nii.get_data()
fd[0, 0, :] = 50
fd[0, -1, :] = 50
fd[-1, 0, :] = 50
fd[-1, -1, :] = 50

nb.Nifti1Image(fd, feat_nii.affine, feat_nii.header).to_filename('newfeat.nii.gz')


# In[5]:


plt.clf()
fig = plt.gcf()
_ = fig.set_size_inches(15, 7)
gs = gridspec.GridSpec(2, 2, width_ratios=[10, 1], height_ratios=[1, 1],  hspace=0.0, wspace=0.03)

b_ax1 = plt.subplot(gs[0, 0])
b_ax2 = plt.subplot(gs[1, 0])

thres = 20
vmin = 50
vmax = 200

disp = plotting.plot_anat(str(fprep_std), display_mode='z', annotate=False,
                          cut_coords=[-15, -5, 10, 20, 40], cmap='cividis', threshold=thres, vmin=vmin, vmax=vmax,
                          axes=b_ax1, resampling_interpolation='nearest')
disp.annotate(size=12, left_right=True, positions=True)
disp.add_contours(str(ATLAS_HOME / '1mm_tpm_csf.nii.gz'), colors=['k'], levels=[0.8])
disp.add_contours(str(ATLAS_HOME / '1mm_tpm_wm.nii.gz'), colors=['w'], levels=[0.8], linewidths=[1], alpha=0.7)
disp.add_contours(str(ATLAS_HOME / '1mm_brainmask.nii.gz'), colors=['k'], levels=[0.8], linewidths=[3], alpha=.7)


disp = plotting.plot_anat('newfeat.nii.gz', display_mode='z', annotate=False,
                          cut_coords=[-15, -5, 10, 20, 40], cmap='cividis', threshold=thres, vmin=vmin, vmax=vmax,
                          axes=b_ax2, resampling_interpolation='nearest')
disp.annotate(size=12, left_right=False, positions=False)
disp.annotate(size=12, left_right=False, positions=False, scalebar=True,
              loc=3, size_vertical=2, label_top=False, frameon=True, borderpad=0.1)

disp.add_contours(str(ATLAS_HOME / '1mm_tpm_csf.nii.gz'), colors=['k'], levels=[0.8])
disp.add_contours(str(ATLAS_HOME / '1mm_tpm_wm.nii.gz'), colors=['w'], levels=[0.8], linewidths=[1], alpha=0.7)
disp.add_contours(str(ATLAS_HOME / '1mm_brainmask.nii.gz'), colors=['k'], levels=[0.8], linewidths=[3], alpha=.7)

b_ax1.annotate(
    'fMRIPrep',
    xy=(0., .5), xycoords='axes fraction', xytext=(-20, .0),
    textcoords='offset points', va='center', color='k', size=15,
    rotation=90)

b_ax2.annotate(
    r'\texttt{feat}',
    xy=(0., .5), xycoords='axes fraction', xytext=(-20, .0),
    textcoords='offset points', va='center', color='k', size=12,
    rotation=90)


inner_grid = gridspec.GridSpecFromSubplotSpec(1, 2, width_ratios=[1.5, 15],
                                              subplot_spec=gs[:, -1], wspace=0.01)

b_ax3 = fig.add_subplot(inner_grid[0])
gradient = np.hstack((np.zeros((50,)), np.linspace(0, 1, 120), np.ones((130,))))[::-1]
gradient = np.vstack((gradient, gradient))
b_ax3.imshow(gradient.T, aspect='auto', cmap=plt.get_cmap('cividis'))
b_ax3.xaxis.set_ticklabels([])
b_ax3.xaxis.set_ticks([])
b_ax3.yaxis.set_ticklabels([])
b_ax3.yaxis.set_ticks([])

b_ax4 = fig.add_subplot(inner_grid[1])
sn.distplot(nb.load(str(fprep_std)).get_data()[mask2mm], label='fMRIPrep',
            vertical=True, ax=b_ax4, kde=False, norm_hist=True)
sn.distplot(nb.load(str(feat_std)).get_data()[mask2mm], label=r'\texttt{feat}', vertical=True,
            color='darkorange', ax=b_ax4, kde=False, norm_hist=True)

# plt.gca().set_ylim((0, 300))
plt.legend(prop={'size': 15}, edgecolor='none')

b_ax4.xaxis.set_ticklabels([])
b_ax4.xaxis.set_ticks([])
b_ax4.yaxis.set_ticklabels([])
b_ax4.yaxis.set_ticks([])

plt.axis('off')
b_ax3.axis('off')
b_ax4.axis('off')


# In[18]:


feat_closeup = nb.load(str(feat_std))
feat_data = feat_closeup.get_data()[15:82, 72:106, :]
nb.Nifti1Image(feat_data, feat_closeup.affine, feat_closeup.header).to_filename('feat_clip.nii.gz')

fprep_closeup = nb.load(str(fprep_std))
fprep_data = fprep_closeup.get_data()[15:82, 72:106, :]
nb.Nifti1Image(fprep_data, fprep_closeup.affine, fprep_closeup.header).to_filename('fprep_clip.nii.gz')


msknii = nb.load(str(ATLAS_HOME / '1mm_tpm_csf.nii.gz'))
fullbox = np.zeros(msknii.shape)
fullbox[30:164, 142:213, 70:76] = 1
fullbox[30:164, 142:213, 88:90] = 2
fullbox[30:164, 142:213, 96:100] = 3
nb.Nifti1Image(fullbox, msknii.affine, msknii.header).to_filename('box.nii.gz')

newcsf = msknii.get_data()
newcsf[:, :, 100:] = 0
nb.Nifti1Image(newcsf, msknii.affine, msknii.header).to_filename('newcsfmsk.nii.gz')

mskdat = msknii.get_data()[30:164, 142:213, :]
nb.Nifti1Image(mskdat, msknii.affine, msknii.header).to_filename('csf_clip.nii.gz')
msknii = nb.load(str(ATLAS_HOME / '1mm_tpm_wm.nii.gz'))
mskdat = msknii.get_data()[30:164, 142:213, :]
nb.Nifti1Image(mskdat, msknii.affine, msknii.header).to_filename('wm_clip.nii.gz')
msknii = nb.load(str(ATLAS_HOME / '1mm_brainmask.nii.gz'))
mskdat = msknii.get_data()[30:164, 142:213, :]
nb.Nifti1Image(mskdat, msknii.affine, msknii.header).to_filename('bm_clip.nii.gz')

mskbox = np.zeros_like(mskdat)
mskbox[1:-1, 1:-1, :] = 1
nb.Nifti1Image(mskbox, msknii.affine, msknii.header).to_filename('box_clip.nii.gz')


# In[19]:


colorboxes = [
    (244/255, 3/255, 23/255),
    (73/255, 213/255, 253/255),
    (245/255, 29/255, 245/255)
]
fs = 25

plt.clf()
fig = plt.gcf()
_ = fig.set_size_inches(15, 19.5)
gs = gridspec.GridSpec(4, 2, width_ratios=[10, 1], height_ratios=[1, 1, 0.4, 3.7],  hspace=0., wspace=0.03)

b_ax1 = plt.subplot(gs[0, 0])
b_ax2 = plt.subplot(gs[1, 0])

thres = 20
vmin = 50
vmax = 200

disp = plotting.plot_anat(str(fprep_std), display_mode='z', annotate=False,
                          cut_coords=[-15, -5, 10, 20, 40], cmap='cividis', threshold=thres, vmin=vmin, vmax=vmax,
                          axes=b_ax1, resampling_interpolation='nearest')
disp.annotate(size=15, left_right=True, positions=True)
disp.add_contours('newcsfmsk.nii.gz', colors=['k'], levels=[0.8])
disp.add_contours(str(ATLAS_HOME / '1mm_tpm_wm.nii.gz'), colors=['w'], levels=[0.8], linewidths=[1], alpha=0.7)
disp.add_contours(str(ATLAS_HOME / '1mm_brainmask.nii.gz'), colors=['k'], levels=[0.8], linewidths=[3], alpha=.7)
disp.add_contours('box.nii.gz', colors=colorboxes, levels=[0.8, 1.8, 2.8], linewidths=[3], alpha=1., zorder=20)


disp = plotting.plot_anat('newfeat.nii.gz', display_mode='z', annotate=False,
                          cut_coords=[-15, -5, 10, 20, 40], cmap='cividis', threshold=thres, vmin=vmin, vmax=vmax,
                          axes=b_ax2, resampling_interpolation='nearest')
disp.annotate(size=15, left_right=False, positions=False, scalebar=True, scale_width=2.5,
              loc=3, size_vertical=2, label_top=False, frameon=True, borderpad=0.01)

disp.add_contours('newcsfmsk.nii.gz', colors=['k'], levels=[0.8])
disp.add_contours(str(ATLAS_HOME / '1mm_tpm_wm.nii.gz'), colors=['w'], levels=[0.8], linewidths=[1], alpha=0.7)
disp.add_contours(str(ATLAS_HOME / '1mm_brainmask.nii.gz'), colors=['k'], levels=[0.8], linewidths=[3], alpha=.7)
disp.add_contours('box.nii.gz', colors=colorboxes, levels=[0.8, 1.8, 2.8], linewidths=[3], alpha=1.)


b_ax1.annotate(
    'fMRIPrep',
    xy=(0., .5), xycoords='axes fraction', xytext=(-fs, .0),
    textcoords='offset points', va='center', color='k', size=fs,
    rotation=90)

b_ax2.annotate(
    r'\texttt{feat}',
    xy=(0., .5), xycoords='axes fraction', xytext=(-fs, .0),
    textcoords='offset points', va='center', color='k', size=fs,
    rotation=90)


inner_grid = gridspec.GridSpecFromSubplotSpec(1, 2, width_ratios=[1.5, 15],
                                              subplot_spec=gs[:2, -1], wspace=0.01)

b_ax3 = fig.add_subplot(inner_grid[0])
gradient = np.hstack((np.zeros((50,)), np.linspace(0, 1, 120), np.ones((130,))))[::-1]
gradient = np.vstack((gradient, gradient))
b_ax3.imshow(gradient.T, aspect='auto', cmap=plt.get_cmap('cividis'))
b_ax3.xaxis.set_ticklabels([])
b_ax3.xaxis.set_ticks([])
b_ax3.yaxis.set_ticklabels([])
b_ax3.yaxis.set_ticks([])

b_ax4 = fig.add_subplot(inner_grid[1])
sn.distplot(nb.load(str(fprep_std)).get_data()[mask2mm], label='fMRIPrep',
            vertical=True, ax=b_ax4, kde=False, norm_hist=True)
sn.distplot(nb.load(str(feat_std)).get_data()[mask2mm], label=r'\texttt{feat}', vertical=True,
            color='darkorange', ax=b_ax4, kde=False, norm_hist=True)

# plt.gca().set_ylim((0, 300))
plt.legend(prop={'size': 15}, edgecolor='none')

b_ax4.xaxis.set_ticklabels([])
b_ax4.xaxis.set_ticks([])
b_ax4.yaxis.set_ticklabels([])
b_ax4.yaxis.set_ticks([])

plt.axis('off')
b_ax3.axis('off')
b_ax4.axis('off')


inner_grid = gridspec.GridSpecFromSubplotSpec(3, 2, width_ratios=[1, 1], hspace=0.04,
                                              subplot_spec=gs[3, :], wspace=0.01)

thres = 20
vmin = 40
vmax = 200


for i, coord in enumerate([-5, 10, 20]):
    ax1 = fig.add_subplot(inner_grid[i, 0])
    disp = plotting.plot_anat(
        'fprep_clip.nii.gz', display_mode='z', annotate=False,
        cut_coords=[coord], cmap='cividis', threshold=thres, vmin=vmin, vmax=vmax,
        resampling_interpolation='nearest', axes=ax1)
    disp.add_contours('csf_clip.nii.gz', colors=['k'], levels=[0.8])
    disp.add_contours('wm_clip.nii.gz', colors=['w'], levels=[0.8], linewidths=[1], alpha=0.7)
    disp.add_contours('bm_clip.nii.gz', colors=['k'], levels=[0.8], linewidths=[3], alpha=.7)
#     disp.annotate(size=42, left_right=False, positions=True, scalebar=False)
    
    ax1.axis('on')
    ax1.set_xticklabels([])
    ax1.set_xticks([])
    ax1.set_yticklabels([])
    ax1.set_yticks([])
    ax1.set_ylabel("z = %d" % coord, size=fs)
    for pos in ['top', 'bottom', 'left', 'right']:
        ax1.spines[pos].set_color(colorboxes[i])
        ax1.spines[pos].set_visible(True)
        ax1.spines[pos].set_linewidth(4)
        ax1.spines[pos].set_position(('outward', 2))
        
    ax2 = fig.add_subplot(inner_grid[i, 1])
    disp = plotting.plot_anat(
        'feat_clip.nii.gz', display_mode='z', annotate=False,
        cut_coords=[coord], cmap='cividis', threshold=thres, vmin=vmin, vmax=vmax,
        resampling_interpolation='nearest', axes=ax2)
    disp.add_contours('csf_clip.nii.gz', colors=['k'], levels=[0.8])
    disp.add_contours('wm_clip.nii.gz', colors=['w'], levels=[0.8], linewidths=[1], alpha=0.7)
    disp.add_contours('bm_clip.nii.gz', colors=['k'], levels=[0.8], linewidths=[3], alpha=.7)
    
    ax2.axis('on')
    ax2.set_xticklabels([])
    ax2.set_xticks([])
    ax2.set_yticklabels([])
    ax2.set_yticks([])
    for pos in ['top', 'bottom', 'left', 'right']:
        ax2.spines[pos].set_color(colorboxes[i])
        ax2.spines[pos].set_visible(True)
        ax2.spines[pos].set_linewidth(4)
        ax2.spines[pos].set_position(('outward', 2))
        
    if i == 0:
        ax1.set_title('fMRIPrep', size=fs+2, position=(0.5, 1.02))
        ax2.set_title(r'\texttt{feat}', size=fs+2, position=(0.5, 1.02))

disp.annotate(size=32, left_right=False, positions=False, scalebar=True, scale_width=2.5,
              loc=4, size_vertical=1.2, label_top=False, frameon=True, borderpad=0.2)

ax1.annotate(
    r'\noindent\parbox{25cm}{Variability mappings presented in this figure are derived from N=257 biologically independent participants}',
    xy=(0.12, -0.15), xycoords='axes fraction', xytext=(.0, .0),
    textcoords='offset points', va='center', color=(0.2, 0.2, 0.2), size=fs - 10,
)

plt.savefig(str(out_folder / 'SF04.pdf'),
            format='pdf', bbox_inches='tight', pad_inches=0.2, dpi=300)


# In[ ]: