from __future__ import annotations

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


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


def draw_orientation(delta: float, eta: float, rng: Generator, size: int = 1) -> ArrayLike:
    z_theta, z_phi = rng.random((2, size))
    theta = np.arccos(1 - 2 * z_theta)
    phi = 2 * np.pi * z_phi

    return omega_q(delta, eta, theta, phi)


class RandomJump:
    def __init__(self, delta: float, eta: float, rng: Generator | None = None):
        self._delta = delta
        self._eta = eta

        self._rng = rng

    def __repr__(self):
        return 'Random Jump'

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