from pyitab.io.loader import DataLoader
from pyitab.analysis.linear_model import LinearModel
from pyitab.preprocessing.pipelines import PreprocessingPipeline

from pyitab.preprocessing.normalizers import FeatureZNormalizer
from pyitab.preprocessing.functions import SampleAttributeTransformer, TargetTransformer
from pyitab.preprocessing.slicers import SampleSlicer

from pyitab.plot.connectivity import plot_connectivity_matrix
from scipy.stats import zscore

import numpy as np
import warnings
warnings.filterwarnings("ignore")

data_path = '/media/robbis/DATA/meg/viviana-hcp/'
conf_file = "/media/robbis/DATA/meg/viviana-hcp/bids.conf"

loader = DataLoader(configuration_file=conf_file,
                    data_path=data_path,
                    subjects="/media/robbis/DATA/meg/viviana-hcp/participants.tsv",
                    loader='bids-meg',
                    task='blp',
                    bids_atlas="complete",
                    bids_correction="corr",
                    bids_derivatives='True',
                    load_fx='hcp-blp')


ds = loader.fetch()

nodes = ds.fa.nodes_1
matrix = np.zeros_like(ds.samples[0])

nanmask = np.logical_not(np.isnan(ds.samples).sum(0))
ds = ds[:, nanmask]

def plot_stats(stat, matrix=matrix, nanmask=nanmask, nodes=nodes, title='stat'):
    matrix[nanmask] = stat
    _, a = plot_connectivity_matrix(matrix, 
                                    networks=nodes, 
                                    cmap=pl.cm.viridis, 
                                    vmin=0.)

    a.set_title(title)


# Transform dataset to have mean 0 and std 1
prepro = [
    SampleSlicer(task=['rest', 'task1', 'task2', 'task4', 'task5']),
    FeatureZNormalizer(),
    SampleAttributeTransformer(attr='dexterity1', fx=('zscore', zscore)),
    SampleAttributeTransformer(attr='dexterity2', fx=('zscore', zscore)),
]

ds = PreprocessingPipeline(nodes=prepro).transform(ds)

##################

# 1. Full model (band + task + subject + dexterity)

n_bands = len(np.unique(ds.sa.mainband))
n_tasks = len(np.unique(ds.sa.maintask))
n_subjects = len(np.unique(ds.sa.subject))
dexterity = 1

band_contrast = np.zeros((n_bands-1, n_bands+n_tasks+n_subjects+1))
band_contrast[0, 1:4] = [1, -1, 0]
band_contrast[1, 1:4] = [0, 1, -1]
band_contrast[:, 4:7] = 1./n_tasks
band_contrast[:, 7:] = 1./n_subjects

task_contrast = np.zeros((n_tasks-1, n_bands+n_tasks+n_subjects+1))
task_contrast[0, 4:5] = [1, -1]
task_contrast[1, 5:7] = [1, -1]

index = 7
subject_contrast = np.zeros((n_subjects-1, n_bands+n_tasks+n_subjects+1))
for i in range(n_subjects-1):
    subject_contrast[i, index:index+2] = [1, -1]
    index += 1




contrasts = {
    't+dexterity': np.hstack([1, np.zeros(n_bands), np.zeros(n_tasks), np.zeros(n_subjects)]),
    'f+band': band_contrast,
    'f+task': task_contrast,
    'f+subject': subject_contrast,
}

lm = LinearModel(attr=['mainband', 'maintask', 'subject', 'dexterity1'])
lm.fit(ds, full_model=True)

lm._contrast(contrast=contrasts)

plot_stats(lm.scores.r_square, title='r2')


# stats
stats = lm.scores.stats_contrasts
for contrast, stats in lm.scores.stats_contrasts.items():
    
    if contrast[0] == 'f':
        test = 'F'
    else:
        test = 't'
    
    s = stats[test]
    p = stats['p_values']
    t = 0.001 / s.shape[0]
    s[p > t] = 0

    plot_stats(s, title=contrast)

