import os
import pandas as pd
import mrestimator as mre
import numpy as np
import matplotlib.pylab as plt
import seaborn as sns
import matplotlib.gridspec as gridspec
from scipy import stats as sps
import mr_utilities


def cm2inch(value):
    return value/2.54


mycol = ['darkgreen', 'mediumaquamarine',  'slategray', 'darkblue','steelblue',
         'firebrick', 'purple', 'orange', "red", "violet", "darkred",
         "darkviolet", "green", "salmon", "black", 'darkgreen',
         'mediumaquamarine',  'silver', 'darkblue','steelblue', 'firebrick',
         'purple', 'orange', "red", "violet", "darkred", "darkviolet"]


def distribution_A_recordings(resultpath):
    data_path = '../Data/preictal/'
    recordings = [dI for dI in os.listdir(data_path) if os.path.isdir(
        os.path.join(data_path, dI))]

    A_mean = []
    A_nonzero = []
    A_nonzero_frac = []

    for recording in recordings:
        outputfile = '{}{}/pkl/activity_SU_4ms.pkl'.format(data_path,
                                                           recording)

        if not os.path.isfile(outputfile):
            print("No binning data for recording ", recording)
            continue

        data = pd.read_pickle(outputfile)
        binnings = data['binning'].tolist()
        for binning in binnings:
            activity = data[data['binning'] == binning]['activity'].tolist()[0]
            A_mean.append(np.mean(activity))
            A_nonzero.append(mr_utilities.count_nonzero(activity))
            A_nonzero_frac.append(mr_utilities.count_nonzero(activity) /
                                  len(activity) * 100)

    A_mean_Hz = [A_m*(1000/4) for A_m in A_mean]
    A_mean = A_mean_Hz
    quantile = 0.50
    A_mean_50 = np.quantile(A_mean, quantile)
    A_nonzero_50 = np.quantile(A_nonzero, quantile)
    A_nonzero_frac_50 = np.quantile(A_nonzero_frac, quantile)

    quantile = 0.25
    A_mean_25 = np.quantile(A_mean, quantile)
    A_nonzero_25 = np.quantile(A_nonzero, quantile)
    A_nonzero_frac_25 = np.quantile(A_nonzero_frac, quantile)

    sns.set(style='white')
    fig, ax = plt.subplots(1, 3, figsize=((cm2inch(21), cm2inch(2.8))))
    bins = np.arange(0, np.max(A_mean), 20)
    ax[0].hist(A_mean, alpha=0.6, color='darkblue')
    ax[0].set_xlabel(r'population firing rate R (Hz)')
    ax[0].set_ylabel(r'# recordings')
    ax[0].text(0.8, 0.8, r'$q_{50}$'+'={:.1f} Hz'.format(A_mean_50),
               horizontalalignment='center', verticalalignment='center',
               transform=ax[0].transAxes, color='green', fontsize=11)
    ax[0].text(0.8, 0.65, r'$q_{25}$'+'={:.1f} Hz'.format(A_mean_25),
               horizontalalignment='center', verticalalignment='center',
               transform=ax[0].transAxes, color='green', fontsize=11)

    bins = np.arange(0, np.max(A_nonzero), 20)
    ax[1].hist(A_nonzero, alpha=0.6, color='darkblue')
    ax[1].set_xlabel(r'$n_{A_t\neq 0}$')
    ax[1].text(0.8, 0.8, r'$q_{50}$'+'={:.0f}'.format(A_nonzero_50),
               horizontalalignment='center', verticalalignment='center',
               transform=ax[1].transAxes, color='green', fontsize=11)
    ax[1].text(0.8, 0.65, r'$q_{25}$'+'={:.0f}'.format(A_nonzero_25),
               horizontalalignment='center', verticalalignment='center',
               transform=ax[1].transAxes, color='green', fontsize=11)

    bins = np.arange(0, np.max(A_nonzero_frac), 20)
    ax[2].hist(A_nonzero_frac, alpha=0.6, color='darkblue')
    ax[2].set_xlabel(r'$n_{A_t\neq 0}$ / $n_{A_t}$ (%)')
    ax[2].text(0.7, 0.8, r'$q_{50}$'+'={:.1f} %'.format(A_nonzero_frac_50),
               horizontalalignment='center',
               verticalalignment='center', transform=ax[2].transAxes,
               color='green', fontsize=11)
    ax[2].text(0.7, 0.65, r'$q_{25}$'+'={:.1f} %'.format(A_nonzero_frac_25),
               horizontalalignment='center',
               verticalalignment='center', transform=ax[2].transAxes,
               color='green', fontsize=11)

    ax[0].tick_params(top='off', bottom='on', left='on', right='off',
                      labelleft='on',labelbottom='on')
    ax[1].tick_params(top='off', bottom='on', left='on', right='off',
                      labelleft='on', labelbottom='on')
    ax[2].tick_params(top='off', bottom='on', left='on', right='off',
                      labelleft='on', labelbottom='on')

    sns.despine()
    plt.subplots_adjust(bottom=None, right=None,
                        wspace=0.6, hspace=0.3, left=0.15, top=1.2)
    figfile = '{}S4_recording_statistics.pdf'.format(resultpath)
    plt.savefig(figfile, bbox_inches='tight')
    plt.show()


