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):
        rotation1, rotation2 = np.broadcast_arrays(rotation1, rotation2)
        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)
        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, 100, endpoint=False)
pi_half_to_pi = np.linspace(np.pi/2, np.pi, 100, endpoint=False)

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, alpha1=zero_to_pi_half[::-1])

DipolePath = PairRotationalPath()
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=zero_to_pi_half)
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)


if __name__ == '__main__':

    params = ModelParams(R=150, kappaR=3)
    # ex1 = expansion.MappedExpansionQuad(0.44, 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.SmearedCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=5, sigma1=0.001, l_max=10)
    ex1 = expansion.SmearedCharges(omega_k=np.array([np.pi, 0]), lambda_k=3, sigma1=0.001, l_max=10)
    ex2 = ex1.clone()

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

    dist = 2

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

    plt.plot(energy)
    plt.show()