import matplotlib.pyplot as plt
import numpy as np
from matplotlib import rc

from threefish.defaults import default_figsize, default_settings
from threefish.load import save_visualization
from threefish.plot.plotstyle import plot_style
from threefish.plot.tags import tag2
from threefish.RAM.calc_analytics import calc_nonlin_analytic_values
from threefish.RAM.reformat_matrix import convert_csv_str_to_float

rc('text', usetex=True)


def plot_chi2():



    default_settings()  # ts=13, ls=13, fs=13, lw = 0.7
    plot_style()
    default_figsize(column=2, length=2.5)
    fig, ax = plt.subplots(1, 2)
    plt.subplots_adjust(bottom=0.2, wspace=0.35, left=0.1, right=0.92)

    ###################################################
    transfer_values, suscept_values = calc_nonlin_analytic_values()

    ###################################################
    # plot transfer function, panel A
    transfer_values = transfer_values.astype(complex)
    ax[0].plot(transfer_values.index, np.abs(transfer_values['transfer']),
               color='black')
    ax[0].set_xlabel(r'$f$')
    ax[0].set_ylabel(r'$|\chi_{1}(f)|$')
    ax[0].set_xlim(0, transfer_values.index[-1])
    ylim = ax[0].get_ylim()
    ax[0].set_ylim(0, ylim[-1])

    ###################################################
    # plot matrix, panel B
    suscept_values.columns = suscept_values.index
    new_keys, stack_plot = convert_csv_str_to_float(suscept_values)
    prss = np.abs(stack_plot)
    vmax = np.quantile(prss, 0.994)
    ten = 10**np.floor(np.log10(vmax))
    for fac, delta in zip([1, 2, 3, 4, 6, 8, 10],
                          [0.5, 1, 1, 2, 3, 4, 5]):
        if fac*ten >= vmax:
            vmax = fac*ten
            ten *= delta
            break
    pc = ax[1].pcolormesh(new_keys, new_keys, prss, rasterized=True,
                          cmap='viridis', vmin=0, vmax=vmax)
    ax[1].set_aspect('equal')
    cax = ax[1].inset_axes([1.04, 0, 0.05, 1])
    cax.set_spines_outward('lrbt', 0)
    fig.colorbar(pc, cax=cax, label=r'$|\chi_2(f_1, f_2)|$')
    cax.set_yticks_delta(50)
    ax[1].set_xlabel(r'$f_{1}$')
    ax[1].set_ylabel(r'$f_{2}$')
    ax[1].set_xlim(0, 1)
    ax[1].set_ylim(0, 1)
    ax[1].set_xticks_delta(0.2)
    ax[1].set_xticks_delta(0.2)
    tag2(axes=ax, xoffs=[-0.5, -5.5], yoffs=1.7)  # ax_ams[3],
    save_visualization()


if __name__ == '__main__':

    plot_chi2()