import numpy as np
from rotational_path import PairRotationalPath, PathEnergyPlot
import expansion
from parameters import ModelParams
from pathlib import Path

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(beta1=pi_half_to_pi[::-1])
DipolePath.add_euler(beta1=zero_to_pi_half[::-1])
DipolePath.add_euler(beta1=zero_to_pi_half, beta2=zero_to_pi_half)
DipolePath.add_euler(beta1=np.pi/2, beta2=np.pi/2, alpha1=zero_to_pi_half)
DipolePath.add_euler(beta1=np.pi/2, alpha1=np.pi/2, beta2=pi_half_to_pi)
DipolePath.add_euler(beta1=np.pi/2, beta2=pi_half_to_pi[::-1], alpha2=np.pi)
DipolePath.add_euler(beta1=zero_to_pi_half[::-1], beta2=pi_half_to_pi, alpha2=np.pi)
DipolePath.add_euler(beta1=zero_to_pi_half, beta2=pi_half_to_pi[::-1], alpha2=np.pi)
DipolePath.add_euler(beta1=pi_half_to_pi, beta2=zero_to_pi_half[::-1], alpha2=np.pi)

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


if __name__ == '__main__':

    params = ModelParams(R=150, kappaR=3)
    # ex1 = expansion.MappedExpansionQuad(np.array([0.35, 0.44, 0.6]), params.kappaR, 0.001, max_l=20)
    # ex1 = expansion.Expansion24(sigma2=0.001, sigma4=0)
    # ex1 = expansion.Expansion(l_array=np.array([1]), coefs=expansion.rot_sym_expansion(np.array([1]), np.array([0.001])))
    # ex1 = expansion.GaussianCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=np.array([5]), sigma1=0.001, l_max=10)
    ex1 = expansion.GaussianCharges(omega_k=np.array([np.pi, 0]), lambda_k=np.array([1, 5, 10, 15]), sigma1=0.001, l_max=10)
    ex2 = ex1.clone()

    path_plot = PathEnergyPlot(ex1, ex2, DipolePath2, dist=2., params=params)
    path_plot.plot(labels=[rf'$\lambda$={l}' for l in [1, 5, 10, 15]],
                   # norm_euler_angles={'beta1': np.pi/2}
                   )
    # path_plot.plot_sections(save_as=Path('/home/andraz/ChargedShells/Figures/dipole_path2.png'))