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

# In[1]:


# %load_ext autoreload
# %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 niworkflows.data import get_template

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 = get_template('MNI152NLin2009cAsym')
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 / 'tpl-MNI152NLin2009cAsym_space-MNI_res-01_T1w.nii.gz'))
mask1mm = nb.load(str(ATLAS_HOME / 'tpl-MNI152NLin2009cAsym_space-MNI_res-01_brainmask.nii.gz')).get_data() > 0
mask2mm = nb.load(str(ATLAS_HOME / 'tpl-MNI152NLin2009cAsym_space-MNI_res-02_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]:


# sn.set_style("whitegrid", {
#     'ytick.major.size': 5,
#     'xtick.major.size': 5,
# })
# sn.set_context("notebook", font_scale=1.5)

# pgf_with_custom_preamble = {
#     'ytick.major.size': 0,
#     'xtick.major.size': 0,
#     'font.size': 30,
#     'font.sans-serif': ['HelveticaLTStd-Light'],
#     'font.family': 'sans-serif', # use serif/main font for text elements
#     'text.usetex': False,    # use inline math for ticks
# }
# mpl.rcParams.update(pgf_with_custom_preamble)


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():
        print('Forced or %s not found' % pipe_mean)
        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


# In[4]:


# Use the WM mask to normalize intensities of EPI means
meanmask = nb.load(str(ATLAS_HOME / 'tpl-MNI152NLin2009cAsym_space-MNI_res-02_class-WM_probtissue.nii.gz')).get_data() > 0.9

# 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]:


df = pd.read_csv(str(ANALYSIS_HOME / 'smoothness.csv'))

plt.clf()
fig = plt.gcf()
_ = fig.set_size_inches(15, 2 * 3.1)
# gs = gridspec.GridSpec(2, 4, width_ratios=[38, 7, 60, 10], height_ratios=[1, 1],  hspace=0.0, wspace=0.03)
gs = gridspec.GridSpec(2, 3, width_ratios=[42, 9, 64], height_ratios=[1, 1],  hspace=0.0, wspace=0.03)

a_ax1 = plt.subplot(gs[0, 0])
a_ax2 = plt.subplot(gs[1, 0])


fmriprep_smooth = df[df.pipeline.str.contains('fmriprep')][['fwhm_pre', 'fwhm_post']]
feat_smooth = df[df.pipeline.str.contains('feat')][['fwhm_pre', 'fwhm_post']]

cols = palettable.tableau.ColorBlind_10.hex_colors
sn.distplot(fmriprep_smooth.fwhm_post, color=cols[0], ax=a_ax2,
            axlabel='Smoothing', label='fMRIPrep')
sn.distplot(feat_smooth.fwhm_post, color=cols[1], ax=a_ax2,
            axlabel='Smoothing', label=r'\texttt{feat}')
sn.distplot(fmriprep_smooth.fwhm_pre, color=cols[0], ax=a_ax1,
            axlabel='Smoothing', label='fMRIPrep')
sn.distplot(feat_smooth.fwhm_pre, color=cols[1], ax=a_ax1,
            axlabel='Smoothing', label=r'\texttt{feat}')

a_ax2.set_xlim([3, 8.8])
a_ax2.set_xticks([])
a_ax2.set_xticklabels([])
a_ax2.xaxis.tick_bottom()
a_ax2.grid(False)
a_ax2.set_xlabel('')

a_ax2.set_ylim([-1.1, 9.9])
a_ax2.set_yticks([])
a_ax2.spines['left'].set_visible(False)

a_ax1.set_ylabel(r'\noindent\parbox{4.8cm}{\centering\textbf{Before smoothing} fraction of images}',
                 size=13)
a_ax1.yaxis.set_label_coords(-0.1, 0.4)
a_ax2.set_ylabel(r'\noindent\parbox{4.8cm}{\centering\textbf{After smoothing} fraction of images}',
                 size=13)
a_ax2.yaxis.set_label_coords(-0.1, 0.6)

# ax4.spines['bottom'].set_position(('outward', 20))
a_ax2.invert_yaxis()
a_ax2.spines['top'].set_visible(False)
a_ax2.spines['bottom'].set_visible(False)
a_ax2.spines['left'].set_visible(False)
a_ax2.spines['right'].set_visible(False)

a_ax1.set_xlim([3, 8.8])
a_ax1.set_ylim([-0.6, 10.4])
a_ax1.grid(False)

a_ax1.set_xlabel('(mm)')
a_ax1.xaxis.set_label_coords(0.95, 0.1)
a_ax1.set_yticks([])
a_ax1.set_yticklabels([])
a_ax1.set_xticks([3, 4, 5, 6 , 7, 8])
a_ax1.set_xticklabels([3, 4, 5, 6 , 7, 8])
a_ax1.tick_params(axis='x', zorder=100, direction='inout')
a_ax1.spines['left'].set_visible(False)
a_ax1.spines['right'].set_visible(False)
a_ax1.spines['top'].set_visible(False)
a_ax1.spines['bottom'].set_visible(True)
a_ax1.zorder = 100

# a_ax2.xaxis.set_label_position('top')
# a_ax2.xaxis.set_label_coords(0.45, 0.95)
a_ax1.annotate(
    r'\noindent\parbox{6.8cm}{\centering\textbf{Estimated smoothness} full~width~half~maximum~(mm)}',
    xy=(0.15, 0.8), xycoords='axes fraction', xytext=(.0, .0),
    textcoords='offset points', va='center', color='k', size=13,
)

legend = a_ax2.legend(ncol=2, loc='upper center', bbox_to_anchor=(0.5, 0.45), prop={'size': 15})
legend.get_frame().set_facecolor('w')
legend.get_frame().set_edgecolor('none')

# a_ax2.annotate(
#     r'\noindent\parbox{15cm}{Panels A, B present statistics derived from N=257 biologically independent participants}',
#     xy=(-0.2, 0.03), xycoords='axes fraction', xytext=(.0, .0),
#     textcoords='offset points', va='center', color='k', size=11,
# )

###### PANEL B

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


thres = 20
vmin = 50
vmax = 200

disp = plotting.plot_anat(str(fprep_std), display_mode='z', annotate=False,
                          cut_coords=[-5, 10, 20], cmap='cividis', threshold=thres, vmin=vmin, vmax=vmax,
                          axes=b_ax1)
disp.annotate(size=12, left_right=True, positions=True)
disp.add_contours(str(ATLAS_HOME / 'tpl-MNI152NLin2009cAsym_space-MNI_res-01_class-CSF_probtissue.nii.gz'), colors=['k'], levels=[0.8])
disp.add_contours(str(ATLAS_HOME / 'tpl-MNI152NLin2009cAsym_space-MNI_res-01_class-WM_probtissue.nii.gz'), colors=['w'], levels=[0.8], linewidths=[1], alpha=0.7)
disp.add_contours(str(ATLAS_HOME / 'tpl-MNI152NLin2009cAsym_space-MNI_res-01_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=[-5, 10, 20], cmap='cividis', threshold=thres, vmin=vmin, vmax=vmax,
                          axes=b_ax2)
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 / 'tpl-MNI152NLin2009cAsym_space-MNI_res-01_class-CSF_probtissue.nii.gz'), colors=['k'], levels=[0.8])
disp.add_contours(str(ATLAS_HOME / 'tpl-MNI152NLin2009cAsym_space-MNI_res-01_class-WM_probtissue.nii.gz'), colors=['w'], levels=[0.8], linewidths=[1], alpha=0.7)
disp.add_contours(str(ATLAS_HOME / 'tpl-MNI152NLin2009cAsym_space-MNI_res-01_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, 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')


a_ax1.set_title('A', fontdict={'fontsize': 24}, loc='left', x=-0.2);
b_ax1.set_title('B', fontdict={'fontsize': 24}, loc='left');

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


# In[ ]:


coords = [-27, 0, 7]

thres = 20
vmin = 50
vmax = 200


# Plot
plt.clf()
fig = plt.figure(figsize=(20,10))

plotting.plot_anat('newfeat.nii.gz', cut_coords=coords, colorbar=True, cmap='cividis',
                   threshold=thres, vmin=vmin, vmax=vmax, title='feat',
                   axes=plt.subplot(2,2,1)
);
plotting.plot_anat(str(fprep_std), cut_coords=coords, colorbar=True, cmap='cividis',
                   threshold=thres, vmin=vmin, vmax=vmax, title='fmriprep',
                   axes=plt.subplot(2,2,3)
);
plotting.plot_glass_brain(str(feat_std), threshold=200, colorbar=True, title='feat',
                          axes=plt.subplot(2,2,2));
plotting.plot_glass_brain(str(fprep_std), threshold=200, colorbar=True, title='fmriprep',
                          axes=plt.subplot(2,2,4));


# In[ ]: