import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib import patches
import matplotlib.patches as mpatches
import scipy.io as sio

# plotting configuration
ratio = 1.5
figure_len, figure_width = 15*ratio, 12*ratio
font_size_1, font_size_2 = 36*ratio, 36*ratio
legend_size = 18*ratio
line_width, tick_len = 3*ratio, 10*ratio
marker_size = 15*ratio
plot_line_width = 5*ratio
hfont = {'fontname': 'Arial'}

b_plotting_activity = True
b_plotting_paradoxical_effect = False
b_plotting_frozen_inhibition = False

if b_plotting_activity:
    # simulation setup
    dt = 0.0001
    T = int(9/dt)

    l_g_e_ini = [1.55, 1.8]

    for g_e_ini in l_g_e_ini:
        Jacobian_mat = np.zeros((2, 2)) * np.nan

        # neuronal parameters
        tau_e, tau_i = 0.020, 0.010
        alpha_e, alpha_i = 2, 2

        # short-term depression
        x, u_d = 1, 1
        tau_x = 0.20

        # network connectivity
        Jee = 1.8
        Jie = 1.0
        Jei = 1.0
        Jii = 0.6

        r_e, r_i = 0, 0
        z_e, z_i = 0, 0

        l_r_e, l_r_i, l_x = [], [], []
        l_ISN_index = []

        for i in range(T):
            if 50000 <= i < 70000:
                g_e, g_i = 3.0, 2
            else:
                g_e, g_i = g_e_ini, 2

            g_e = g_e * (g_e > 0)
            g_i = g_i * (g_i > 0)

            # SSN part
            z_e = Jee * x * r_e - Jei * r_i + g_e
            z_i = Jie * r_e - Jii * r_i + g_i

            z_e = z_e * (z_e > 0)
            z_i = z_i * (z_i > 0)

            r_e = r_e + (-r_e + np.power(z_e, alpha_e)) / tau_e * dt
            r_i = r_i + (-r_i + np.power(z_i, alpha_i)) / tau_i * dt

            r_e = r_e * (r_e > 0)
            r_i = r_i * (r_i > 0)

            x = x + ((1 - x) / tau_x - u_d * x * r_e) * dt
            x = np.clip(x, 0, 1)

            Jacobian_mat[0, 0] = 1.0 / tau_e * (x * Jee * alpha_e * np.power(r_e, (alpha_e - 1.0) / alpha_e) - 1)
            Jacobian_mat[0, 1] = 1.0 / tau_e * Jee * alpha_e * np.power(r_e, (2 * alpha_e - 1.0) / alpha_e)
            Jacobian_mat[1, 0] = - u_d * x
            Jacobian_mat[1, 1] = -1.0 / tau_x - u_d * r_e
            lambda_1 = np.linalg.eig(Jacobian_mat)[0][0]
            lambda_2 = np.linalg.eig(Jacobian_mat)[0][1]
            l_ISN_index.append(np.max([lambda_1.real, lambda_2.real]))

            l_r_e.append(r_e)
            l_r_i.append(r_i)
            l_x.append(x)

        l_r_e = np.asarray(l_r_e)
        l_r_i = np.asarray(l_r_i)
        l_x = np.asarray(l_x)
        l_ISN_index = np.asarray(l_ISN_index)


        # plotting
        plt.figure(figsize=(figure_len, figure_width))
        ax = plt.gca()
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['bottom'].set_visible(True)
        ax.spines['left'].set_visible(True)
        for axis in ['top', 'bottom', 'left', 'right']:
            ax.spines[axis].set_linewidth(line_width)
        plt.tick_params(width=line_width, length=tick_len)
        plt.yscale('symlog', linthreshy=0.1)

        plt.plot(l_r_e, color='blue', linewidth=plot_line_width)
        plt.plot(l_r_i, color='red', linewidth=plot_line_width)

        plt.xticks(np.arange(30000, 90000 + 5000, 20000), np.arange(0, 6 + 0.5, 2), fontsize=font_size_1, **hfont)
        plt.yticks([0, 1, 100, 10000], fontsize=font_size_1, **hfont)
        plt.xlabel('Time (s)', fontsize=font_size_1, **hfont)
        plt.ylabel('Firing rate (Hz)', fontsize=font_size_1, **hfont)
        plt.xlim([30000, 90000])
        plt.ylim([0, 10000])
        plt.legend(['Exc', 'Inh'], prop={"family": "Arial", 'size': font_size_1}, loc='upper right')
        if g_e_ini == 1.55:
            plt.savefig('paper_figures/png/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_initial_non_ISN.png')
            plt.savefig('paper_figures/pdf/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_initial_non_ISN.pdf')
        else:
            plt.savefig('paper_figures/png/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_initial_ISN.png')
            plt.savefig('paper_figures/pdf/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_initial_ISN.pdf')

        # plotting
        plt.figure(figsize=(figure_len, figure_width))
        ax = plt.gca()
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['bottom'].set_visible(True)
        ax.spines['left'].set_visible(True)
        for axis in ['top', 'bottom', 'left', 'right']:
            ax.spines[axis].set_linewidth(line_width)
        plt.tick_params(width=line_width, length=tick_len)

        plt.plot(l_x, linewidth=plot_line_width)
        plt.plot(np.sqrt(1.0/(Jee * alpha_e) * np.power(l_r_e, 1.0/alpha_e -1)), '--', linewidth=plot_line_width)

        plt.xticks(np.arange(30000, 90000 + 5000, 20000), np.arange(0, 6 + 0.5, 2), fontsize=font_size_1, **hfont)
        plt.yticks([0, 0.2, 0.4, 0.6, 0.8, 1.0], fontsize=font_size_1, **hfont)
        plt.xlabel('Time (s)', fontsize=font_size_1, **hfont)
        plt.ylabel('x', fontsize=font_size_1, **hfont)
        plt.xlim([30000, 90000])
        plt.ylim([0, 1])
        if g_e_ini == 1.55:
            plt.savefig('paper_figures/png/Revision_Fig_Point_1_1_Supralinear_network_x_initial_non_ISN.png')
            plt.savefig('paper_figures/pdf/Revision_Fig_Point_1_1_Supralinear_network_x_initial_non_ISN.pdf')
        else:
            plt.savefig('paper_figures/png/Revision_Fig_Point_1_1_Supralinear_network_x_initial_ISN.png')
            plt.savefig('paper_figures/pdf/Revision_Fig_Point_1_1_Supralinear_network_x_initial_ISN.pdf')


        # plotting
        plt.figure(figsize=(figure_len, figure_width))
        ax = plt.gca()
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['bottom'].set_visible(True)
        ax.spines['left'].set_visible(True)
        for axis in ['top', 'bottom', 'left', 'right']:
            ax.spines[axis].set_linewidth(line_width)
        plt.tick_params(width=line_width, length=tick_len)
        plt.yscale('symlog', linthreshy=10)

        plt.plot(l_ISN_index, linewidth=plot_line_width)

        plt.xticks(np.arange(30000, 90000 + 5000, 20000), np.arange(0, 6 + 0.5, 2), fontsize=font_size_1, **hfont)
        plt.yticks([-10000, -100, 0, 100, 10000], fontsize=font_size_1, **hfont)
        plt.xlabel('Time (s)', fontsize=font_size_1, **hfont)
        plt.ylabel('ISN index', fontsize=font_size_1, **hfont)
        plt.xlim([30000, 90000])
        plt.ylim([-10000, 10000])
        plt.hlines(y=0, xmin=30000, xmax=90000, colors='k', linestyles=[(0, (6, 6, 6, 6))], linewidth=line_width)
        if g_e_ini == 1.55:
            plt.savefig('paper_figures/png/Revision_Fig_Point_1_1_Supralinear_network_ISN_index_initial_non_ISN.png')
            plt.savefig('paper_figures/pdf/Revision_Fig_Point_1_1_Supralinear_network_ISN_index_initial_non_ISN.pdf')
        else:
            plt.savefig('paper_figures/png/Revision_Fig_Point_1_1_Supralinear_network_ISN_index_initial_ISN.png')
            plt.savefig('paper_figures/pdf/Revision_Fig_Point_1_1_Supralinear_network_ISN_index_initial_ISN.pdf')







