#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on 2020-07-20
@author: Anne Urai
"""

import os
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import patsy # to build design matrix
import statsmodels.api as sm
from ibl_pipeline import behavior, subject, reference
from paper_behavior_functions import (query_sessions_around_criterion,
                                      seaborn_style, institution_map,
                                      group_colors, figpath, EXAMPLE_MOUSE,
                                      FIGURE_WIDTH, FIGURE_HEIGHT,
                                      dj2pandas, plot_psychometric)

# Load some things from paper_behavior_functions
figpath = figpath()
seaborn_style()
institution_map, col_names = institution_map()
pal = group_colors()
#cmap = sns.diverging_palette(20, 220, n=3, center="dark")
cmap = sns.color_palette([[0.8984375,0.37890625,0.00390625],
                          [0.3, 0.3, 0.3], [0.3671875,0.234375,0.59765625]])

# ========================================== #
#%% 1. LOAD DATA - just from example mouse
# ========================================== #

# Query sessions: before and after full task was first introduced
use_sessions, _ = query_sessions_around_criterion(criterion='biased',
                                                  days_from_criterion=[2, 3],
                                                  as_dataframe=False,
                                                  force_cutoff=True)
use_sessions = (subject.Subject & 'subject_nickname = "%s"' % EXAMPLE_MOUSE) * use_sessions

trial_fields = ('trial_stim_contrast_left', 'trial_stim_contrast_right',
                'trial_response_time', 'trial_stim_prob_left',
                'trial_feedback_type', 'trial_stim_on_time', 'trial_response_choice')

# query trial data for sessions and subject name and lab info
trials = use_sessions.proj('task_protocol') * behavior.TrialSet.Trial.proj(*trial_fields)

# only grab the example mouse
subject_info = (subject.Subject) * \
               (subject.SubjectLab * reference.Lab).proj('institution_short')

# Fetch, join and sort data as a pandas DataFrame
behav = dj2pandas(trials.fetch(format='frame')
                         .join(subject_info.fetch(format='frame'))
                         .sort_values(by=['institution_short', 'subject_nickname',
                                          'session_start_time', 'trial_id'])
                         .reset_index())
# split the two types of task protocols (remove the pybpod version number)
behav['task'] = behav['task_protocol'].str[14:20].copy()

# RECODE SOME THINGS JUST FOR PATSY
behav['contrast'] = np.abs(behav.signed_contrast)
behav['stimulus_side'] = np.sign(behav.signed_contrast)
behav['block_id'] = behav['probabilityLeft'].map({80:-1, 50:0, 20:1})

# ========================================== #
#%% 2. DEFINE THE GLM
# ========================================== #


# DEFINE THE MODEL
def fit_glm(behav, prior_blocks=False, n_sim=10000):

    # drop trials with contrast-level 50, only rarely present (should not be its own regressor)
    behav = behav[np.abs(behav.signed_contrast) != 50]

    # use patsy to easily build design matrix
    if not prior_blocks:
        behav = behav.dropna(subset=['trial_feedback_type', 'choice',
                                  'previous_choice', 'previous_outcome']).reset_index()
        endog, exog = patsy.dmatrices('choice ~ 1 + stimulus_side:C(contrast, Treatment)'
                                      '+ previous_choice:C(previous_outcome)',
                               data=behav, return_type='dataframe')
    else:
        behav = behav.dropna(subset=['trial_feedback_type', 'choice',
                                  'previous_choice', 'previous_outcome', 'block_id']).reset_index()
        endog, exog = patsy.dmatrices('choice ~ 1 + stimulus_side:C(contrast, Treatment)'
                                      '+ previous_choice:C(previous_outcome) '
                                      '+ block_id',
                               data=behav, return_type='dataframe')

    # remove the one column (with 0 contrast) that has no variance
    if 'stimulus_side:C(contrast, Treatment)[0.0]' in exog.columns:
        exog.drop(columns=['stimulus_side:C(contrast, Treatment)[0.0]'], inplace=True)

    # recode choices for logistic regression
    endog['choice'] = endog['choice'].map({-1:0, 1:1})

    # rename columns
    exog.rename(columns={'Intercept': 'bias',
              'stimulus_side:C(contrast, Treatment)[6.25]': '6.25',
             'stimulus_side:C(contrast, Treatment)[12.5]': '12.5',
             'stimulus_side:C(contrast, Treatment)[25.0]': '25',
             'stimulus_side:C(contrast, Treatment)[50.0]': '50',
             'stimulus_side:C(contrast, Treatment)[100.0]': '100',
             'previous_choice:C(previous_outcome)[-1.0]': 'unrewarded',
             'previous_choice:C(previous_outcome)[1.0]': 'rewarded'},
             inplace=True)

    # NOW FIT THIS WITH STATSMODELS - ignore NaN choices
    logit_model = sm.Logit(endog, exog)
    res = logit_model.fit_regularized(disp=False)  # run silently

    # what do we want to keep?
    params = pd.DataFrame(res.params).T

    # USE INVERSE HESSIAN TO CONSTRUCT MULTIVARIATE GAUSSIAN
    cov = -np.linalg.inv(logit_model.hessian(res.params))
    samples = np.random.multivariate_normal(res.params, cov, n_sim)

    # sanity check: the mean of those samples should not be too different from the params
    assert(np.allclose(params, np.mean(samples, axis=0), atol=0.1))

    # NOW SIMULATE THE MODEL X TIMES
    simulated_choices = []
    for n in range(n_sim):
        # plug sampled parameters into the model - predict sequence of choices
        z = np.dot(exog, samples[n])

        # then compute the mean choice fractions at each contrast, save and append
        behav['simulated_choice'] = 1 / (1 + np.exp(-z))
        if not prior_blocks:
            simulated_choices.append(behav.groupby(['signed_contrast'])['simulated_choice'].mean().values)
        else: # split by probabilityLeft block
            gr = behav.groupby(['probabilityLeft', 'signed_contrast'])['simulated_choice'].mean().reset_index()
            simulated_choices.append([gr.loc[gr.probabilityLeft == 20, 'simulated_choice'].values,
                                      gr.loc[gr.probabilityLeft == 50, 'simulated_choice'].values,
                                      gr.loc[gr.probabilityLeft == 80, 'simulated_choice'].values])

    return params, simulated_choices  # wide df

# ========================================== #
#%% 3. FIT
# ========================================== #

print('fitting GLM to BASIC task...')
params_basic, simulation_basic = fit_glm(behav.loc[behav.task == 'traini', :],
                                         prior_blocks=False)

print('fitting GLM to FULL task...')
params_full, simulation_full = fit_glm(behav.loc[behav.task == 'biased', :],
                                       prior_blocks=True)

# ========================================== #
#%% 4. PLOT PSYCHOMETRIC FUNCTIONS
# ========================================== #

# for plotting, replace 100 with -35
behav['signed_contrast'] = behav['signed_contrast'].replace(-100, -35)
behav['signed_contrast'] = behav['signed_contrast'].replace(100, 35)

# BASIC TASK
plt.close('all')
# prep the figure with psychometric layout
fig = sns.FacetGrid(behav.loc[behav.task == 'traini', :],
                    sharex=True, sharey=True,
                    height=FIGURE_HEIGHT, aspect=(FIGURE_WIDTH/4)/FIGURE_HEIGHT)
fig.map(plot_psychometric, "signed_contrast", "choice_right", "subject_nickname",
        color='k', linewidth=0) # this will be empty, hack
# now plot the datapoints, no errorbars
sns.lineplot(data=behav.loc[behav.task == 'traini', :],
             x='signed_contrast', y='choice2', marker='o', err_style='bars',
             color='k', linewidth=0, ci=95, ax=fig.ax)# overlay the simulated
# confidence intervals from the model - shaded regions
fig.ax.fill_between(sorted(behav.signed_contrast.unique()),
                    np.quantile(np.array(simulation_basic), q=0.025, axis=0),
                    np.quantile(np.array(simulation_basic), q=0.975, axis=0),
                    alpha=0.5, facecolor='k')
fig.set_axis_labels(' ', 'Rightward choices (%)')
fig.despine(trim=True)
fig.savefig(os.path.join(figpath, "figure5b_basic_psychfunc.pdf"))

# FULL TASK
plt.close('all')
fig = sns.FacetGrid(behav.loc[behav.task == 'biased', :],
                    hue="probabilityLeft", palette=cmap,
                    sharex=True, sharey=True,
                    height=FIGURE_HEIGHT, aspect=(FIGURE_WIDTH/4)/FIGURE_HEIGHT)
fig.map(plot_psychometric, "signed_contrast", "choice_right", "subject_nickname", linewidth=0) # just for axis layout,
# hack
# now plot the datapoints, no errorbars
sns.lineplot(data=behav.loc[behav.task == 'biased', :],
             x='signed_contrast', y='choice2', marker='o', err_style='bars',
             hue='probabilityLeft', palette=cmap, linewidth=0, ci=95, ax=fig.ax, legend=None)# overlay the simulated
# confidence intervals from the model - shaded regions
for cidx, c in enumerate(cmap):
    simulation_full_perblock = [sim[cidx] for sim in simulation_full] # grab what we need, not super elegant
    fig.ax.fill_between(sorted(behav.signed_contrast.unique()),
                    np.quantile(np.array(simulation_full_perblock), q=0.025, axis=0),
                    np.quantile(np.array(simulation_full_perblock), q=0.975, axis=0),
                    alpha=0.5, facecolor=cmap[cidx])

fig.ax.annotate('20:80', xy=(-5, 0.6), xytext=(-25, 0.8), color=cmap[0], fontsize=7)
fig.ax.annotate('80:20', xy=(5, 0.4), xytext=(13, 0.18), color=cmap[2], fontsize=7)

fig.set_axis_labels('\u0394 Contrast (%)', 'Rightward choices (%)')
fig.despine(trim=True)
fig.savefig(os.path.join(figpath, "figure5b_full_psychfunc.pdf"))