import numpy as np
import seaborn as sns
from matplotlib import pyplot as plt
import scipy.integrate as integrate

sns.set(style="white")
import os
from glob import glob
import pandas as pd
from sklearn.preprocessing import StandardScaler

# load Essen timeseries data and scrub resting state data
_fd_thr = 0.15
_perc_scrub_thr = 0.5
all_ts = []
root_dir = 'data/ex_timeseries/'
for f in glob(root_dir + 'essen-*.tsv'):  # change to bochum-*.tsv for replication with the Bochum dataset
    path, filename = os.path.split(f)
    ts = pd.read_csv(f, sep='\t').iloc[:, 1:].values
    ts = StandardScaler().fit_transform(ts)
    fd = pd.read_csv(path + '/FD_' + os.path.splitext(filename)[0] + '.txt', sep=',').values.flatten()
    fd = np.hstack(([0], fd))
    if np.sum(fd>_fd_thr)/len(fd) < _perc_scrub_thr:
        all_ts.append(ts[fd<_fd_thr])
    else:
        print('perc. scrubbed:', np.sum(fd>_fd_thr)/len(fd))
len(all_ts)

# comopute the negative precision matrix via nilearn
from nilearn.connectome import ConnectivityMeasure
correlation_measure = ConnectivityMeasure(kind='precision',  vectorize=False, discard_diagonal=False) # uses Ledoit-Wolf estimator
correlation_measure.fit_transform(all_ts)
J = -1 * correlation_measure.mean_

def compute_W_from_J(J, num_attractors=None):
    if num_attractors is None:
        num_attractors = J.shape[0]
    # Eigenvalue decomposition of J
    eigenvalues, eigenvectors = np.linalg.eigh(J)
    # Sort eigenvalues and corresponding eigenvectors in descending order
    idx = np.argsort(eigenvalues)[::-1]
    eigenvalues = eigenvalues[idx]
    eigenvectors = eigenvectors[:, idx]
    # Select the first `num_hidden_nodes` eigenvalues and eigenvectors
    selected_eigenvalues = eigenvalues[:num_attractors]
    selected_eigenvectors = eigenvectors[:, :num_attractors]
    # Compute the weight matrix W
    W = selected_eigenvectors * np.sqrt(-selected_eigenvalues)
    return W

# load taks-based data and scrub resting state data
_fd_thr = 0.15
_perc_scrub_thr = 0.5

timeseries = []
events_heat = []
events_pain = []
events_rating = []

_TR_ = 2.0  # seconds

for sub in range(1, 33):
    sub_pad = f'{sub + 1:02}'

    fd_path = 'data/ex_timeseries_task/sub-' + sub_pad + '_FD.txt'
    fd = pd.read_csv(fd_path).values.flatten()
    fd = np.hstack(([0], fd))

    ts_path = 'data/ex_timeseries_task/sub-' + sub_pad + '_ts.txt'
    ts = pd.read_csv(ts_path, sep='\t').iloc[:, 1:].values

    events_path = 'data/ex_timeseries_task/sub-' + sub_pad + '_task-heatpainwithregulationandratings_run-01_events.tsv'
    event = pd.read_csv(events_path, sep='\t')
    ratings = event['ratings'].values
    event_heatpain = event[~ np.isnan(event['temperature'])].copy()  # remove rating blocks, only retain heat pain blocks
    event_heatpain['ratings'] = ratings[~ np.isnan(ratings)] 

    event_heat = event_heatpain[event_heatpain['ratings']<100]
    event_pain = event_heatpain[event_heatpain['ratings']>=100]
    event_rating = event[np.isnan(event['temperature'])].copy()

    ts = StandardScaler().fit_transform(ts)

    if np.sum(fd > _fd_thr) / len(fd) < _perc_scrub_thr:
        ts[fd >= _fd_thr] = np.nan
    else:  
        print(sub, 'perc. scrubbed:', np.sum(fd > _fd_thr) / len(fd))
        continue

    timeseries.append(ts)
    events_heat.append(event_heat)
    events_pain.append(event_pain)
    events_rating.append(event_rating)

len(timeseries)

#1st level analysis

from tqdm import tqdm
from nilearn.glm.first_level import compute_regressor
import statsmodels.api as sm
import statsmodels.formula.api as smf

