import os

import numpy as np
import mpmath as mp
import matplotlib.pyplot as plt
import pandas as pd
from IPython import embed
from matplotlib import rc

from threefish.core import find_code_vs_not
from threefish.load import save_visualization
from threefish.plot.colorbar import colorbar_outside
from threefish.plot.limits import set_clim_same
from threefish.plot.tags import tag2
from threefish.RAM.plot_labels import nonlin_title, set_xlabel_arrow, set_ylabel_arrow
from threefish.RAM.plot_subplots import plt_RAM_perc
from threefish.RAM.reformat_matrix import convert_csv_str_to_float, find_model_names_susept
from threefish.defaults import default_figsize, default_settings
from threefish.plot.plotstyle import plot_style
from threefish.RAM.values import perc_model_full

rc('text', usetex=True)


def FiringRate(mu, D, tau_ref, vR, vT):
    x_start = (mu - vT)/mp.sqrt(2.0*D)
    x_end = (mu - vR)/mp.sqrt(2.0*D)
    dx = 0.0001
    r = 0.0

    for i in np.arange(x_start, x_end, dx):
        integrand = mp.exp(i**2) * mp.erfc(i)
        r += integrand*dx
    
    return 1.0/(tau_ref + mp.sqrt(mp.pi)*r)

def Susceptibility_1(omega, r0, mu, D, tau_ref, vR, vT):
    delta = (vR**2 - vT**2 + 2.0*mu*(vT - vR))/(4.0*D)
    a = (r0 * omega*1.0j)/(mp.sqrt(D) * (omega*1.0j - 1.0))
    b = mp.pcfd(omega*1.0j - 1.0, (mu - vT)/mp.sqrt(D)) - mp.exp(delta) * mp.pcfd(omega*1.0j - 1.0, (mu - vR)/mp.sqrt(D))
    c = mp.pcfd(omega*1.0j, (mu - vT)/mp.sqrt(D)) - mp.exp(delta) * mp.exp(omega*1.0j*tau_ref) * mp.pcfd(omega*1.0j, (mu - vR)/mp.sqrt(D))
    return a * b/c

def Susceptibility_2(omega1, omega2, chi1_1, chi1_2, r0, mu, D, tau_ref, vR, vT):
    delta = (vR**2 - vT**2 + 2.0*mu*(vT - vR))/(4.0*D)
    a1 = r0*(1.0 - omega1*1.0j - omega2*1.0j)*(omega1*1.0j + omega2*1.0j)/(2.0*D*(omega1*1.0j - 1.0)*(omega2*1.0j - 1.0))
    a2 = (omega1*1.0j + omega2*1.0j)/(2.0*mp.sqrt(D))
    a3 = chi1_1/(omega2*1.0j - 1.0) + chi1_2/(omega1*1.0j - 1.0)
    b1 = mp.pcfd(omega1*1.0j + omega2*1.0j - 2.0, (mu - vT)/mp.sqrt(D)) - mp.exp(delta) * mp.pcfd(omega1*1.0j + omega2*1.0j - 2.0, (mu - vR)/mp.sqrt(D))
    b2 = mp.pcfd(omega1*1.0j + omega2*1.0j - 1.0, (mu - vT)/mp.sqrt(D))
    b3 = mp.exp(delta) * mp.pcfd(omega1*1.0j + omega2*1.0j - 1.0, (mu - vR)/mp.sqrt(D))
    c = mp.pcfd(omega1*1.0j + omega2*1.0j, (mu - vT)/mp.sqrt(D)) - mp.exp(delta) * mp.exp(1.0j*(omega1 + omega2)*tau_ref) * mp.pcfd(omega1*1.0j + omega2*1.0j, (mu - vR)/mp.sqrt(D))

    return a1 * b1/c + a2*a3*b2/c - a2*a3*b3/c