#######################################################
# 2. Modulation of tasks within bands.

contrasts = {

    #'t+restvstask': [1, -1/4, -1/4, -1/4, -1/4, 0],
    'f+restvstask': [1, -1/4, -1/4, -1/4, -1/4, 0],

    'f+task':       [[1,-1, 0, 0, 0, 0], 
                     [0, 1,-1, 0, 0, 0],
                     [0, 0, 1,-1, 0, 0],
                     [0, 0, 0, 1,-1, 0]],
    't+rest':  [1, 0, 0, 0, 0, 0],
    #'t+1task': [0, 1, 0, 0, 0, 0],
    #'t+2task': [0, 0, 1, 0, 0, 0],
    #'t+4task': [0, 0, 0, 1, 0, 0],
    #'t+5task': [0, 0, 0, 0, 1, 0],
    
    
    #'t+handvsfoot': [0, 1/2, -1/2, 1/2, -1/2, 0],
    #'f+handvsfoot': [0, 1/2, -1/2, 1/2, -1/2, 0],

    #'t+movement': [0, 1/4, 1/4, 1/4, 1/4, 0],
    #'f+movement': [0, 1/4, 1/4, 1/4, 1/4, 0],

    't+dexterity':  [0, 0, 0, 0, 0, 1],
    #'t+taskvsdext': [1/5, 1/5, 1/5, 1/5, 1/5, -1],
    #'f+taskvsdext': [1/5, 1/5, 1/5, 1/5, 1/5, -1],
}
import seaborn as sns
color2 = "#F21A00"
color1 = "#3B9AB2"
tpalette = sns.blend_palette([color1, "#EEEEEE", color2], n_colors=100, as_cmap=True)
fpalette = sns.blend_palette(["#EEEEEE", color2], n_colors=100, as_cmap=True)


for band in ['alpha', 'beta', 'gamma']:
    ds_ = SampleSlicer(mainband=[band]).transform(ds)
    ds_ = PreprocessingPipeline(nodes=prepro).transform(ds_)
    lm = LinearModel(attr=['task', 'dexterity1'])
    
    lm.fit(ds_, formula='task + dexterity1 - 1')
    lm._contrast(contrast=contrasts)

    title = band
    """
    r2 = lm.scores.r_square
    matrix[nanmask] = r2
    _, a = plot_connectivity_matrix(matrix, 
                                    networks=nodes, 
                                    cmap=pl.cm.viridis, 
                                    vmin=0.)

    a.set_title(band+" | r2")
    """
    # stats
    stats = lm.scores.stats_contrasts
    for contrast, stats in lm.scores.stats_contrasts.items():
        
        if contrast[0] == 'f':
            test = 'F'
            cmap = fpalette
            vmin = 0
        else:
            test = 't'
            cmap = tpalette

        s = stats[test]
        p = stats['p_values']
        t = 0.05 / s.shape[0]

        if contrast[0] == 't':
            vmin = -1*np.max(np.abs(s))

        tmask = np.logical_and(p>t, p<(1-t))
        s[tmask] = 0

        matrix[nanmask] = s
        _, a = plot_connectivity_matrix(matrix, 
                                        networks=nodes, 
                                        cmap=cmap,
                                        vmin=vmin)

        a.set_title(title+" | "+contrast)


####
# 3. Tests
lm = LinearModel(attr=['mainband', 'maintask', 'dexterity1'])
lm.fit(ds, formula='mainband:task + dexterity1 - 1')


contrast = np.zeros(16)

contrasts = {

    't+alpha': [1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0],
    't+beta':  [0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0],
    't+gamma': [0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0],

    't+arestvstask': [1, 0, 0, -1/4, 0, 0, -1/4, 0, 0, -1/4, 0, 0, -1/4, 0, 0, 0],
    't+brestvstask': [0, 1, 0, 0, -1/4, 0, 0, -1/4, 0, 0, -1/4, 0, 0, -1/4, 0, 0],
    't+grestvstask': [0, 0, 1, 0, 0, -1/4, 0, 0, -1/4, 0, 0, -1/4, 0, 0, -1/4, 0],

    't+dexterity1': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]
}

