import nipype.algorithms.modelgen as model
import pandas
import nipype.pipeline.engine as pe
from nipype.interfaces.base import Bunch

import nipype.interfaces.utility as util
import nipype.interfaces.io as nio
import nipype.interfaces.fsl as fsl

import glob, os

copes = sorted(glob.glob('/home/gdholla1/data/in_limbo/birte/copes_mni/*/cope1.nii.gz'))
varcopes = sorted(glob.glob('/home/gdholla1/data/in_limbo/birte/varcopes_mni/*/varcope1.nii.gz'))
dof_files = [os.path.abspath('dof')] * len(copes)

thresholds = [2.3, 3.1, 3.7]

from nipype.workflows.fmri.fsl import create_fixed_effects_flow
# Create mixed effects work flow
mixedfx = create_fixed_effects_flow('in_limbo_level2_birte_fsl')
mixedfx.base_dir = '/home/gdholla1/workflow_folders/'


# Set up smoothing nodes
smoother_copes  = pe.MapNode(fsl.IsotropicSmooth(), 
                      iterfield=['in_file'],
                      name='smoother_copes')

smoother_copes.inputs.in_file = copes
smoother_copes.inputs.fwhm = 8.0

smoother_varcopes  = pe.MapNode(fsl.IsotropicSmooth(), 
                      iterfield=['in_file'],
                      name='smoother_varcopes')

smoother_varcopes.inputs.in_file = varcopes
smoother_varcopes.inputs.fwhm = 8.0

mixedfx.connect(smoother_copes, 'out_file', mixedfx.get_node('inputspec'), 'copes')
mixedfx.connect(smoother_varcopes, 'out_file', mixedfx.get_node('inputspec'), 'varcopes')


mixedfx.inputs.inputspec.dof_files = dof_files
mixedfx.inputs.flameo.mask_file = '/usr/share/fsl/data/standard/MNI152_T1_2mm_brain_mask.nii.gz'
mixedfx.inputs.l2model.num_copes = len(copes)

### Mixed effects (flame1) instead of fixed effects (fe)
mixedfx.inputs.flameo.run_mode = 'flame1'


# Set up a datasink to write data to
ds = pe.Node(nio.DataSink(), name='datasink')
ds.inputs.base_directory = '/home/gdholla1/data/in_limbo/birte/level2_fsl'
mixedfx.connect(mixedfx.get_node('outputspec'), 'zstats', ds, 'zstats')

# Estimate smoothness for Gaussian Random Field Theory
smooth_est = pe.Node(fsl.SmoothEstimate(), name='smooth_est')
smooth_est.inputs.mask_file = '/usr/share/fsl/data/standard/MNI152_T1_2mm_brain_mask.nii.gz'

get_one = lambda x: x[0]

mixedfx.connect(mixedfx.get_node('outputspec'), ('zstats', get_one), smooth_est, 'zstat_file')

cluster = pe.Node(fsl.Cluster(), name='cluster')
cluster.inputs.out_threshold_file = True

# Iterate over different thresholds
cluster.iterables = [('threshold',
                      thresholds)]

# Set up GRF and write threhsolded z-map to datasink
mixedfx.connect(smooth_est, 'dlh', cluster, 'dlh')
mixedfx.connect(mixedfx.get_node('outputspec'), ('zstats', get_one), cluster, 'in_file')
mixedfx.connect(cluster, 'threshold_file', ds, 'thresholded_z')

def argmin(image, threshold=0):
    import nibabel as nb
    import numpy as np
    
    image = nb.load(image)
    data = np.ma.masked_less_equal(image.get_data(), threshold)
    return np.unravel_index(np.ma.argmin(data), data.shape)

get_comparison_voxel = pe.Node(util.Function(function=argmin,
                                             input_names=['image'],
                                             output_names=['index']),
                               name='get_comparison_voxel')

get_min_value = pe.MapNode(util.Function(function=get_value_at,
                                         input_names=['image', 'index'],
                                         output_names=['value']),
                           iterfield=['image'],
                           name='get_min_value')