def mt_all_supp(save_path,fit_method, kmin, kmax, windowsize, windowstep, dt):

    patients = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17,
                18, 19, 20]
    fig = plt.figure(figsize=(cm2inch(2*13.5), cm2inch(1.2*10*5)))
    outer = gridspec.GridSpec(10, 2, wspace=0.2, hspace=0.9)
    laterality = ["SOZ", "nSOZ"]

    for ip, patient in enumerate(patients):
        inner = gridspec.GridSpecFromSubplotSpec(2, 1, subplot_spec=outer[ip],
                                                 wspace=0.1, hspace=0.9)

        for j, soz in enumerate(laterality):
            ax = plt.Subplot(fig, inner[j])
            ax = mt_plot(ax, patient, soz, save_path,fit_method, kmin, kmax,
                          windowsize, windowstep, dt)

            if j == 0:
                ax.set_title('Patient {}'.format(int(patient)), fontsize=12)
                ax.tick_params(top='off', bottom='on', left='off',
                               right='off', labelleft='on',
                               labelbottom='off')
                ax.set_xlabel('')

            fig.add_subplot(ax)
    sns.despine(fig)
    plt.savefig('{}/S6_allmt.pdf'.format(save_path), bbox_inches='tight')
    plt.show()
    plt.close()


def binning_label(patient, soz):
    focifile = '../Data/patients.txt'
    f = open(focifile, 'r')
    foci = [line.rstrip('\n').split('\t') for line in f.readlines()]

    focus = 'NA'
    for i, idf in enumerate(foci[1:]):
        if int(idf[0]) == int(patient):
            focus = idf[1]

    if soz == "SOZ":
        if focus == "L":
            binning = "left"
        elif focus == "R":
            binning = "right"
        else:
            raise Exception("No clear focus found.")

    elif soz == "nSOZ":
        if focus == "L":
            binning = "right"
        elif focus == "R":
            binning = "left"
        else:
            raise Exception("No clear focus found.")
    return binning


