import matplotlib.pyplot as plt
from charged_shells import expansion, interactions, mapping, functions
from charged_shells.parameters import ModelParams
import numpy as np
from typing import Literal
from pathlib import Path
from functools import partial
import json
from itertools import cycle

Array = np.ndarray
Expansion = expansion.Expansion
RotConfig = Literal['ep', 'pp']


def peak_energy_plot(kappaR: Array,
                     a_bar: Array,
                     which: Literal['ep', 'pp'],
                     R: float = 150,
                     dist: float = 2.,
                     l_max=20,
                     save_as: Path = None):

    params = ModelParams(R=R, kappaR=kappaR)

    ex1 = expansion.MappedExpansionQuad(a_bar[:, None], params.kappaR[None, :], 0.001, l_max=l_max)
    ex2 = ex1.clone()
    if which == 'ep':
        ex1.rotate_euler(alpha=0, beta=np.pi / 2, gamma=0)

    energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=dist), 1)
    energy = energy_fn(ex1, ex2, params)

    fig, ax = plt.subplots()
    for en, ab in zip(energy, a_bar):
        ax.plot(kappaR, en / en[0], label=rf'$\bar a = {ab:.1f}$')
        # ax.plot(kappaR, en, label=rf'$\bar a = {ab:.1f}$')
    ax.legend(fontsize=12)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
    ax.set_xlabel(r'$\kappa R$', fontsize=15)
    ax.set_ylabel(rf'$\bar V$', fontsize=15)
    plt.tight_layout()
    if save_as is not None:
        plt.savefig(save_as, dpi=600)
    plt.show()


def IC_peak_energy_plot(config_data: dict,
                        a_bar: list,
                        which: RotConfig,
                        max_kappaR: float = 30.,
                        R: float = 150,
                        save_as: Path = None,
                        save_data: bool = False,
                        quad_correction: bool = False):

    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_11"))
    em_data = np.load(em_data_path.joinpath("pair_energy_fig11.npz"))
    data = em_data['fixA']

    if which == 'ep':
        column_idx = 4
    elif which == 'pp':
        column_idx = 3
    else:
        raise ValueError

    abar, inverse, counts = np.unique(data[:, 1], return_counts=True, return_inverse=True)
    all_kappaR = np.unique(data[:, 2])

    kappaR = np.geomspace(np.min(all_kappaR), max_kappaR, 30)
    params = ModelParams(R=R, kappaR=kappaR)

    ex1 = expansion.MappedExpansionQuad(np.array(a_bar)[:, None], params.kappaR[None, :], 0.001, l_max=20)
    ex2 = ex1.clone()
    if which == 'ep':
        ex1.rotate_euler(alpha=0, beta=np.pi / 2, gamma=0)

    energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=2), 1)
    energy = energy_fn(ex1, ex2, params)

    data_dict_ic = {}
    data_dict_cs = {}
    k = 0
    for i in range(len(abar)):
        ab = np.around(abar[i], 5)
        if ab in np.around(a_bar, 5):
            idx, = np.nonzero(inverse == i)
            kR = data[idx, 2][data[idx, 2] <= max_kappaR]
            en = data[idx, column_idx][data[idx, 2] <= max_kappaR]
            if quad_correction:
                en *= extra_correction(ab, kR)
            sort = np.argsort(kR)
            data_dict_ic[ab] = np.stack((kR[sort], np.abs(en)[sort])).T
            data_dict_cs[ab] = np.stack((kappaR, np.abs(energy[k]))).T
            k += 1

    if save_data:
        ic_save = {str(key): val for key, val in data_dict_ic.items()}
        cs_save = {str(key): val for key, val in data_dict_cs.items()}
        ic_save['abar'] = a_bar
        cs_save['abar'] = a_bar
        with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
            config_data = json.load(config_file)
        np.savez(Path(config_data["figure_data"]).joinpath("fig_11_IC.npz"), **ic_save)
        np.savez(Path(config_data["figure_data"]).joinpath("fig_11_CS.npz"), **cs_save)

    colors = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])

    fig, ax = plt.subplots()
    for (key, data1), data2 in zip(data_dict_ic.items(), data_dict_cs.values()):
        current_color = next(colors)
        ax.plot(data1[:, 0], data1[:, 1], label=rf'$\bar a = {key:.2f}$', c=current_color)
        ax.plot(data2[:, 0], data2[:, 1], ls='--', c=current_color)
    ax.legend(fontsize=14, ncol=2)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=15)
    ax.set_xlabel(r'$\kappa R$', fontsize=15)

    y_label = r'$U_{pp}$' if which == 'pp' else r'$|U_{ep}|$'
    ax.set_ylabel(y_label, fontsize=15)
    ax.set_yscale('log')
    ax.set_xscale('log')
    plt.tight_layout()
    if save_as is not None:
        plt.savefig(save_as, dpi=600)
    plt.show()


def extra_correction(a_bar, kappaR):
    y = a_bar * kappaR
    # return (kappaR ** 2 * a_bar ** 4 * functions.sph_bessel_k(1, kappaR) / functions.sph_bessel_k(3, kappaR) *
    #         np.sinh(y) / (-3 * y * np.cosh(y) + (3 + y ** 2) * np.sinh(y)))
    return (a_bar ** 2 * functions.sph_bessel_k(1, kappaR) / functions.sph_bessel_k(3, kappaR)
            * functions.sph_bessel_i(0, y) / functions.sph_bessel_i(2, y))


def main():
    with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
        config_data = json.load(config_file)

    # kappaR = np.arange(0.5, 50, 0.1)
    # a_bar = np.arange(0.2, 0.8, 0.2)
    # peak_energy_plot(kappaR, a_bar, which='pp',
    #                  # save_as=Path('/home/andraz/ChargedShells/Figures/nonmonotonicity_check_ep.pdf')
    #                  )

    a_bar = [0.2, 0.32, 0.52, 0.8]
    # a_bar = list(np.arange(0.2, 0.81, 0.04))
    IC_peak_energy_plot(config_data, a_bar=a_bar, which='pp', max_kappaR=50,
                        save_as=Path('/home/andraz/ChargedShells/Figures/Emanuele_data/both_nonmonotonicity_check_pp.png'),
                        save_data=False,
                        quad_correction=False
                        )

    # kr = np.linspace(0.1, 100, 1000)
    # for ab in [0.2, 0.32, 0.52, 0.8]:
    #     plt.plot(kr, extra_correction(ab, kr))
    # plt.xscale('log')
    # plt.yscale('log')
    # plt.show()


if __name__ == '__main__':
    
    main()