lm._contrast(contrast=contrasts)

stats = lm.scores.stats_contrasts
for contrast, stats in lm.scores.stats_contrasts.items():
    
    if contrast[0] == 'f':
        test = 'F'
        cmap = fpalette
        vmin = 0
    else:
        test = 't'
        cmap = tpalette

    s = stats[test]
    p = stats['p_values']
    t = 0.05 / s.shape[0]

    if contrast[0] == 't':
        vmin = -1*np.max(np.abs(s))

    tmask = np.logical_and(p>t, p<(1-t))
    s[tmask] = 0

    matrix[nanmask] = s
    _, a = plot_connectivity_matrix(matrix, 
                                    networks=nodes, 
                                    cmap=cmap,
                                    vmin=vmin)

    a.set_title(contrast)

####
# 4. ICC
from pingouin import intraclass_corr
import pandas as pd
from joblib import Parallel, delayed


def regression(ds, i, attrs):

    values = [ds.samples[:, i]]
    keys = ['y']

    for a in attrs:
        values.append(ds.sa[a].value)
        keys.append(a)


    df = pd.DataFrame(np.vstack(values).T, 
                      columns=keys)

    df['y'] = np.float_(df['y'])
    df['y'] -= df['y'].mean()
    df['y'] /= df['y'].std()

    return df

def icc(ds, i, attr):

    df = regression(ds, i, attr)
    icc_index = intraclass_corr(df, targets='subject', raters='task', ratings='y')

    return icc_index['ICC'][0]


icc_ = {}

for band in ['alpha', 'beta', 'gamma']:
    ds_ = SampleSlicer(mainband=[band]).transform(ds)
    ds_ = PreprocessingPipeline(nodes=prepro).transform(ds_)

    icc_[band] = Parallel(n_jobs=-1, verbose=1)\
        (delayed(icc)(ds_, i, attr=['task', 'subject']) for i in range(ds_.shape[1]))
    
for band, score in icc_.items():

    matrix[nanmask] = score
    _, a = plot_connectivity_matrix(matrix, 
                                    networks=nodes, 
                                    cmap=pl.cm.viridis, 
                                    vmin=0.)

    a.set_title(band+" | icc")
    

# 4.1 Explain 





#####
# 5. Full model with uncollapsed bands

lm = LinearModel(attr=['band', 'task', 'subject', 'dexterity1'])
lm.fit(ds, full_model=True)

theta_subj = np.zeros_like(lm.scores.theta[0, :])
for i, variable in enumerate(lm.scores.design_info):

    #theta = lm.scores.theta[:, i]

    if 'subject' in variable:
        theta_subj += lm.scores.theta[i, :]

    matrix[nanmask] = theta_subj
    _, a = plot_connectivity_matrix(matrix, 
                                    networks=nodes, 
                                    cmap=pl.cm.bwr)

    a.set_title(variable)
from pyitab.io.loader import DataLoader
from pyitab.analysis.linear_model import LinearModel
from pyitab.preprocessing.pipelines import PreprocessingPipeline

from pyitab.preprocessing.normalizers import FeatureZNormalizer
from pyitab.preprocessing.functions import SampleAttributeTransformer, TargetTransformer
from pyitab.preprocessing.slicers import SampleSlicer

from pyitab.plot.connectivity import plot_connectivity_matrix
from scipy.stats import zscore

import numpy as np
import warnings
warnings.filterwarnings("ignore")

data_path = '/media/robbis/DATA/meg/viviana-hcp/'
conf_file = "/media/robbis/DATA/meg/viviana-hcp/bids.conf"

