import numpy as np
from charged_shells.rotational_path import PairRotationalPath, PathEnergyPlot
from charged_shells import expansion, patch_size, interactions
from charged_shells.parameters import ModelParams
import matplotlib.pyplot as plt
from pathlib import Path
import json

Array = np.ndarray


zero_to_pi_half = np.linspace(0, np.pi/2, 100, endpoint=True)
pi_half_to_pi = np.linspace(np.pi/2, np.pi, 100, endpoint=True)

QuadPath = PairRotationalPath()
QuadPath.set_default_x_axis(zero_to_pi_half)
QuadPath.add_euler(beta1=np.pi/2, beta2=zero_to_pi_half)
QuadPath.add_euler(beta1=zero_to_pi_half[::-1], beta2=zero_to_pi_half[::-1])
QuadPath.add_euler(beta1=zero_to_pi_half)
QuadPath.add_euler(beta1=zero_to_pi_half[::-1], beta2=zero_to_pi_half)
QuadPath.add_euler(beta1=np.pi/2, beta2=zero_to_pi_half, alpha2=np.pi/2)
QuadPath.add_euler(beta1=np.pi/2, beta2=np.pi/2, alpha1=zero_to_pi_half[::-1])

DipolePath = PairRotationalPath()
DipolePath.set_default_x_axis(zero_to_pi_half)
DipolePath.add_euler(beta2=pi_half_to_pi[::-1])
DipolePath.add_euler(beta2=zero_to_pi_half[::-1])
DipolePath.add_euler(beta2=zero_to_pi_half, beta1=zero_to_pi_half)
DipolePath.add_euler(beta2=np.pi/2, beta1=np.pi/2, alpha2=zero_to_pi_half)
DipolePath.add_euler(beta2=np.pi/2, alpha2=np.pi/2, beta1=pi_half_to_pi)
DipolePath.add_euler(beta2=np.pi/2, beta1=pi_half_to_pi[::-1], alpha1=np.pi)
DipolePath.add_euler(beta2=zero_to_pi_half[::-1], beta1=pi_half_to_pi, alpha1=np.pi)
DipolePath.add_euler(beta2=zero_to_pi_half, beta1=pi_half_to_pi[::-1], alpha1=np.pi)
DipolePath.add_euler(beta2=pi_half_to_pi, beta1=zero_to_pi_half[::-1], alpha1=np.pi)

DipolePath2 = PairRotationalPath()
DipolePath2.set_default_x_axis(zero_to_pi_half)
DipolePath2.add_euler(beta2=pi_half_to_pi[::-1])
DipolePath2.add_euler(beta2=zero_to_pi_half[::-1])
DipolePath2.add_euler(beta2=zero_to_pi_half, beta1=zero_to_pi_half)
DipolePath2.add_euler(beta2=np.pi/2, beta1=np.pi/2, alpha2=zero_to_pi_half)
DipolePath2.add_euler(beta2=np.pi/2, alpha2=np.pi/2, beta1=pi_half_to_pi)
DipolePath2.add_euler(beta2=zero_to_pi_half[::-1], beta1=pi_half_to_pi[::-1])
DipolePath2.add_euler(beta2=zero_to_pi_half[::-1], beta1=np.pi)
DipolePath2.add_euler(beta2=zero_to_pi_half, beta1=pi_half_to_pi[::-1], alpha1=np.pi)
DipolePath2.add_euler(beta2=pi_half_to_pi, beta1=zero_to_pi_half[::-1], alpha1=np.pi)