J = -1 * correlation_measure.mean_
W = compute_W_from_J(J)

all_proj = []
models = {}

for sub in tqdm(range(len(timeseries))):
    models[sub] = {}

    # attractor timeseries
    proj = W.T @ timeseries[sub].T
    all_proj.append(proj)

    # HRF-convolved EV from the events: heat
    ev = events_heat[sub][['onset', 'duration', 'ratings']].copy()
    ev.loc[:, 'ratings'] = 1 
    hrf_regressor_heat = compute_regressor(ev.values.transpose(),
                                    hrf_model='spm',
                                    frame_times=np.arange(0, len(timeseries[sub][:,0])*_TR_, _TR_))[0].flatten()

    # parametric modulation with heat rating
    ev = events_heat[sub][['onset', 'duration', 'ratings']].copy()
    # center the ratings
    ev.loc[:, 'ratings'] = (ev['ratings'] - np.mean(ev['ratings']))/100
    hrf_regressor_heatint = compute_regressor(ev.values.transpose(),
                                    hrf_model='spm',
                                    frame_times=np.arange(0, len(timeseries[sub][:,0])*_TR_, _TR_))[0].flatten()
    
    # HRF-convolved EV from the events: pain
    ev = events_pain[sub][['onset', 'duration', 'ratings']].copy()
    ev.loc[:, 'ratings'] = 1 
    hrf_regressor_pain = compute_regressor(ev.values.transpose(),
                                    hrf_model='spm',
                                    frame_times=np.arange(0, len(timeseries[sub][:,0])*_TR_, _TR_))[0].flatten()

    # parametric modulation with pain rating    
    ev = events_pain[sub][['onset', 'duration', 'ratings']].copy()
    # center the ratings
    ev.loc[:, 'ratings'] = (ev['ratings'] - np.mean(ev['ratings']))/100
    hrf_regressor_painint = compute_regressor(ev.values.transpose() ,
                                    hrf_model='spm',
                                    frame_times=np.arange(0, len(timeseries[sub][:,0])*_TR_, _TR_))[0].flatten()
    
    # HRF-convolved EV from the events: rating events
    ev = events_rating[sub][['onset', 'duration', 'ratings']].copy()
    ev.loc[:, 'ratings'] = 1 
    hrf_regressor_rating = compute_regressor(ev.values.transpose(),
                                    hrf_model='spm',
                                    frame_times=np.arange(0, len(timeseries[sub][:,0])*_TR_, _TR_))[0].flatten()

    # plot regressors on a single plot
    #plt.figure(figsize=(10, 6))
    #plt.plot(hrf_regressor_heat, label='HRF_Regressor_Heat')
    #plt.plot(hrf_regressor_heatint, label='HRF_Regressor_HeatInt')
    #plt.plot(hrf_regressor_pain, label='HRF_Regressor_Pain')
    #plt.plot(hrf_regressor_painint, label='HRF_Regressor_PainInt')
    #plt.plot(hrf_regressor_rating, label='HRF_Regressor_Rating')
    #plt.legend()
    #plt.show()

    #plot attractor activity
    #plt.figure(figsize=(10, 6))
    #plt.plot(proj[sub, :])
    #plt.show()
    
    # fit 1st level model with statsmodels
    for attr in range(W.shape[0]):

        # Create a DataFrame for the model
        HRF_df = pd.DataFrame({
            'Attractor_Activity': proj[attr, :],
            'HRF_Regressor_Heat': hrf_regressor_heat,
            'HRF_Regressor_HeatInt': hrf_regressor_heatint,
            'HRF_Regressor_Pain': hrf_regressor_pain,  
            'HRF_Regressor_PainInt': hrf_regressor_painint, 
            'HRF_Regressor_Rating': hrf_regressor_rating,
        })

        data.dropna(inplace=True)

        # Define the formula for the model
        formula = 'Attractor_Activity ~ HRF_Regressor_Heat + HRF_Regressor_HeatInt + HRF_Regressor_Pain + HRF_Regressor_PainInt + HRF_Regressor_Rating'
        # Fit the model
        model = smf.ols(formula, data=HRF_df).fit()

        # hack in differece contrast for heat and pain
        params = model.params.copy()
        params['Pain>Heat'] = params['HRF_Regressor_Pain'] - params['HRF_Regressor_Heat']
        # Get the parameter estimates
        models[sub][attr] = params     