loader = DataLoader(configuration_file=conf_file,
                    data_path=data_path,
                    subjects="/media/robbis/DATA/meg/viviana-hcp/participants.tsv",
                    loader='bids-meg',
                    task='blp',
                    bids_atlas="complete",
                    bids_correction="corr",
                    bids_derivatives='True',
                    load_fx='hcp-blp')


ds = loader.fetch()

nodes = ds.fa.nodes_1
matrix = np.zeros_like(ds.samples[0])

nanmask = np.logical_not(np.isnan(ds.samples).sum(0))
ds = ds[:, nanmask]

def plot_stats(stat, matrix=matrix, nanmask=nanmask, nodes=nodes, title='stat'):
    matrix[nanmask] = stat
    _, a = plot_connectivity_matrix(matrix, 
                                    networks=nodes, 
                                    cmap=pl.cm.viridis, 
                                    vmin=0.)

    a.set_title(title)


# Transform dataset to have mean 0 and std 1
prepro = [
    SampleSlicer(task=['rest', 'task1', 'task2', 'task4', 'task5']),
    FeatureZNormalizer(),
    SampleAttributeTransformer(attr='dexterity1', fx=('zscore', zscore)),
    SampleAttributeTransformer(attr='dexterity2', fx=('zscore', zscore)),
]

ds = PreprocessingPipeline(nodes=prepro).transform(ds)

##################

# 1. Full model (band + task + subject + dexterity)

n_bands = len(np.unique(ds.sa.mainband))
n_tasks = len(np.unique(ds.sa.maintask))
n_subjects = len(np.unique(ds.sa.subject))
dexterity = 1

band_contrast = np.zeros((n_bands-1, n_bands+n_tasks+n_subjects+1))
band_contrast[0, 1:4] = [1, -1, 0]
band_contrast[1, 1:4] = [0, 1, -1]
band_contrast[:, 4:7] = 1./n_tasks
band_contrast[:, 7:] = 1./n_subjects

task_contrast = np.zeros((n_tasks-1, n_bands+n_tasks+n_subjects+1))
task_contrast[0, 4:5] = [1, -1]
task_contrast[1, 5:7] = [1, -1]

index = 7
subject_contrast = np.zeros((n_subjects-1, n_bands+n_tasks+n_subjects+1))
for i in range(n_subjects-1):
    subject_contrast[i, index:index+2] = [1, -1]
    index += 1




contrasts = {
    't+dexterity': np.hstack([1, np.zeros(n_bands), np.zeros(n_tasks), np.zeros(n_subjects)]),
    'f+band': band_contrast,
    'f+task': task_contrast,
    'f+subject': subject_contrast,
}

lm = LinearModel(attr=['mainband', 'maintask', 'subject', 'dexterity1'])
lm.fit(ds, full_model=True)

lm._contrast(contrast=contrasts)

plot_stats(lm.scores.r_square, title='r2')


# stats
stats = lm.scores.stats_contrasts
for contrast, stats in lm.scores.stats_contrasts.items():
    
    if contrast[0] == 'f':
        test = 'F'
    else:
        test = 't'
    
    s = stats[test]
    p = stats['p_values']
    t = 0.001 / s.shape[0]
    s[p > t] = 0

    plot_stats(s, title=contrast)

#######################################################
# 2. Modulation of tasks within bands.