def mt_plot(ax, patient, soz, save_path, fit_method, kmin, kmax,
            windowsize, windowstep, dt):

    recordings = mr_utilities.list_preictrecordings(patient)[0]
    binning = binning_label(patient, soz)

    for irec, rec in enumerate(recordings):

        mt_file = '{}{}/mt_results_{}_{}_kmin{}_kmax{}_winsize{}_winstep{}' \
                  '_dt{}.pkl'.format(save_path, rec, binning, fit_method, kmin,
                                     kmax, windowsize, windowstep, dt)

        if not os.path.isfile(mt_file):
            print('{} not existing'.format(mt_file))
            continue

        # -----------------------------------------------------------------#
        # read m(t) data
        # -----------------------------------------------------------------#
        mt_frame = pd.read_pickle(mt_file)

        original_length = len(mt_frame.m.tolist())
        rejected_inconsistencies = ['H_nonzero', 'H_linear',
                                    'not_converged']
        for incon in rejected_inconsistencies:
            if not mt_frame.empty:
                mt_frame = mt_frame[[incon not in mt_frame[
                    'inconsistencies'].tolist()[i] for i in range(
                        len(mt_frame['inconsistencies'].tolist()))]]
        if len(mt_frame.m.tolist()) < 0.95 * original_length:
            continue

        kmin = mt_frame.kmin.unique()[0]
        kmax = mt_frame.kmax.unique()[0]
        fit_method = mt_frame.fit_method.unique()[0]
        winsize = mt_frame.windowsize.unique()[0]
        t0_list = mt_frame.t0.tolist()
        mt = mt_frame.m.tolist()
        t0_list = [t0_ms * 0.001 for t0_ms in t0_list]

        half_win_sec = int((winsize / 2) / 1000 * dt)
        t_middle = [t0 + half_win_sec for t0 in t0_list]

        # -----------------------------------------------------------------#
        # Plotting
        # -----------------------------------------------------------------#
        ax.plot(t_middle, mt, '.', label='Rec {}'.format(irec + 1),
                color=mycol[irec], markersize=0.5)
        ax.axhline(y=1, color='black', linestyle='--', alpha=0.3)
        ax.set_ylabel(r"$\hat{m}$")
        ax.set_xlim((t0_list[0] - 10, t0_list[-1] + 2 * half_win_sec + 10))
        binning_focus = mr_utilities.binning_focus(binning, int(patient))
        if binning_focus == 'focal':
            focus_label = 'ipsi'
        else:
            focus_label = 'contra'
        ax.set_title("{}".format(focus_label), loc='left')
        ax.set_xlim((-10, 610))

    handles, labels = ax.get_legend_handles_labels()
    if len(labels) == 0:
        if soz == 'SOZ':
            ax.set_title("ipsi", loc='left')
        if soz == 'nSOZ':
            ax.set_title("contra", loc='left')
        ax.axhline(y=1, color='black', linestyle='--', alpha=0.3)
        ax.set_ylabel(r"$\hat{m}$")
        ax.set_ylim((0.8, 1.05))
        ax.set_xlim((-10, 610))

    ax.xaxis.grid(False)
    ax.set_xlabel("time (s)")
    ax.tick_params(top='off', bottom='off', left='off', right='off',
                   labelleft='on', labelbottom='off')
    ax.tick_params(top='off', bottom='on', left='off', right='off',
                   labelleft='on', labelbottom='on')
    sns.despine()
    return ax