if __name__ == '__main__':

    # LIF model in dimensionless units: dv/dt = -v + mu + sqrt(2D)*xi
    ## v: membrane voltage
    ## mu: mean input voltage
    ## D: noise intensity
    ## xi: white Gaussian noise
    ## tau_mem = 1 (membrane time constant, skipped here)


    save_name = '-transfer'
    load_function, name1, save_name = find_model_names_susept(save_name)
    save_name2 = '-matrix'
    load_function, name2, save_name = find_model_names_susept(save_name2)
    versions_comp, subfolder, mod_name_slash, mod_name, subfolder_path = find_code_vs_not()

    redo = False
    if versions_comp == 'develop':
        if (not os.path.exists(name1)) | (redo == True):
            mu = 1.1
            D = 0.001

            # threshold voltage
            vT = 1.0
            # reset voltage
            vR = 0.0
            # refractory period
            tau_ref = 0.0

            # stationary firing rate
            r0 = FiringRate(mu, D, tau_ref, vR, vT)

            x_min = 0.01
            x_max = 1.0
            dx = 0.01
            y_min = 0.01
            y_max = 1.0
            dy = 0.01

            f1_range = np.arange(x_min, x_max + 0.5 * dx, dx)
            f2_range = np.arange(y_min, y_max + 0.5 * dy, dy)

            chi1_data = np.zeros(f2_range.size, dtype=complex)
            chi2_data = np.zeros((f2_range.size, f1_range.size), dtype=complex)

            for f2 in range(f2_range.size):
                omega2 = 2.0 * np.pi * f2_range[f2]
                chi1_2 = Susceptibility_1(omega2, r0, mu, D, tau_ref, vR, vT)
                print('Step: ' + str(f2 + 1) + ' von ' + str(int((y_max - y_min) / dy) + 1))
                chi1_data[f2] = chi1_2

                for f1 in range(f1_range.size):
                    omega1 = 2.0 * np.pi * f1_range[f1]
                    chi1_1 = Susceptibility_1(omega1, r0, mu, D, tau_ref, vR, vT)
                    chi2 = Susceptibility_2(omega1, omega2, chi1_1, chi1_2, r0, mu, D, tau_ref, vR, vT)
                    chi2_data[f2][f1] = chi2
                    chi1_data[f1] =  chi1_1

            #embed()
            #chi2_data.index = chi2_data.index.map(lambda x: "{:.2f}".format(float(x)))
            #chi2_data.columns = chi2_data.columns.map(lambda x: "{:.2f}".format(float(x)))
            frame = pd.DataFrame(chi1_data, index=f1_range, columns = ['transfer'])
            frame.index = frame.index.map(lambda x: "{:.2f}".format(float(x)))
            frame.to_csv(name1)

            frame2 = pd.DataFrame(chi2_data, index=f2_range, columns = f1_range)
            frame2.index = frame2.index.map(lambda x: "{:.2f}".format(float(x)))
            frame2.columns = frame2.columns.map(lambda x: "{:.2f}".format(float(x)))
            frame2.to_csv(name2)

            # name1 =
            print('transfer embed')
        else:
            frame = pd.read_csv(name1, index_col=0)
            frame2 = pd.read_csv(name2, index_col=0)
    else:
        frame = pd.read_csv(name1, index_col = 0)
        frame2 = pd.read_csv(name2, index_col = 0)
    #embed()
    # colorbar plot of chi_2
    #def vary_contrasts_big_with_tuning3(freqs=[(39.5, -135.5)], cells_here='2011-10-25-ad-invivo-1'):
    default_settings()  # ts=13, ls=13, fs=13, lw = 0.7
    plot_style()
    default_figsize(column=2, length=2.5)


    fig1, ax = plt.subplots(1,2)
    plt.subplots_adjust(bottom=0.2, wspace=0.35, left=0.1, right = 0.92)
    kwargs = {'cmap': 'turbo', 'extent': [frame2.columns[0], frame2.columns[-1], frame2.index[0], frame2.index[-1]]}#x_min - 0.5*dx, x_max + 0.5*dx, y_min - 0.5*dy, y_max + 0.5*dy
    #im = ax[0].imshow(np.abs(frame2), **kwargs, origin='lower')
    #clim = im.get_clim()
    #im.set_clim(0, clim[-1])
    #embed()

    #cbar.set_label(nonlin_title(add_nonlin_title=' ['),
    #               rotation=90,
    #               labelpad=8)
    #set_clim_same([im], mats=[np.abs(frame2)], lim_type='up', nr_clim='perc', clims='', percnr=99)#perc_model_full()
    #plt.colorbar(pos, fraction=0.046)#, pad=0.04

    ###################################################
    # transfer
    #embed()
    #new_keys, frame0 = convert_csv_str_to_float(frame)
    frame = frame.astype(complex)
    ax[0].plot(frame.index, np.abs(frame['transfer']), color = 'black')
    ax[0].set_xlabel(r'$f_{1}$')
    ax[0].set_ylabel(r'$|\chi_{1}|$')
    ax[0].set_xlim(0, frame.index[-1])
    ylim = ax[0].get_ylim()
    ax[0].set_ylim(0, ylim[-1])



    ###################################################
    # matrix
    frame2.columns = frame2.index
    new_keys, stack_plot = convert_csv_str_to_float(frame2)
    im = plt_RAM_perc(ax[1], 'perc', np.abs(stack_plot))
    set_clim_same([im],
                  mats=[np.abs(stack_plot)],
                  lim_type='up',
                  nr_clim='perc',
                  clims='',
                  percnr=perc_model_full())

    pos_rel = -0.12
    #set_xlabel_arrow(ax[0], xpos=0.95, ypos=pos_rel)
    #set_ylabel_arrow(ax[0], xpos=pos_rel, ypos=0.95)
    cbar, left, bottom, width, height = colorbar_outside(ax[1],
                                                         im,
                                                         add=5,
                                                         width=0.01)
    ax[1].set_xlabel(r'$f_{1}$')
    ax[1].set_ylabel(r'$f_{2}$')
    #ax[0].set_title(r'$|\chi_{2}(\omega_{1}, \omega_{2})|$' + ': ' + r'$\mu =$ ' + str(mu) + ', D = ' + str(D) + ', ' + r'$r_{0} =$ ' + str(round(r0,3)), pad = 10)
    ax[1].set_xticks_delta(0.2)
    ax[1].set_xticks_delta(0.2)
    ax[1].set_xlim(frame2.columns[0]-0.01, frame2.columns[-1])
    ax[1].set_ylim(frame2.columns[0] - 0.01, frame2.columns[-1])
    cbar, left, bottom, width, height = colorbar_outside(ax[1], im, add=5, width=0.01)
    cbar.set_label(r'$|\chi_{2}|$', rotation=90, labelpad=8)

    fig = plt.gcf()
    tag2(axes = ax, xoffs=[-0.5, -5.5], yoffs=1.7)  # ax_ams[3],

    save_visualization()
    #plt.show()

    # 3d plot of chi_2
    surface = False
    if surface:
        fig2 = plt.figure(figsize=(6,6))
        ax = plt.axes(projection='3d')

        x = frame2.index
        y = frame2.columns
        X,Y = np.meshgrid(x,y)

        Z = np.abs(frame2.values)

        surf = ax.plot_surface(X,Y,Z, cmap = plt.cm.turbo)
        fig2.colorbar(surf, shrink=0.5)

        ax.set(xlabel=r'$f_{1}$', ylabel=r'$f_{2}$', zlabel=r'$|\chi_{2}(\omega_{1}, \omega_{2})|$')