from __future__ import annotations

from abc import ABC, abstractmethod

import numpy as np
from numpy.random import Generator
from numpy.typing import ArrayLike

from .helper import xyz_to_spherical, spherical_to_xyz, omega_q, draw_orientation, get_rotation_matrix


class BaseMotion(ABC):
    def __init__(self, delta: float, eta: float, rng: Generator):
        self._delta = delta
        self._eta = eta

        self._rng = rng

    @abstractmethod
    def __repr__(self):
        pass

    def start(self):
        pass

    @abstractmethod
    def jump(self, size: int = 1) -> ArrayLike:
        pass


class RandomJump(BaseMotion):

    def __repr__(self):
        return 'Random Jump'

    def jump(self, size: int = 1) -> ArrayLike:
        return omega_q(self._delta, self._eta, *draw_orientation(self._rng, size=size))


class TetrahedralJump(BaseMotion):
    def __init__(self, delta: float, eta: float, rng: Generator):
        super().__init__(delta, eta, rng)

        self._orientation = None
        self._start = None

    def __repr__(self):
        return 'Tetrahedral Jump'

    def start(self):
        self._orientation = self._make_tetrahedron()
        self._start = self._rng.choice([0, 1, 2, 3])

    def _make_tetrahedron(self) -> np.ndarray:
        beta = np.arccos(-1/3)  # tetrahedral angle 109.5 degrees
        sin_beta = np.sin(beta)
        cos_beta = np.cos(beta)

        # corners of a tetrahedron
        alpha = 2 * np.pi * self._rng.random()
        corners = np.array([
            [0, 0, 1],
            [sin_beta * np.cos(alpha),           sin_beta * np.sin(alpha),           cos_beta],
            [sin_beta * np.cos(alpha+2*np.pi/3), sin_beta * np.sin(alpha+2*np.pi/3), cos_beta],
            [sin_beta * np.cos(alpha+4*np.pi/3), sin_beta * np.sin(alpha+4*np.pi/3), cos_beta]
        ])

        # orientation in lab system
        theta0, phi0 = draw_orientation(self._rng)

        rot = get_rotation_matrix(
            corners[0],
            np.array(spherical_to_xyz(1., theta0, phi0)),
        )
        orientations = np.zeros(4)
        for i in range(4):
            corner_lab = np.dot(rot, corners[i])
            _, theta_i, phi_i = xyz_to_spherical(*corner_lab)
            orientations[i] = omega_q(self._delta, self._eta, theta_i, phi_i)

        return orientations

    def jump(self, size: int = 1) -> ArrayLike:
        jumps = self._rng.choice([1, 2, 3], size=size)
        jumps = np.cumsum(jumps) + self._start
        jumps %= 4

        self._start = jumps[-1]

        return self._orientation[jumps]