def model_comparison(save_as=None):
    params = ModelParams(R=150, kappaR=3)
    abar = 0.5
    sigma_tilde = 0.001

    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
    ex2 = ex1.clone()

    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params)
    energy = path_plot.evaluate_path()
    x_axis = path_plot.rot_path.stack_x_axes()

    # peak_energy_sanity_check
    ex1new = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
    ex2new = ex1new.clone()
    pp_energy = interactions.charged_shell_energy(ex1new, ex2new, params)
    print(f'PP energy: {pp_energy}')

    # Emanuele data
    with open(Path("/home/andraz/ChargedShells/charged-shells/config.json")) as config_file:
        config_data = json.load(config_file)
    em_data = np.load(Path(config_data["emanuele_data"]).joinpath("FIG_7").joinpath("pathway.npz"))['arr_0']

    fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
    ax.plot(em_data[:, 0], em_data[:, 1], label='ICi')
    ax.plot(x_axis, np.squeeze(energy), label='CSp')
    # ax.plot(x_axis, em_data[:, 1] / np.squeeze(energy), label='CSp')
    PathEnergyPlot.plot_style(fig, ax)
    if save_as is not None:
        plt.savefig(save_as, dpi=600)
    plt.show()


def kappaR_dependence(kappaR: Array, abar: float, sigma_tilde=0.001, save_as=None):
    params = ModelParams(R=150, kappaR=kappaR)

    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
    ex2 = ex1.clone()

    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params, match_expansion_axis_to_params=0)

    path_plot.plot(labels=[rf'$\kappa R$={kR}' for kR in kappaR],
                   # norm_euler_angles={'beta2': np.pi},
                   save_as=save_as)


def abar_dependence(abar: Array, kappaR: float, sigma_tilde=0.001, save_as=None):
    params = ModelParams(R=150, kappaR=kappaR)

    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
    ex2 = ex1.clone()

    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params, match_expansion_axis_to_params=None)

    path_plot.plot(labels=[rf'$\bar a$={a}' for a in abar],
                   # norm_euler_angles={'beta2': np.pi},
                   save_as=save_as)


def sigma0_dependence(sigma0: Array, kappaR: float, abar: float, sigma_tilde=0.001, save_as=None):
    params = ModelParams(R=150, kappaR=kappaR)

    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30, sigma0=sigma0)
    ex2 = ex1.clone()

    path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=2., params=params, match_expansion_axis_to_params=None)

    path_plot.plot(labels=[rf'$\tilde \sigma_0={s0}$' for s0 in sigma0],
                   # norm_euler_angles={'beta2': np.pi},
                   save_as=save_as)


def distance_dependence(dist: Array, kappaR: float, abar: float, sigma_tilde=0.001, save_as=None):
    params = ModelParams(R=150, kappaR=kappaR)

    ex1 = expansion.MappedExpansionQuad(abar, params.kappaR, sigma_tilde, l_max=30)
    ex2 = ex1.clone()

    plots = []
    for d in dist:
        path_plot = PathEnergyPlot(ex1, ex2, QuadPath, dist=d, params=params)
        x = d * kappaR
        plots.append(path_plot.evaluate_path() * np.exp(x) * x)

    x_axis = path_plot.rot_path.stack_x_axes()
    labels = [rf'$\rho/R ={d}$' for d in dist]

    fig, ax = plt.subplots(figsize=plt.figaspect(0.5))
    for pl, lbl in zip(plots, labels):
        ax.plot(x_axis, pl, label=lbl)
    PathEnergyPlot.plot_style(fig, ax)
    ax.set_ylabel(r'$U \kappa\rho e^{\kappa\rho}$', fontsize=15)
    if save_as is not None:
        plt.savefig(save_as, dpi=600)
    plt.show()


def IC_kappaR_dependence(config_data: dict, save_as=None):

    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_8").joinpath("FIXEDCHARGE")
                    .joinpath("FIX_A").joinpath("ECC_0.25"))
    kR1 = np.load(em_data_path.joinpath("EMME_1.").joinpath("pathway.npz"))['arr_0']
    kR3 = np.load(em_data_path.joinpath("EMME_3.").joinpath("pathway.npz"))['arr_0']
    kR10 = np.load(em_data_path.joinpath("EMME_10.").joinpath("pathway.npz"))['arr_0']

    labels = [rf'$\kappa R$={kR}' for kR in [1, 3, 10]]

    fig, ax = plt.subplots()
    ax.plot(kR1[:, 0], kR1[:, 1], label=labels[0])
    ax.plot(kR3[:, 0], kR3[:, 1], label=labels[1])
    ax.plot(kR10[:, 0], kR10[:, 1], label=labels[2])
    PathEnergyPlot.plot_style(fig, ax)
    if save_as is not None:
        plt.savefig(save_as, dpi=600)
    plt.show()