def get_value_at(image, index):
    import nibabel as nb    
    
    return nb.load(image).get_data()[index]

mixedfx.connect(cluster, 'threshold_file', get_comparison_voxel, 'image')
mixedfx.connect(get_comparison_voxel, 'index', get_min_value, 'index')
mixedfx.connect(smoother_copes, 'out_file', get_min_value, 'image')

def subtract_and_inverse(image, value):
    import nibabel as nb
    import os
    
    image = nb.load(image)
    data = image.get_data()
    
    data -= value
    data *= -1
    
    fn = os.path.abspath('cope_in_limbo.nii.gz')
    nb.save(nb.Nifti1Image(data, image.get_affine()), fn)
    
    return fn


subtract_cope = pe.MapNode(util.Function(function=subtract_and_inverse,
                                         input_names=['image', 'value'],
                                         output_names=['in_limbo_cope'],
                                         ),
                           iterfield=['image',
                                      'value'],
                           name='subtract_from_cope')


mixedfx.connect(get_min_value, 'value', subtract_cope, 'value')
mixedfx.connect(smoother_copes, 'out_file', subtract_cope, 'image')

in_limbo_mixedfx = create_fixed_effects_flow('in_limbo_mixedfx')
mixedfx.connect(smoother_varcopes, 'out_file', in_limbo_mixedfx, 'inputspec.varcopes')

in_limbo_mixedfx.inputs.inputspec.dof_files = dof_files
in_limbo_mixedfx.inputs.flameo.mask_file = '/usr/share/fsl/data/standard/MNI152_T1_2mm_brain_mask.nii.gz'
in_limbo_mixedfx.inputs.l2model.num_copes = len(copes)

mixedfx.connect(subtract_cope, 'in_limbo_cope', in_limbo_mixedfx, 'inputspec.copes')


def get_in_limbo_area(thresholded_z, in_limbo_t, threshold):
    
    import nibabel as nb
    import numpy as np
    import os
    
    # Load the two level 2 results
    z_image = nb.load(thresholded_z)
    z = z_image.get_data()
    in_limbo_t = nb.load(in_limbo_t).get_data()
    
    # Set up outputs
    in_limbo = np.zeros_like(z)
    unactivated = np.zeros_like(z)
    
    # Apply logic
    in_limbo[(z == 0) & (in_limbo_t < threshold) & (in_limbo_t != 0)] = 1
    unactivated[(z == 0) & (in_limbo_t > threshold)] = 1
    
    nb.save(nb.Nifti1Image(in_limbo, z_image.get_affine(),), 
            os.path.abspath('in_limbo.nii.gz'))
    
    nb.save(nb.Nifti1Image(unactivated, z_image.get_affine(),), 
        os.path.abspath('in_limbo_unactivated.nii.gz'))
    
    return os.path.abspath('in_limbo.nii.gz'), os.path.abspath('in_limbo_unactivated.nii.gz')

in_limbo_mapper = pe.Node(util.Function(function=get_in_limbo_area,
                                        input_names=['thresholded_z', 
                                                     'in_limbo_t',
                                                     'threshold'],
                                        output_names=['in_limbo', 'unactivated']),
                          name='in_limbo_mapper')

import scipy as sp
from scipy import stats

dof = len(copes) - 1
in_limbo_mapper.inputs.threshold = sp.stats.t(dof).ppf(0.95)


mixedfx.connect(in_limbo_mixedfx.get_node('outputspec'), ('tstats', get_one), in_limbo_mapper, 'in_limbo_t')
mixedfx.connect(cluster, 'threshold_file', in_limbo_mapper, 'thresholded_z')

from IPython.display import Image
mixedfx.write_graph()
Image('/home/gdholla1/workflow_folders/in_limbo_level2_birte_fsl/graph.dot.png')

mixedfx.run()

import seaborn as sns
palette = sns.blend_palette(["seagreen", "lightblue"]);

sns.set_style('white')

# Don't put warnings in my plots, please, Python
import warnings
warnings.filterwarnings('ignore')