if b_plotting_paradoxical_effect:
    # simulation setup
    dt = 0.0001
    T = int(9 / dt)

    # neuronal parameters
    tau_e, tau_i = 0.020, 0.010
    alpha_e, alpha_i = 2, 2

    # short-term depression
    x, u_d = 1, 1
    tau_x = 0.20

    # network connectivity
    Jee = 1.8
    Jie = 1.0
    Jei = 1.0
    Jii = 0.6

    l_b_before_stimulation = [True, False]

    for b_before_stimulation in l_b_before_stimulation:
        x = 1
        r_e, r_i = 0, 0
        z_e, z_i = 0, 0

        l_r_e, l_r_i, l_x = [], [], []

        for i in range(T):
            if 50000 <= i < 70000:
                g_e, g_i = 3.0, 2
            else:
                g_e, g_i = 1.8, 2

            if b_before_stimulation:
                if 42000 < i <= 49000:
                    g_i = 2.1
                else:
                    pass
            else:
                if 62000 < i <= 69000:
                    g_i = 2.1
                else:
                    pass

            g_e = g_e * (g_e > 0)
            g_i = g_i * (g_i > 0)

            # SSN part
            z_e = Jee * x * r_e - Jei * r_i + g_e
            z_i = Jie * r_e - Jii * r_i + g_i

            z_e = z_e * (z_e > 0)
            z_i = z_i * (z_i > 0)

            r_e = r_e + (-r_e + np.power(z_e, alpha_e)) / tau_e * dt
            r_i = r_i + (-r_i + np.power(z_i, alpha_i)) / tau_i * dt

            r_e = r_e * (r_e > 0)
            r_i = r_i * (r_i > 0)

            x = x + ((1 - x) / tau_x - u_d * x * r_e) * dt
            x = np.clip(x, 0, 1)

            l_r_e.append(r_e)
            l_r_i.append(r_i)
            l_x.append(x)

        l_r_e = np.asarray(l_r_e)
        l_r_i = np.asarray(l_r_i)
        l_x = np.asarray(l_x)

        if b_before_stimulation:
            plt.figure(figsize=(figure_len, figure_width))
            ax = plt.gca()
            ax.spines['top'].set_visible(False)
            ax.spines['right'].set_visible(False)
            ax.spines['bottom'].set_visible(True)
            ax.spines['left'].set_visible(True)
            for axis in ['top', 'bottom', 'left', 'right']:
                ax.spines[axis].set_linewidth(line_width)
            plt.tick_params(width=line_width, length=tick_len)

            mean_e = l_r_e / np.mean(l_r_e[40000:42000])
            mean_i = l_r_i / np.mean(l_r_i[40000:42000])

            plt.plot(mean_e, color='blue', linewidth=plot_line_width)
            plt.plot(mean_i, color='red', linewidth=plot_line_width)

            plt.xticks([40000, 42000, 44000, 46000, 48000], [1.0, 1.2, 1.4, 1.6, 1.8], fontsize=font_size_1, **hfont)
            plt.yticks([0, 0.2, 0.4, 0.6, 0.8, 1.0, 1.2], fontsize=font_size_1, **hfont)
            plt.xlabel('Time (s)', fontsize=font_size_1, **hfont)
            plt.ylabel('Normalized firing rate', fontsize=font_size_1, **hfont)
            plt.xlim([40000, 48000])
            plt.ylim([0, 1.2])
            plt.legend(['Exc', 'Inh'], prop={"family": "Arial", 'size': font_size_1})
            plt.hlines(y=1, xmin=42000, xmax=50000, colors='k', linestyles=[(0, (6, 6, 6, 6))], linewidth=line_width)
            plt.savefig(
                'paper_figures/png/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_normalized_activity_paradoxical_effect_before_stimulation_ISN.png')
            plt.savefig(
                'paper_figures/pdf/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_normalized_activity_paradoxical_effect_before_stimulation_ISN.pdf')

        else:
            plt.figure(figsize=(figure_len, figure_width))
            ax = plt.gca()
            ax.spines['top'].set_visible(False)
            ax.spines['right'].set_visible(False)
            ax.spines['bottom'].set_visible(True)
            ax.spines['left'].set_visible(True)
            for axis in ['top', 'bottom', 'left', 'right']:
                ax.spines[axis].set_linewidth(line_width)
            plt.tick_params(width=line_width, length=tick_len)

            mean_e = l_r_e / np.mean(l_r_e[60000:62000])
            mean_i = l_r_i / np.mean(l_r_i[60000:62000])

            plt.plot(mean_e, color='blue', linewidth=plot_line_width)
            plt.plot(mean_i, color='red', linewidth=plot_line_width)

            plt.xticks([60000, 62000, 64000, 66000, 68000], [3.0, 3.2, 3.4, 3.6, 3.8], fontsize=font_size_1,
                       **hfont)
            plt.yticks([0.85, 0.9, 0.95, 1.0, 1.05], fontsize=font_size_1, **hfont)
            plt.xlabel('Time (s)', fontsize=font_size_1, **hfont)
            plt.ylabel('Normalized firing rate', fontsize=font_size_1, **hfont)
            plt.xlim([60000, 68000])
            plt.ylim([0.85, 1.05])
            plt.legend(['Exc', 'Inh'], prop={"family": "Arial", 'size': font_size_1})
            plt.hlines(y=1, xmin=62000, xmax=70000, colors='k', linestyles=[(0, (6, 6, 6, 6))], linewidth=line_width)
            plt.savefig(
                'paper_figures/png/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_normalized_activity_paradoxical_effect_during_stimulation_ISN.png')
            plt.savefig(
                'paper_figures/pdf/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_normalized_activity_paradoxical_effect_during_stimulation_ISN.pdf')



