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.colorbar import colorbar_outside
from threefish.plot.limits import set_clim_same
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.plot_subplots import plt_RAM_perc
from threefish.RAM.reformat_matrix import convert_csv_str_to_float
from threefish.RAM.values import perc_model_full

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)
    fig1, 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_{1}$')
    ax[0].set_ylabel(r'$|\chi_{1}|$')
    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)
    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())

    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(suscept_values.columns[0] - 0.01, suscept_values.columns[-1])
    ax[1].set_ylim(suscept_values.columns[0] - 0.01, suscept_values.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()