import matplotlib.pyplot as plt
import expansion
import interactions
from parameters import ModelParams
import numpy as np
from typing import Literal
from pathlib import Path

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)

    energy = []
    for params in params.unravel():
        ex1 = expansion.MappedExpansionQuad(a_bar, params.kappaR, 0.001, l_max=l_max)
        ex2 = ex1.clone()
        if which == 'ep':
            ex1.rotate_euler(alpha=0, beta=np.pi / 2, gamma=0)
        energy.append(interactions.charged_shell_energy(ex1, ex2, dist, params))

    energy = np.array(energy)

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


if __name__ == '__main__':

    kappaR = np.arange(1, 10, 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_pp.pdf')
                     )