if b_plotting_frozen_inhibition:
    # simulation setup
    dt = 0.0001
    T = int(9/dt)

    # neuronal parameters
    tau_e, tau_i = 0.020, 0.010
    alpha_e, alpha_i = 2, 2

    # adaptation
    x, u_d = 1, 1
    tau_x = 0.20

    # network connectivity
    Jee = 1.8
    Jie = 1.0
    Jei = 1.0
    Jii = 0.6

    l_b_before_stimulation = [True, False]

    l_g_e_ini = [1.55, 1.8]

    for g_e_ini in l_g_e_ini:

        for b_before_stimulation in l_b_before_stimulation:
            x = 1
            r_e, r_i = 0, 0
            z_e, z_i = 0, 0

            l_r_e, l_r_i, l_x = [], [], []

            for i in range(T):
                if b_before_stimulation:
                    g_e, g_i = g_e_ini, 2
                else:
                    if 50000 <= i:
                        g_e, g_i = 3.0, 2
                    else:
                        g_e, g_i = g_e_ini, 2

                if b_before_stimulation:
                    if 42000 <= i < 42001:
                        r_e = r_e + 0.01
                    else:
                        pass
                else:
                    if 62000 <= i < 62001:
                        r_e = r_e + 0.01
                    else:
                        pass

                g_e = g_e * (g_e > 0)
                g_i = g_i * (g_i > 0)

                # SSN part
                z_e = Jee * x * r_e - Jei * r_i + g_e
                z_i = Jie * r_e - Jii * r_i + g_i

                z_e = z_e * (z_e > 0)
                z_i = z_i * (z_i > 0)

                r_e = r_e + (-r_e + np.power(z_e, alpha_e)) / tau_e * dt
                if b_before_stimulation:
                    if 42000 < i:
                        pass
                    else:
                        r_i = r_i + (-r_i + np.power(z_i, alpha_i)) / tau_i * dt
                else:
                    if 62000 < i:
                        pass
                    else:
                        r_i = r_i + (-r_i + np.power(z_i, alpha_i)) / tau_i * dt

                r_e = r_e * (r_e > 0)
                r_i = r_i * (r_i > 0)

                x = x + ((1 - x) / tau_x - u_d * x * r_e) * dt
                x = np.clip(x, 0, 1)

                l_r_e.append(r_e)
                l_r_i.append(r_i)
                l_x.append(x)

            l_r_e = np.asarray(l_r_e)
            l_r_i = np.asarray(l_r_i)
            l_x = np.asarray(l_x)

            if b_before_stimulation:
                plt.figure(figsize=(figure_len, figure_width))
                ax = plt.gca()
                ax.spines['top'].set_visible(False)
                ax.spines['right'].set_visible(False)
                ax.spines['bottom'].set_visible(True)
                ax.spines['left'].set_visible(True)
                for axis in ['top', 'bottom', 'left', 'right']:
                    ax.spines[axis].set_linewidth(line_width)
                plt.tick_params(width=line_width, length=tick_len)
                plt.yscale('symlog', linthreshy=0.1)

                plt.plot(l_r_e, color='blue', linewidth=plot_line_width)
                plt.plot(l_r_i, color='red', linewidth=plot_line_width)

                plt.xticks(np.arange(30000, 90000 + 5000, 20000), np.arange(0, 6 + 0.5, 2), fontsize=font_size_1, **hfont)
                plt.yticks([0, 1, 100, 10000], fontsize=font_size_1, **hfont)
                plt.xlabel('Time (s)', fontsize=font_size_1, **hfont)
                plt.ylabel('Firing rate (Hz)', fontsize=font_size_1, **hfont)
                plt.xlim([30000, 90000])
                plt.ylim([0, 10000])
                plt.legend(['Exc', 'Inh'], prop={"family": "Arial", 'size': font_size_1}, loc='upper right')
                plt.vlines(x=42001, ymin=0, ymax=10000, colors='k', linestyles=[(0, (6, 6, 6, 6))],
                           linewidth=line_width)

                if g_e_ini == 1.55:
                    plt.savefig(
                        'paper_figures/png/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_normalized_activity_frozen_inhibition_before_stimulation_non_ISN.png')
                    plt.savefig(
                        'paper_figures/pdf/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_normalized_activity_frozen_inhibition_before_stimulation_non_ISN.pdf')

                else:
                    plt.savefig(
                        'paper_figures/png/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_normalized_activity_frozen_inhibition_before_stimulation_ISN.png')
                    plt.savefig(
                        'paper_figures/pdf/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_normalized_activity_frozen_inhibition_before_stimulation_ISN.pdf')

            else:
                plt.figure(figsize=(figure_len, figure_width))
                ax = plt.gca()
                ax.spines['top'].set_visible(False)
                ax.spines['right'].set_visible(False)
                ax.spines['bottom'].set_visible(True)
                ax.spines['left'].set_visible(True)
                for axis in ['top', 'bottom', 'left', 'right']:
                    ax.spines[axis].set_linewidth(line_width)
                plt.tick_params(width=line_width, length=tick_len)
                plt.yscale('symlog', linthreshy=0.1)

                plt.plot(l_r_e, color='blue', linewidth=plot_line_width)
                plt.plot(l_r_i, color='red', linewidth=plot_line_width)

                plt.xticks(np.arange(30000, 90000 + 5000, 20000), np.arange(0, 6 + 0.5, 2), fontsize=font_size_1, **hfont)
                plt.yticks([0, 1, 100, 10000], fontsize=font_size_1, **hfont)
                plt.xlabel('Time (s)', fontsize=font_size_1, **hfont)
                plt.ylabel('Firing rate (Hz)', fontsize=font_size_1, **hfont)
                plt.xlim([30000, 90000])
                plt.ylim([0, 10000])
                plt.legend(['Exc', 'Inh'], prop={"family": "Arial", 'size': font_size_1}, loc='upper right')
                plt.vlines(x=62001, ymin=0, ymax=10000, colors='k', linestyles=[(0, (6, 6, 6, 6))],
                           linewidth=line_width)

                if g_e_ini == 1.55:
                    plt.savefig(
                        'paper_figures/png/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_normalized_activity_frozen_inhibition_during_stimulation_non_ISN.png')
                    plt.savefig(
                        'paper_figures/pdf/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_normalized_activity_frozen_inhibition_during_stimulation_non_ISN.pdf')
                else:
                    plt.savefig(
                        'paper_figures/png/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_normalized_activity_frozen_inhibition_during_stimulation_ISN.png')
                    plt.savefig(
                        'paper_figures/pdf/Revision_Fig_Point_1_1_Supralinear_network_2D_EE_STP_normalized_activity_frozen_inhibition_during_stimulation_ISN.pdf')