def windowsize_analysis_recs(data_path, recording, binning, outputfilename,
                             resultpath):

    outputfile = '{}{}/pkl/{}'.format(data_path, recording, outputfilename)

    if not os.path.isfile(outputfile):
        print("No binning data for recording ", recording)
        print(outputfile)
        return
    data = pd.read_pickle(outputfile)

    dt = data['dt'].unique()[0]
    activity = data[data['binning'] == binning]['activity'].tolist()[0]

    winstep = 100
    ksteps = (1, 400)
    windowsize = np.arange(2500, 40000, 2500)
    N_estimates = len(activity) - np.max(windowsize)
    tau_estimate = []
    tau_conf = []
    m_estimate = []
    m_conf = []
    for Lw in windowsize:
        tau_L = []
        m_L = []
        for iw in range(0, N_estimates, winstep):
            activity_window = activity[iw:iw+Lw]
            input = mre.input_handler(activity_window)
            rk = mre.coefficients(input, steps=ksteps, dt=dt)
            fitres_offset = mre.fit(rk, fitfunc=mre.f_exponential_offset)
            tau_L.append(fitres_offset.tau)
            m_L.append(fitres_offset.mre)
        conf_int_off = sps.mstats.mquantiles(tau_L, prob=[0.025, 0.975],
                                             alphap=0, betap=1, axis=None,
                                             limit=())
        conf_int_off = [abs(np.mean(tau_L) - ci) for ci in conf_int_off]
        tau_conf.append(conf_int_off)
        tau_estimate.append(np.mean(tau_L))

        conf_int_off = sps.mstats.mquantiles(m_L, prob=[0.025, 0.975],
                                             alphap=0, betap=1, axis=None,
                                             limit=())
        conf_int_off = [abs(np.mean(m_L) - ci) for ci in conf_int_off]
        m_conf.append(conf_int_off)
        m_estimate.append(np.mean(m_L))

    windowsize_seconds = [w * 4 / 1000 for w in windowsize]
    fig, ax = plt.subplots(1, 1, figsize=((cm2inch(9), cm2inch(0.9 * 4))))
    ax.scatter(windowsize_seconds, tau_estimate,
               label=r'$\hat{\tau}_{offset}$', c='steelblue',
               marker='.', s=9)
    ax.errorbar(windowsize_seconds, tau_estimate,
                yerr=np.transpose(np.array(tau_conf)),
                linestyle="None", c='steelblue', elinewidth=0.9,
                capsize=2.5)

    ax.set_ylim((-10, 800))
    T = [20, 40, 60, 80, 100, 120, 140]
    ax.set_xticks(T)
    ax.set_xlabel('window size (s)')
    ax.set_ylabel(r'$\hat{\tau}$ (ms)')
    sns.despine()
    plt.subplots_adjust(bottom=None, right=None,
                        wspace=0.4, hspace=0.3, left=0.15, top=1.2)
    figfile = '{}S5_estimates_winsizes_rec_{}_{}.pdf'.format(
        resultpath, recording, binning)
    plt.savefig(figfile, bbox_inches='tight')

    fig, ax = plt.subplots(1, 1, figsize=((cm2inch(9), cm2inch(0.9 * 4))))
    ax.scatter(windowsize_seconds, m_estimate,
               label=r'$\hat{\tau}_{offset}$', c='royalblue',
               marker='.', s=9)
    ax.errorbar(windowsize_seconds, m_estimate,
                yerr=np.transpose(np.array(m_conf)),
                linestyle="None", c='royalblue', elinewidth=0.9,
                capsize=2.5)

    T = [20, 40, 60, 80, 100, 120, 140]
    ax.set_xticks(T)
    ax.set_ylim((0.5, 1.25))
    ax.set_xlabel('window size (s)')
    ax.set_ylabel(r'$\hat{m}$')
    sns.despine()
    plt.subplots_adjust(bottom=None, right=None,
                        wspace=0.4, hspace=0.3, left=0.15, top=1.2)
    figfile = '{}S5_estimates_m_winsizes_rec_{}_{}.pdf'.format(
        resultpath, recording, binning)
    plt.savefig(figfile, bbox_inches='tight')
    plt.show()


def windowsize_analysis():
    binning = 'left'
    outputfilename = 'activity_SU_4ms.pkl'
    resultpath = '../Results/preictal/singleunits/'

    distribution_A_recordings(resultpath)

    # warning: analysis below has long computation time
    # data_path = '../Data/interictal/'
    # recording = '13ref'
    # windowsize_analysis_recs(data_path, recording, binning, outputfilename,
    #                          resultpath)
    #
    # data_path = '../Data/preictal/'
    # recording = '13_02'
    # windowsize_analysis_recs(data_path, recording, binning, outputfilename,
    #                          resultpath)


def all_mt():
    result_path = '../Results/preictal/singleunits/'
    ksteps = (1, 400)
    kmin = ksteps[0]
    kmax=ksteps[1]
    windowsize = 20000
    windowstep = 500
    dt = 4
    mt_all_supp(result_path, 'offset', kmin, kmax, windowsize, windowstep, dt)