contrasts = {

    #'t+restvstask': [1, -1/4, -1/4, -1/4, -1/4, 0],
    'f+restvstask': [1, -1/4, -1/4, -1/4, -1/4, 0],

    'f+task':       [[1,-1, 0, 0, 0, 0], 
                     [0, 1,-1, 0, 0, 0],
                     [0, 0, 1,-1, 0, 0],
                     [0, 0, 0, 1,-1, 0]],
    't+rest':  [1, 0, 0, 0, 0, 0],
    #'t+1task': [0, 1, 0, 0, 0, 0],
    #'t+2task': [0, 0, 1, 0, 0, 0],
    #'t+4task': [0, 0, 0, 1, 0, 0],
    #'t+5task': [0, 0, 0, 0, 1, 0],
    
    
    #'t+handvsfoot': [0, 1/2, -1/2, 1/2, -1/2, 0],
    #'f+handvsfoot': [0, 1/2, -1/2, 1/2, -1/2, 0],

    #'t+movement': [0, 1/4, 1/4, 1/4, 1/4, 0],
    #'f+movement': [0, 1/4, 1/4, 1/4, 1/4, 0],

    't+dexterity':  [0, 0, 0, 0, 0, 1],
    #'t+taskvsdext': [1/5, 1/5, 1/5, 1/5, 1/5, -1],
    #'f+taskvsdext': [1/5, 1/5, 1/5, 1/5, 1/5, -1],
}
import seaborn as sns
color2 = "#F21A00"
color1 = "#3B9AB2"
tpalette = sns.blend_palette([color1, "#EEEEEE", color2], n_colors=100, as_cmap=True)
fpalette = sns.blend_palette(["#EEEEEE", color2], n_colors=100, as_cmap=True)


for band in ['alpha', 'beta', 'gamma']:
    ds_ = SampleSlicer(mainband=[band]).transform(ds)
    ds_ = PreprocessingPipeline(nodes=prepro).transform(ds_)
    lm = LinearModel(attr=['task', 'dexterity1'])
    
    lm.fit(ds_, formula='task + dexterity1 - 1')
    lm._contrast(contrast=contrasts)

    title = band
    """
    r2 = lm.scores.r_square
    matrix[nanmask] = r2
    _, a = plot_connectivity_matrix(matrix, 
                                    networks=nodes, 
                                    cmap=pl.cm.viridis, 
                                    vmin=0.)

    a.set_title(band+" | r2")
    """
    # stats
    stats = lm.scores.stats_contrasts
    for contrast, stats in lm.scores.stats_contrasts.items():
        
        if contrast[0] == 'f':
            test = 'F'
            cmap = fpalette
            vmin = 0
        else:
            test = 't'
            cmap = tpalette

        s = stats[test]
        p = stats['p_values']
        t = 0.05 / s.shape[0]

        if contrast[0] == 't':
            vmin = -1*np.max(np.abs(s))

        tmask = np.logical_and(p>t, p<(1-t))
        s[tmask] = 0

        matrix[nanmask] = s
        _, a = plot_connectivity_matrix(matrix, 
                                        networks=nodes, 
                                        cmap=cmap,
                                        vmin=vmin)

        a.set_title(title+" | "+contrast)


####
# 3. Tests
lm = LinearModel(attr=['mainband', 'maintask', 'dexterity1'])
lm.fit(ds, formula='mainband:task + dexterity1 - 1')


contrast = np.zeros(16)

contrasts = {

    't+alpha': [1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0],
    't+beta':  [0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0],
    't+gamma': [0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0],

    't+arestvstask': [1, 0, 0, -1/4, 0, 0, -1/4, 0, 0, -1/4, 0, 0, -1/4, 0, 0, 0],
    't+brestvstask': [0, 1, 0, 0, -1/4, 0, 0, -1/4, 0, 0, -1/4, 0, 0, -1/4, 0, 0],
    't+grestvstask': [0, 0, 1, 0, 0, -1/4, 0, 0, -1/4, 0, 0, -1/4, 0, 0, -1/4, 0],

    't+dexterity1': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]
}

lm._contrast(contrast=contrasts)

stats = lm.scores.stats_contrasts
for contrast, stats in lm.scores.stats_contrasts.items():
    
    if contrast[0] == 'f':
        test = 'F'
        cmap = fpalette
        vmin = 0
    else:
        test = 't'
        cmap = tpalette

    s = stats[test]
    p = stats['p_values']
    t = 0.05 / s.shape[0]

    if contrast[0] == 't':
        vmin = -1*np.max(np.abs(s))

    tmask = np.logical_and(p>t, p<(1-t))
    s[tmask] = 0

    matrix[nanmask] = s
    _, a = plot_connectivity_matrix(matrix, 
                                    networks=nodes, 
                                    cmap=cmap,
                                    vmin=vmin)

    a.set_title(contrast)