def IC_abar_dependence(config_data: dict, save_as=None):

    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_8").joinpath("FIXEDCHARGE").joinpath("FIX_M"))
    a03 = np.load(em_data_path.joinpath("ECC_0.15").joinpath("EMME_3.").joinpath("pathway.npz"))['arr_0']
    a04 = np.load(em_data_path.joinpath("ECC_0.2").joinpath("EMME_3.").joinpath("pathway.npz"))['arr_0']
    a05 = np.load(em_data_path.joinpath("ECC_0.25").joinpath("EMME_3.").joinpath("pathway.npz"))['arr_0']

    labels =[rf'$\bar a$={a}' for a in [0.3, 0.4, 0.5]]

    fig, ax = plt.subplots()
    ax.plot(a03[:, 0], a03[:, 1], label=labels[0])
    ax.plot(a04[:, 0], a04[:, 1], label=labels[1])
    ax.plot(a05[:, 0], a05[:, 1], label=labels[2])
    PathEnergyPlot.plot_style(fig, ax)
    if save_as is not None:
        plt.savefig(save_as, dpi=600)
    plt.show()


def IC_sigma0_dependence(config_data: dict, save_as=None):
    em_data_path = (Path(config_data["emanuele_data"]).joinpath("FIG_8").joinpath("CHARGE_ZC"))
    undercharged = np.load(em_data_path.joinpath("ZC_-277.27").joinpath("pathway.npz"))['arr_0']
    neutral = np.load(em_data_path.joinpath("ZC_-560").joinpath("pathway.npz"))['arr_0']
    overchargerd = np.load(em_data_path.joinpath("ZC_-842.74").joinpath("pathway.npz"))['arr_0']

    labels = [rf'$\tilde \sigma_0={s0}$' for s0 in [-0.001, 0, 0.001]]

    fig, ax = plt.subplots()
    ax.plot(overchargerd[:, 0], overchargerd[:, 1], label=labels[0])
    ax.plot(neutral[:, 0], neutral[:, 1], label=labels[1])
    ax.plot(undercharged[:, 0], undercharged[:, 1], label=labels[2])
    PathEnergyPlot.plot_style(fig, ax)
    if save_as is not None:
        plt.savefig(save_as, dpi=600)
    plt.show()


def main():

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

    # model_comparison(
    #     # save_as=Path(config_data["figures"]).joinpath("Emanuele_data").joinpath('IC_CS_quadrupole_path.pdf')
    # )

    # kappaR_dependence(np.array([1, 3, 10]), 0.5,
    #                   # save_as=Path("/home/andraz/ChargedShells/Figures/quadrupole_kappaR_dep.png")
    #                   )
    #
    # abar_dependence(np.array([0.3, 0.4, 0.5]), 3,
    #                 save_as=Path("/home/andraz/ChargedShells/Figures/quadrupole_abar_dep.png"))

    sigma0_dependence(np.array([0.001, 0.002, 0.003]), 3, 0.5,
                      # save_as=Path("/home/andraz/ChargedShells/Figures/quadrupole_charge_dep_abar05_kappaR3.png")
    )

    # distance_dependence(dist=np.array([2, 3, 4, 6, 10, 20]), kappaR=3, abar=0.5,
    #                     # save_as=Path(config_data["figures"]).joinpath('quadrupole_distance_dep.png')
    #                     )

    # IC_kappaR_dependence(config_data,
    #                      # save_as=Path(config_data["figures"]).joinpath("Emanuele_data"). joinpath('quadrupole_kappaR_dep.png')
    #                      )
    #
    # IC_abar_dependence(config_data, save_as=Path(config_data["figures"]).joinpath("Emanuele_data").
    #                    joinpath('quadrupole_abar_dep.png'))
    #
    # IC_sigma0_dependence(config_data, save_as=Path(config_data["figures"]).joinpath("Emanuele_data").
    #                      joinpath('quadrupole_charge_dep_abar05_kappaR3.png'))


if __name__ == '__main__':

    main()