from scipy.stats import wilcoxon

second_level_results = {}

for attr in range(W.shape[0]):
    second_level_results[attr] = {}
    # for each event type
    for event_type in models[0][attr].index:
        # skip Intercept
        if event_type == 'Intercept':
            continue
        betas = np.array([models[sub][attr][event_type] for sub in range(len(models))])
        statistic, p_value = wilcoxon(betas)
        second_level_results[attr][event_type] = {'statistic': statistic, 'p_value': p_value}

from nilearn import datasets, surface, plotting
from matplotlib.patches import Rectangle
from connattractor import network
from nilearn.plotting.cm import _cmap_d as nilearn_cmaps

mist_regions = pd.read_csv('data/MIST122_relabeled.tsv', sep='\t')
mist_regions.columns = ['label', 'MIST_index', 'region', 'module']
mist_mapper=mist_regions[['label', 'region']]
mist_mapper.set_index('region', inplace=True)

fsaverage = datasets.fetch_surf_fsaverage('fsaverage6')
surf = surface.load_surf_mesh('data/full_brain_left.surf.gii')
bg_map = np.hstack((surface.load_surf_data(fsaverage.sulc_left),
                    np.zeros(28935))) # magic number: number of cerebellum vertices (SUIT cerebellum template)

for event_type in second_level_results[attr].keys():
    # Set the number of attractors to plot
    k = 1

    # find the k attractors with the lowest p-values
    p_values = []
    for attr in range(W.shape[0]):
        p_value = second_level_results[attr][event_type]['p_value']
        p_values.append((p_value, attr))
    p_values.sort()

    for i in range(k):
        min_p_values, min_attr = p_values[i]
        print('===========================================')
        print(f'** {event_type}, Attractor {min_attr}')
        print(f'{"Attractor":<120} {"Event Type":<25} {"Statistic":<20} {"p-value":<20} {"Mean Beta":<20}')
        for event_types in second_level_results[attr].keys():
            statistics = second_level_results[min_attr][event_types]['statistic']
            mean_betas = np.mean([models[sub][min_attr][event_types] for sub in range(len(models))])
            p_values = second_level_results[min_attr][event_types]['p_value']
            print(f'{min_attr:<20} {event_types:<25} {statistics:<20} {p_values:<20} {mean_betas:<20}')

        # print out the top 5 lowest and highest regions, based on mist_mapper
        # Get the mean activation for each region

        mean_beta = np.mean([models[sub][min_attr][event_type] for sub in range(len(models))])
        
        sign = 1
        region_indices = np.argsort(sign * W[:, min_attr])
        top_5_lowest = mist_regions.loc[region_indices[:5], 'region']
        top_5_highest = mist_regions.loc[region_indices[-5:], 'region']
        print('* Top 5 highest weights:')
        for region in region_indices[-5:]:
            print(f'{mist_regions.loc[region, "region"]}: {W[region, min_attr]}')
        print('* Top 5 lowest weights:')
        for region in region_indices[:5]:
            print(f'{mist_regions.loc[region, "region"]}: {W[region, min_attr]}')


        # surface plot
        #if event_type == 'Rating':
        #    img = network.State(sign * W[:,min_attr]).to_Nifti1Image() # visualization hack
        #else:
        img = network.State(sign * W[:,min_attr]).to_Nifti1Image()
        texture = surface.vol_to_surf(img, surf)
        fig = plotting.plot_surf_stat_map(surf, texture, bg_map=bg_map, bg_on_data=True, 
                                          hemi='left', colorbar=False, cmap='coolwarm', 
                                          threshold=0, symmetric_cbar=True, vmax=0.6,
                                          title='Attractor {}'.format(min_attr))

        fig = plotting.plot_surf_stat_map(surf, texture, bg_map=bg_map, bg_on_data=True, 
                                          hemi='right', colorbar=False, cmap='coolwarm', 
                                          threshold=0, symmetric_cbar=True, vmax=0.6)
        plt.show()

        fig, ax = plt.subplots(figsize=(9, 2))
        # blocks
        for j in events_heat[-1].index:
            ax.add_patch(Rectangle((events_heat[-1]['onset'][j], -2), 
                                        events_heat[-1]['duration'][j],
                                        4, 
                                        alpha=0.2,
                                        linewidth=0,
                                        color='purple'))
            
        for j in events_pain[-1].index:
            ax.add_patch(Rectangle((events_pain[-1]['onset'][j], -2), 
                                        events_pain[-1]['duration'][j],
                                        4, 
                                        alpha=0.2,
                                        linewidth=0,
                                        color='purple'))

        for j in events_rating[-1].index:
            ax.add_patch(Rectangle((events_rating[-1]['onset'][j], -2), 
                                        events_rating[-1]['duration'][j],
                                        4, 
                                        alpha=0.2,
                                        linewidth=0,
                                        color='green'))

        ax.set_xlabel('Time (s)')
        ax.set_ylabel('Attractor activation')

        # line plot
        df = pd.DataFrame(columns=['time', 'subject', 'activation'])
        for sub in range(0, len(timeseries)):
            # append this and the time to the df:
            # sign * all_proj[sub][strongest_attractors[i],:]
            df = pd.concat([df, pd.DataFrame({'time': np.arange(0, len(timeseries[sub][:,0])*_TR_, _TR_),
                                              'subject': sub,
                                              'activation': sign * all_proj[sub][min_attr,:]})])
            df.reset_index(drop=True, inplace=True)

        sns.lineplot(data=df, x='time', y='activation', ax=ax)
        sns.despine(ax=ax)
        ax.set_ylim(-3, 3)
        plt.savefig(f'fig/task-{event_type}_attractor-{min_attr}_lineplot.pdf')
        plt.show()

        if event_type != 'Pain>Heat':
            slopes = np.array([sign * models[sub][min_attr][event_type] for sub in range(len(models))])
        else:
            slopes = np.array([models[sub][min_attr]['HRF_Regressor_Pain'] for sub in range(len(models))])
            slopes -= np.array([models[sub][min_attr]['HRF_Regressor_Heat'] for sub in range(len(models))])
            slopes *= sign

        fig, ax = plt.subplots(figsize=(2, 2))
        limit = 0.8 #np.max(np.abs(slopes[sub]))*0.8
        norm = plt.Normalize(-limit, limit)
        cmap = matplotlib.colors.LinearSegmentedColormap.from_list("", ["green", "purple"])
        
        for sub, slope in enumerate(slopes):
            intercept = models[sub][min_attr]['Intercept']
            if event_type != 'Pain>Heat':
                x_vals = HRF_df[event_type]
            else:
                x_vals = HRF_df['HRF_Regressor_Pain'] - HRF_df['HRF_Regressor_Heat']
            y_vals = slope * x_vals + intercept
            ax.plot(x_vals, y_vals, alpha=0.2, color=cmap(norm(np.sign(slopes[sub]))), linewidth=2)
        
        # plot the mean slope and mean intercept
        mean_slope = np.mean(slopes)
        if event_type != 'Pain>Heat':
            mean_intercepts = sign * np.mean(np.array([models[sub][min_attr]['Intercept'] for sub in range(len(models))]))
        else:
            mean_intercepts = sign * np.mean(np.array([models[sub][min_attr]['HRF_Regressor_Pain'] for sub in range(len(models))]))
            mean_intercepts -= sign * np.mean(np.array([models[sub][min_attr]['HRF_Regressor_Heat'] for sub in range(len(models))]))
        x_vals = np.array(ax.get_xlim())
        y_vals = mean_slope * x_vals + mean_intercepts
        ax.plot(x_vals, y_vals, color='black', linewidth=4)
        

        ax.set_ylim(-2, 2)
        plt.xlabel(event_type)
        plt.ylabel('Activation')
        sns.despine(ax=ax)
        plt.savefig(f'fig/task-{event_type}_attractor-{min_attr}_slopeplot.pdf')
        plt.show()
        
        #ax.set_xlabel('Time (s)')
        #ax.set_ylabel('Attractor Activity')
        #ax.set_title(f'Attractor {min_attr} Activity')