def example_autocorrelations(patient, pkl_file_inter, pkl_file_pre,
                             resultpath):
    mr_results_pre = pd.read_pickle(pkl_file_pre)
    mr_results_inter = pd.read_pickle(pkl_file_inter)
    mr_results_pre['rectype'] = 'preictal'
    mr_results_inter['rectype'] = 'interictal'
    mr_results = pd.concat([mr_results_pre, mr_results_inter],
                           ignore_index=True)

    mr_results = mr_results[mr_results['patient'] == patient]
    binning_foci = ['focal', 'contra-lateral']

    sns.set(style="white")
    fig, ax = plt.subplots(2, 4, figsize=(cm2inch(24), cm2inch(6)))
    fit_method = 'offset'
    mr_rec_pre = mr_results[mr_results['rectype'] == 'preictal']
    pre_recordings_choice = mr_rec_pre.recording.unique()[:3]

    for si, soz in enumerate(binning_foci):
        # plot interictal in first position
        mr_rec_inter = mr_results[mr_results['rectype'] == 'interictal']
        mr_rec_inter = mr_rec_inter[mr_rec_inter['binning_focus'] == soz]
        rk = mr_rec_inter['rk'].item()[0]
        k_steps = mr_rec_inter['rk_steps'].item()
        dt = mr_rec_inter['dt'].item()
        time_steps = np.array(k_steps[0]) * dt
        ax[si, 0].plot(time_steps, rk, '-', linewidth=0.8,
                      c='slategray', alpha=0.4)
        fitfunc = mr_utilities.string_to_func(fit_method)
        fitparams = mr_rec_inter['fit_params'].item()
        fitparams = tuple(fitparams[0])
        m = mr_rec_inter['m'].item()
        ax[si, 0].plot(time_steps, fitfunc(time_steps, *fitparams),
                       label=r'$m$ = {:.2f}'.format(m), color='slategray')

        ax[si, 0].tick_params(top='off', bottom='on', left='on',
                              right='off', labelleft='on', labelbottom='on')

        ax[si, 0].set_xlabel(r"time lag $k$ (ms)")
        ax[si, 0].legend(loc=1, fontsize='x-small')
        ax[si, 0].ticklabel_format(axis='y', style='sci', scilimits=(-1, 1),
                                   useMathText=True)

        # plot 3 preictal recordings next to it
        mr_rec_pre = mr_results[mr_results['rectype'] == 'preictal']
        mr_rec_pre = mr_rec_pre[mr_rec_pre['binning_focus'] == soz]
        for irec, rec in enumerate(pre_recordings_choice):
            mr_rec = mr_rec_pre[mr_rec_pre['recording'] == rec]
            rk = mr_rec['rk'].item()[0]
            k_steps = mr_rec['rk_steps'].item()
            dt = mr_rec['dt'].item()
            time_steps = np.array(k_steps[0]) * dt
            ax[si, irec+1].plot(time_steps, rk, '-', linewidth=0.8,
                           c='darkblue', alpha=0.4)
            fitfunc = mr_utilities.string_to_func(fit_method)
            fitparams = mr_rec['fit_params'].item()
            fitparams = tuple(fitparams[0])
            m = mr_rec['m'].item()
            ax[si, irec+1].plot(time_steps, fitfunc(time_steps, *fitparams),
                           label=r'$m$ = {:.2f}'.format(m), color='darkblue')

            ax[si, irec+1].tick_params(top='off', bottom='on', left='on',
                                  right='off', labelleft='on',labelbottom='on')

            ax[si, irec+1].set_xlabel(r"time lag $k$ (ms)")
            ax[si, irec+1].legend(loc=1, fontsize='x-small')
            ax[si, irec+1].ticklabel_format(axis='y', style='sci',
                                       scilimits=(-1, 1), useMathText=True)
    ax[0, 0].set_ylabel("ACF")
    ax[1, 0].set_ylabel("ACF")
    plt.subplots_adjust(bottom=None, right=None,
                        wspace=0.5, hspace=1, left=0.15, top=1.2)
    sns.despine(fig)
    resultpath_rec = '{}S1_autocorrelations_patient{}.pdf'.format(resultpath, patient)
    plt.savefig(resultpath_rec, bbox_inches='tight')
    plt.close()


def create_example_ACF_plot():
    pkl_file_pre = '../Results/preictal/singleunits/mr_results.pkl'
    pkl_file_inter = '../Results/interictal/singleunits/mr_results.pkl'
    resultpath = '../Results/preictal/singleunits/'
    patient = 12
    example_autocorrelations(patient, pkl_file_inter, pkl_file_pre, resultpath)
    patient = 13
    example_autocorrelations(patient, pkl_file_inter, pkl_file_pre, resultpath)


def main():
    create_example_ACF_plot()
    windowsize_analysis()
    all_mt()


if __name__ == '__main__':
    main()