#!/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, colors
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[8]:


res_shape = np.array(mask1mm.shape[:3])
bbox = np.argwhere(mask1mm)
new_origin = np.clip(bbox.min(0) - 5, a_min=0, a_max=None)
new_end = np.clip(bbox.max(0) + 5, a_min=0,
                  a_max=res_shape - 1)

# Find new origin, and set into new affine
new_affine_4 = atlas.affine.copy()
new_affine_4[:3, 3] = new_affine_4[:3, :3].dot(
    new_origin) + new_affine_4[:3, 3]

cropped_atlas = atlas.__class__(
    atlas.get_data()[new_origin[0]:new_end[0], new_origin[1]:new_end[1], new_origin[2]:new_end[2]],
    new_affine_4,
    atlas.header
)


# In[11]:


plt.clf()
fig = plt.gcf()
_ = fig.set_size_inches(12, 10)
gs = gridspec.GridSpec(4, 3, height_ratios=[15, 15, 4, 0.5],  hspace=0.0, wspace=0.)
ax1 = plt.subplot(gs[0, :])
ax2 = plt.subplot(gs[1, :])

cut_coords = [0, 15, 30]
disp = plotting.plot_stat_map(str(ANALYSIS_HOME / 'acm_fpre.nii.gz'),
                              bg_img=cropped_atlas, threshold=0.25, display_mode='z',
                              cut_coords=cut_coords, vmax=0.8, alpha=0.8,
                              axes=ax1, colorbar=False, annotate=False)
disp.annotate(size=20, left_right=True, positions=True)

disp = plotting.plot_stat_map(str(ANALYSIS_HOME / 'acm_feat.nii.gz'),
                       bg_img=cropped_atlas, threshold=0.25, display_mode='z',
                       cut_coords=cut_coords, vmax=0.8, alpha=0.8,
                       axes=ax2, colorbar=False, annotate=False)
disp.annotate(size=24, left_right=False, positions=False, scalebar=True,
              loc=3, size_vertical=2, label_top=False, frameon=True, borderpad=0.1,
              bg_color='None')


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

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

ax3 = fig.add_subplot(gs[3, 2])

cmap = plotting.cm.cold_hot
gradient = np.linspace(-0.8, 0.8, cmap.N)

# istart = int(norm(-offset, clip=True) * (our_cmap.N - 1))
# istop = int(norm(offset, clip=True) * (our_cmap.N - 1))
GRAY = (0.85, 0.85, 0.85, 1.)
cmaplist = []
for i in range(cmap.N):
    cmaplist += [cmap(i)] if not -0.25 < gradient[i] < 0.25 else [GRAY]  # just an average gray color

cmap = colors.LinearSegmentedColormap.from_list('Custom cmap', cmaplist, cmap.N)

th_index = cmaplist.index(GRAY)

gradient = np.vstack((gradient, gradient))
ax3.imshow(gradient, aspect='auto', cmap=cmap)
ax3.set_title(r'\noindent\parbox{7.5cm}{\centering\textbf{Fraction of participants} \\ with significant response}',
              size=18, position=(0.5, 3.0))
ax3.xaxis.set_ticklabels(['80\%', '25\%', '25\%', '80\%'], size=20)
ax3.xaxis.set_ticks([0, th_index, cmap.N - th_index - 1, cmap.N])
ax3.yaxis.set_ticklabels([])
ax3.yaxis.set_ticks([])
for pos in ['top', 'bottom', 'left', 'right']:
#     ax3.spines[pos].set_visible(False)
    ax3.spines[pos].set_color(GRAY)

# ax3.annotate(
#     r'\noindent\parbox{15cm}{'
#     r'Cold hues represent negative activation (response inhibition). '
#     r'Warm hues represent positive activation. '
#     r'Activation count maps are derived from N=257 biologically independent participants.}',
#     xy=(-2.05, 0.07), xycoords='axes fraction', xytext=(.0, 10.0),
#     textcoords='offset points', va='center', color='k', size=14,
# )


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


# In[ ]: