from __future__ import annotations

from abc import ABC, abstractmethod

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


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

    @property
    def name(self) -> str:
        return self.__class__.__name__

    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
        cos_theta0, phi0 = draw_orientation(self._rng)

        rot = get_rotation_matrix(
            corners[0],
            np.array(spherical_to_xyz(1., np.arccos(cos_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]


# Helper functions

def xyz_to_spherical(x_in: float, y_in: float, z_in: float) -> tuple[float, float, float]:
    r = np.linalg.norm([x_in, y_in, z_in])
    theta = np.arccos(z_in)
    phi = np.arctan2(y_in, x_in)

    return r, theta, phi


def spherical_to_xyz(r: float, theta: float, phi: float) -> tuple[float, float, float]:
    sin_theta = np.sin(theta)
    return r*np.cos(phi)*sin_theta, r*np.sin(phi)*sin_theta, r*np.cos(theta)


def get_rotation_matrix(vec_in: np.ndarray, vec_out: np.ndarray):
    rotation = np.eye(3)

    # rotation by angle around given axis
    cos_angle = np.dot(vec_in, vec_out)

    # check for parallel / anti-parallel vectors
    if cos_angle == 1:
        return rotation
    elif cos_angle == -1:
        return -rotation

    else:
        axis = np.cross(vec_in, vec_out)
        scale = np.linalg.norm(axis)
        axis /= scale
        sin_angle = scale / np.linalg.norm(vec_in) / np.linalg.norm(vec_out)

        v_cross = np.array([
            [0,       -axis[2], axis[1]],
            [axis[2], 0,        -axis[0]],
            [-axis[1], axis[0], 0],
        ])

        rotation += sin_angle * v_cross
        rotation += (1-cos_angle) * v_cross @ v_cross

        return rotation


def omega_q(delta: float, eta: float, cos_theta: ArrayLike, phi: ArrayLike) -> ArrayLike:
    # sin_theta = np.sin(cos_theta)
    # cos_theta = np.cos(cos_theta)
    sin_theta_sq = 1 - cos_theta**2
    return np.pi * delta * (3 * cos_theta**2 - 1 + eta * sin_theta_sq * np.cos(2*phi))


def draw_orientation(rng: Generator, size: int | None = None) -> tuple[ArrayLike, ArrayLike]:
    if size is not None:
        z_theta, z_phi = rng.random((2, size))
    else:
        z_theta, z_phi = rng.random(2)
    cos_theta = 1 - 2 * z_theta
    phi = 2 * np.pi * z_phi

    return cos_theta, phi