import numpy as np
from dataclasses import dataclass, field
import quaternionic
import expansion
from parameters import ModelParams
import interactions
import matplotlib.pyplot as plt


Quaternion = quaternionic.array
Array = np.ndarray


@dataclass
class PairRotationalPath:

    rotations1: list[Quaternion] = field(default_factory=list)
    rotations2: list[Quaternion] = field(default_factory=list)

    def add(self, rotation1: Quaternion, rotation2: Quaternion):
        self.rotations1.append(rotation1)
        self.rotations2.append(rotation2)

    def add_euler(self, alpha1: Array = 0, beta1: Array = 0, gamma1: Array = 0,
                  alpha2: Array = 0, beta2: Array = 0, gamma2: Array = 0):
        R1_euler = quaternionic.array.from_euler_angles(alpha1, beta1, gamma1)
        R2_euler = quaternionic.array.from_euler_angles(alpha2, beta2, gamma2)
        R1_euler, R2_euler = np.broadcast_arrays(R1_euler, R2_euler)
        self.add(Quaternion(R1_euler), Quaternion(R2_euler))

    def get_rotations(self) -> (Quaternion, Quaternion):
        return Quaternion(np.vstack(self.rotations1)), Quaternion(np.vstack(self.rotations2))


zero_to_pi_half = np.linspace(0, np.pi/2, 1000)

QuadPath = PairRotationalPath()
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, alpha2=zero_to_pi_half[::-1])


if __name__ == '__main__':

    params = ModelParams(R=150, kappaR=3)
    ex1 = expansion.MappedExpansion(1, params.kappaR, 0.001, max_l=15)
    # ex1 = expansion.Expansion24(sigma2=0.001, sigma4=0)
    ex2 = ex1.clone()

    rotations1, rotations2 = QuadPath.get_rotations()
    ex1.rotate(rotations1)
    ex2.rotate(rotations2)

    dist = 3

    energy = interactions.charged_shell_energy(ex1, ex2, dist, params)

    plt.plot(energy)
    plt.show()