####
# 4. ICC
from pingouin import intraclass_corr
import pandas as pd
from joblib import Parallel, delayed


def regression(ds, i, attrs):

    values = [ds.samples[:, i]]
    keys = ['y']

    for a in attrs:
        values.append(ds.sa[a].value)
        keys.append(a)


    df = pd.DataFrame(np.vstack(values).T, 
                      columns=keys)

    df['y'] = np.float_(df['y'])
    df['y'] -= df['y'].mean()
    df['y'] /= df['y'].std()

    return df

def icc(ds, i, attr):

    df = regression(ds, i, attr)
    icc_index = intraclass_corr(df, targets='subject', raters='task', ratings='y')

    return icc_index['ICC'][0]


icc_ = {}

for band in ['alpha', 'beta', 'gamma']:
    ds_ = SampleSlicer(mainband=[band]).transform(ds)
    ds_ = PreprocessingPipeline(nodes=prepro).transform(ds_)

    icc_[band] = Parallel(n_jobs=-1, verbose=1)\
        (delayed(icc)(ds_, i, attr=['task', 'subject']) for i in range(ds_.shape[1]))
    
for band, score in icc_.items():

    matrix[nanmask] = score
    _, a = plot_connectivity_matrix(matrix, 
                                    networks=nodes, 
                                    cmap=pl.cm.viridis, 
                                    vmin=0.)

    a.set_title(band+" | icc")
    

# 4.1 Explain 





#####
# 5. Full model with uncollapsed bands

lm = LinearModel(attr=['band', 'task', 'subject', 'dexterity1'])
lm.fit(ds, full_model=True)

theta_subj = np.zeros_like(lm.scores.theta[0, :])
for i, variable in enumerate(lm.scores.design_info):

    #theta = lm.scores.theta[:, i]

    if 'subject' in variable:
        theta_subj += lm.scores.theta[i, :]

    matrix[nanmask] = theta_subj
    _, a = plot_connectivity_matrix(matrix, 
                                    networks=nodes, 
                                    cmap=pl.cm.bwr)

    a.set_title(variable)
# Transform dataset to have mean 0 and std 1 ################## # 1. Full model (band + task + subject + dexterity) # stats ####################################################### # 2. Modulation of tasks within bands. #'t+restvstask': [1, -1/4, -1/4, -1/4, -1/4, 0], #'t+1task': [0, 1, 0, 0, 0, 0], #'t+2task': [0, 0, 1, 0, 0, 0], #'t+4task': [0, 0, 0, 1, 0, 0], #'t+5task': [0, 0, 0, 0, 1, 0], #'t+handvsfoot': [0, 1/2, -1/2, 1/2, -1/2, 0], #'f+handvsfoot': [0, 1/2, -1/2, 1/2, -1/2, 0], #'t+movement': [0, 1/4, 1/4, 1/4, 1/4, 0], #'f+movement': [0, 1/4, 1/4, 1/4, 1/4, 0], #'t+taskvsdext': [1/5, 1/5, 1/5, 1/5, 1/5, -1], #'f+taskvsdext': [1/5, 1/5, 1/5, 1/5, 1/5, -1], r2 = lm.scores.r_square     matrix[nanmask] = r2     _, a = plot_connectivity_matrix(matrix,                                      networks=nodes,                                      cmap=pl.cm.viridis,                                      vmin=0.)      a.set_title(band+" | r2") # stats #### # 3. Tests #### # 4. ICC # 4.1 Explain ##### # 5. Full model with uncollapsed bands #theta = lm.scores.theta[:, i]