x_slice = 41
y_slice = 49
z_slice = 31


for threshold in [2.3, 3.1, 3.7]:    
    
    ## STANDARD FSL AR(1)-model
    in_limbo = nb.load('/home/gdholla1/workflow_folders/in_limbo_level2_birte_fsl/_threshold_%s/in_limbo_mapper/in_limbo.nii.gz' % threshold).get_data()
    z_map = nb.load('/home/gdholla1/workflow_folders/in_limbo_level2_birte_fsl/_threshold_%s/cluster/zstat1_threshold.nii.gz' % threshold).get_data()
    brain = nb.load('/usr/share/fsl/data/standard/MNI152_T1_2mm_brain.nii.gz').get_data()
    
    z_map = np.ma.masked_less_equal(z_map, 0)
    in_limbo = np.ma.masked_less_equal(in_limbo, 0)
    plt.figure(figsize=(20, 5))
    plt.suptitle('AR(1) estimation - z > %s' % threshold, fontsize=25)
    plt.subplot(131)
    plt.imshow(brain[x_slice, :, :].T, origin='lower', cmap=plt.cm.gray)
    plt.imshow(z_map[x_slice, :, :].T, origin='lower', cmap=plt.cm.hot)
    plt.imshow(in_limbo[x_slice, :, :].T, origin='lower', cmap=plt.cm.Greens_r)
    plt.axis('off')   
    
    plt.subplot(132)
    plt.imshow(brain[:, y_slice, :].T, origin='lower', cmap=plt.cm.gray)
    plt.imshow(z_map[:, y_slice, :].T, origin='lower', cmap=plt.cm.hot)
    plt.imshow(in_limbo[:, y_slice, :].T, origin='lower', cmap=plt.cm.Greens_r)
    plt.axis('off')    
    
    plt.subplot(133)
    plt.imshow(brain[:, :, z_slice].T, origin='lower', cmap=plt.cm.gray)
    plt.imshow(z_map[:, :, z_slice].T, origin='lower', cmap=plt.cm.hot)
    plt.imshow(in_limbo[:, :, z_slice].T, origin='lower', cmap=plt.cm.Greens_r)
    
    plt.axis('off')
    
    
    ## Sandwich approach
    plt.figure(figsize=(20, 5))
    in_limbo = nb.load('/home/gdholla1/projects/in_limbo/final_results/%s_in_limbo_mask_brain_masked.nii.gz' % threshold).get_data()
    z_map = nb.load('/home/gdholla1/projects/in_limbo/final_results/%s_z_threshold.nii.gz' % threshold).get_data()    
    
    z_map = np.ma.masked_less_equal(z_map, 0)
    in_limbo = np.ma.masked_less_equal(in_limbo, 0)
    
    plt.suptitle('Sandwich- estimation z > %s' % threshold, fontsize=25)
    plt.subplot(131)
    plt.imshow(brain[x_slice, :, :].T, origin='lower', cmap=plt.cm.gray)
    plt.imshow(z_map[x_slice, :, :].T, origin='lower', cmap=plt.cm.hot)
    plt.imshow(in_limbo[x_slice, :, :].T, origin='lower', cmap=plt.cm.Greens_r)
    plt.axis('off')   
    
    plt.subplot(132)
    plt.imshow(brain[:, y_slice, :].T, origin='lower', cmap=plt.cm.gray)
    plt.imshow(z_map[:, y_slice, :].T, origin='lower', cmap=plt.cm.hot)
    plt.imshow(in_limbo[:, y_slice, :].T, origin='lower', cmap=plt.cm.Greens_r)
    plt.axis('off')    
    
    plt.subplot(133)
    plt.imshow(brain[:, :, z_slice].T, origin='lower', cmap=plt.cm.gray)
    plt.imshow(z_map[:, :, z_slice].T, origin='lower', cmap=plt.cm.hot)
    plt.imshow(in_limbo[:, :, z_slice].T, origin='lower', cmap=plt.cm.Greens_r)    
    plt.axis('off')