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


def calc_nonlin_analytic_values():

    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):
            frame, frame2 = develop_analytic_nonlin(name1, name2)
        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)
    return frame, frame2


def develop_analytic_nonlin(name1, name2):
    # 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)
    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')
    return frame, frame2


def plot_chi2():
    global ax, frame2
    frame, frame2 = calc_nonlin_analytic_values()
    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)
    ###################################################
    # plot transfer function
    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])
    ###################################################
    # plot 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
    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[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)
    tag2(axes=ax, xoffs=[-0.5, -5.5], yoffs=1.7)  # ax_ams[3],
    save_visualization()


if __name__ == '__main__':



    plot_chi2()