from __future__ import annotations

import math
from typing import Iterator

import numpy as np
from scipy.spatial.transform import Rotation as R


class Vector3:
    """A 3D vector defined by x, y, and z values.

    By convention, units for world space are in meters.

    Args:
        x (float, optional): The x value. Defaults to 0.
        y (float, optional): The y value. Defaults to 0.
        z (float, optional): The z value. Defaults to 0.
    """

    def __init__(self, x: float = 0, y: float = 0, z: float = 0) -> None:
        self.x = x
        self.y = y
        self.z = z

    @staticmethod
    def zero() -> Vector3:
        """Returns a zero vector `Vector3(0, 0, 0)`."""
        return Vector3(0, 0, 0)

    @staticmethod
    def one() -> Vector3:
        """Returns a one vector `Vector3(1, 1, 1)`.

        This is useful for setting a default scale.
        """
        return Vector3(1, 1, 1)

    def to_direction(self) -> Vector3:
        """Converts a rotation Vector3 (roll, pitch, yaw) to a directional unit Vector3."""
        pitch = math.radians(self.y)
        yaw = math.radians(self.z)
        x = math.cos(yaw) * math.cos(pitch)
        y = math.sin(yaw) * math.cos(pitch)
        z = math.sin(pitch)
        return Vector3(x, y, z).normalized()

    def to_rotation(self) -> Vector3:
        """Converts a directional Vector3 to a rotation Vector3 (roll, pitch, yaw).

        Roll is always 0.
        """
        roll = 0
        pitch = math.degrees(math.atan2(self.z, math.sqrt(self.x**2 + self.y**2)))
        yaw = math.degrees(math.atan2(self.y, self.x))
        return Vector3(roll, pitch, yaw)

    def numpy(self) -> np.ndarray:
        """Returns the vector as a NumPy array of shape (3,)."""
        return np.array([self.x, self.y, self.z])

    def magnitude(self) -> float:
        """Returns the magnitude of the vector."""
        sq = self.x**2 + self.y**2 + self.z**2
        return sq**0.5

    def normalized(self) -> Vector3:
        """Returns a new normalized vector."""
        mag = self.magnitude()
        if mag == 0:
            return Vector3.zero()
        return self / mag

    def cross(self, other: Vector3) -> Vector3:
        """Returns the cross product of two vectors."""
        return Vector3(
            self.y * other.z - self.z * other.y,
            self.z * other.x - self.x * other.z,
            self.x * other.y - self.y * other.x,
        )

    def dot(self, other: Vector3) -> float:
        """Returns the dot product of two vectors."""
        return self.x * other.x + self.y * other.y + self.z * other.z

    def isclose(
        self, other: Vector3, rel_tol: float = 1.0e-5, abs_tol: float = 1.0e-8
    ) -> bool:
        """Returns True if two vectors are element-wise equal within a tolerance.

        The tolerance values are positive, typically very small numbers.  The
        relative difference (`rel_tol` * abs(`b`)) and the absolute difference
        `abs_tol` are added together to compare against the absolute difference
        between `a` and `b`.

        NaNs are treated as equal if they are in the same place and if
        ``equal_nan=True``.  Infs are treated as equal if they are in the same
        place and of the same sign in both arrays.

        Args:
            other (Vector3): The other vector to compare.
            rel_tol (float, optional): The relative tolerance. Defaults to 1.e-5.
            abs_tol (float, optional): The absolute tolerance. Defaults to 1.e-8.

        """

        return (
            math.isclose(self.x, other.x, rel_tol=rel_tol, abs_tol=abs_tol)
            and math.isclose(self.y, other.y, rel_tol=rel_tol, abs_tol=abs_tol)
            and math.isclose(self.z, other.z, rel_tol=rel_tol, abs_tol=abs_tol)
        )

    def __add__(self, other: Vector3) -> Vector3:
        return Vector3(self.x + other.x, self.y + other.y, self.z + other.z)

    def __sub__(self, other: Vector3) -> Vector3:
        return Vector3(self.x - other.x, self.y - other.y, self.z - other.z)

    def __mul__(self, other: float | Vector3) -> Vector3:
        if isinstance(other, Vector3):
            return Vector3(self.x * other.x, self.y * other.y, self.z * other.z)
        elif isinstance(other, (int, float)):
            return Vector3(self.x * other, self.y * other, self.z * other)
        else:
            return NotImplemented

    def __rmul__(self, other: float | Vector3) -> Vector3:
        return self.__mul__(other)

    def __truediv__(self, other: float | Vector3) -> Vector3:
        if isinstance(other, Vector3):
            return Vector3(self.x / other.x, self.y / other.y, self.z / other.z)
        elif isinstance(other, (int, float)):
            return Vector3(self.x / other, self.y / other, self.z / other)
        else:
            return NotImplemented

    def __matmul__(self, other: Vector3) -> float:
        return self.x * other.x + self.y * other.y + self.z * other.z

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Vector3):
            return NotImplemented
        return self.isclose(other)

    def __ne__(self, other: object) -> bool:
        if not isinstance(other, Vector3):
            return NotImplemented
        return not self.__eq__(other)

    def __repr__(self) -> str:
        return f"Vector3({round(self.x, 2)}, {round(self.y, 2)}, {round(self.z, 2)})"

    def __iter__(self) -> Iterator[float]:
        return iter((self.x, self.y, self.z))

    def __round__(self, n: int = 8) -> Vector3:
        return Vector3(round(self.x, n), round(self.y, n), round(self.z, n))


class Transform:
    """A transform in 3D space that represents position, rotation, and scale.

    By convention, units for world space are in meters.

    Args:
        position (Vector3, optional): The position. Defaults to Vector3.zero().
        rotation (Vector3, optional): The rotation. Defaults to Vector3.zero().
        scale (Vector3, optional): The scale. Defaults to Vector3(1, 1, 1).
    """

    def __init__(
        self,
        position: Vector3 = Vector3.zero(),
        rotation: Vector3 = Vector3.zero(),
        scale: Vector3 = Vector3(1, 1, 1),
    ) -> None:
        self._position = self._get_translation_matrix(position)
        self._rotation = self._get_rotation_matrix(rotation)
        self._scale = self._get_scale_matrix(scale)

    @property
    def position(self) -> Vector3:
        return Vector3(*self._position[:3, 3])

    @property
    def rotation(self) -> Vector3:
        return self._matrix_to_euler(self._rotation)

    @property
    def scale(self) -> Vector3:
        return Vector3(self._scale[0, 0], self._scale[1, 1], self._scale[2, 2])

    @property
    def matrix(self) -> np.ndarray:
        return self._position @ self._rotation @ self._scale

    def _get_translation_matrix(self, position: Vector3) -> np.ndarray:
        """Returns a translation matrix for the given position.

        Args:
            position (Vector3): The position.

        Returns:
            np.ndarray: The translation matrix.
        """
        translation_matrix = np.identity(4)
        translation_matrix[:3, 3] = position.numpy()
        return translation_matrix

    def _get_rotation_matrix_x(self, angle_radians: float) -> np.ndarray:
        """Returns the 4x4 rotation matrix around the X-axis.

        Args:
            angle_radians (float): The angle in radians.
        """
        c = math.cos(angle_radians)
        s = math.sin(angle_radians)
        return np.array(
            [
                [1, 0, 0, 0],
                [0, c, -s, 0],
                [0, s, c, 0],
                [0, 0, 0, 1],
            ]
        )

    def _get_rotation_matrix_y(self, angle_radians: float) -> np.ndarray:
        """Returns the 4x4 rotation matrix around the Y-axis.

        Args:
            angle_radians (float): The angle in radians.
        """
        c = math.cos(angle_radians)
        s = math.sin(angle_radians)
        return np.array(
            [
                [c, 0, s, 0],
                [0, 1, 0, 0],
                [-s, 0, c, 0],
                [0, 0, 0, 1],
            ]
        )

    def _get_rotation_matrix_z(self, angle_radians: float) -> np.ndarray:
        """Returns the 4x4 rotation matrix around the Z-axis.

        Args:
            angle_radians (float): The angle in radians.
        """
        c = math.cos(angle_radians)
        s = math.sin(angle_radians)
        return np.array(
            [
                [c, -s, 0, 0],
                [s, c, 0, 0],
                [0, 0, 1, 0],
                [0, 0, 0, 1],
            ]
        )

    def _matrix_to_euler(self, mtx: np.ndarray) -> Vector3:
        """Returns the rotation in degrees for the given rotation matrix.
        x - bank / roll / tilt / psi ψ [0],
        y - attitude / pitch / elevation / phi φ [1],
        z - heading / yaw / azimuth / theta θ [2]

        Args:
            mtx (np.ndarray): The rotation matrix.
        Returns:
            Vector3: Roll, pitch, yaw.
        """
        # https://eecs.qmul.ac.uk/~gslabaugh/publications/euler.pdf
        if math.isclose(mtx[2, 0], 1):
            psi = math.atan2(-mtx[1, 2], mtx[1, 1])
            phi = -math.pi / 2
            theta = 0.0
        elif math.isclose(mtx[2, 0], -1):
            psi = -math.atan2(-mtx[1, 2], mtx[1, 1])
            phi = math.pi / 2
            theta = 0.0
        else:
            psi = math.atan2(mtx[2, 1], mtx[2, 2])
            phi = math.asin(-mtx[2, 0])
            theta = math.atan2(mtx[1, 0], mtx[0, 0])
        return Vector3(math.degrees(psi), math.degrees(phi), math.degrees(theta))

    def _get_rotation_matrix(self, rotation: Vector3) -> np.ndarray:
        """Returns the combined 4x4 rotation matrix for the given rotation.

        Args:
            rotation (Vector3): The rotation in degrees.

        Returns:
            np.ndarray: The rotation matrix.
        """
        rx = self._get_rotation_matrix_x(math.radians(rotation.x))
        ry = self._get_rotation_matrix_y(math.radians(rotation.y))
        rz = self._get_rotation_matrix_z(math.radians(rotation.z))
        return rz @ ry @ rx  # Combine rotation matrices in ZYX order

    def _get_scale_matrix(self, scale: Vector3) -> np.ndarray:
        """Returns a scale matrix for the given scale.

        Args:
            scale (Vector3): The scale.

        Returns:
            np.ndarray: The scale matrix.
        """
        scale_matrix = np.identity(4)
        scale_matrix[0, 0] = scale.x
        scale_matrix[1, 1] = scale.y
        scale_matrix[2, 2] = scale.z
        return scale_matrix

    def apply_relative_transform(self, relative_transform: Transform) -> Transform:
        """Applies the relative transform to this transform.

        Args:
            relative_transform (Transform): The relative transform.

        Returns:
            Transform: The transformed transform.
        """
        position = self.transform_position(relative_transform.position)
        rotation = self.transform_rotation(relative_transform.rotation)
        scale = self.transform_scale(relative_transform.scale)
        return Transform(position, rotation, scale)

    def transform_position(self, position: Vector3) -> Vector3:
        """Transforms a position vector by this transform accounting its scale.

        Args:
            position (Vector3): The position to transform.

        Returns:
            Vector3: The transformed position.
        """

        # Convert UE left-handed to Numpy right-handed coordinate system
        coord_conversion = np.array([[1, 1, -1], [1, 1, -1], [-1, -1, 1]])
        rot = self._rotation[:3, :3] * coord_conversion

        scaled = np.array(
            [
                position.x * self.scale.x,
                position.y * self.scale.y,
                position.z * self.scale.z,
            ]
        )
        r = np.dot(rot, scaled) + self.position.numpy()
        return Vector3(r[0], r[1], r[2])

    def transform_rotation(self, relative_rotation: Vector3) -> Vector3:
        """Transforms a rotation vector by this transform.

        Args:
            relative_rotation (Vector3): The rotation to transform.

        Returns:
            Vector3: The transformed rotation.
        """
        return self._matrix_to_euler(
            self._rotation @ self._get_rotation_matrix(relative_rotation)
        )

    def transform_scale(self, relative_scale: Vector3) -> Vector3:
        """Transforms a scale vector by this transform.

        Args:
            relative_scale (Vector3): The scale to transform.

        Returns:
            Vector3: The transformed scale.
        """
        return Vector3(
            self.scale.x * relative_scale.x,
            self.scale.y * relative_scale.y,
            self.scale.z * relative_scale.z,
        )

    def relative_to(self, reference: Transform) -> Transform:
        """Calculates the transform of this transform relative to another transform.

        Args:
            reference (Transform): The reference transform against which the relative
                                transformation is calculated.

        Returns:
            Transform: A new Transform object representing this transform relative to `reference`,
                        with calculated position, rotation, and scale.
        """
        position = reference.get_relative_position(self.position)
        rotation = reference.get_relative_rotation(self.rotation)
        scale = reference.get_relative_scale(self.scale)
        return Transform(position, rotation, scale)

    def get_relative_position(self, other_position: Vector3) -> Vector3:
        """Calculates the relative position of other position relative to this transform.

        Args:
            other_position (Vector3): The other position.

        Returns:
            Vector3: The relative position.
        """
        # Convert UE left-handed to Numpy right-handed coordinate system
        coord_conversion = np.array([[1, 1, -1], [1, 1, -1], [-1, -1, 1]])
        rot = self._rotation[:3, :3] * coord_conversion

        offset = other_position.numpy() - self.position.numpy()

        # apply inverse of rot to offset to rotate it into self's coordinate system
        rotated_position = np.dot(rot.T, offset)

        # divide by self's scale to get the relative position
        relative_position = rotated_position / self.scale.numpy()

        return Vector3(*relative_position)

    def get_relative_rotation(self, other_rotation: Vector3) -> Vector3:
        """Calculates the relative rotation of other rotation relative to this transform.

        Args:
            other_rotation (Vector3): The other rotation.

        Returns:
            Vector3: The relative rotation.
        """
        return self._matrix_to_euler(
            self._rotation.T @ self._get_rotation_matrix(other_rotation)
        )

    def get_relative_scale(self, other_scale: Vector3) -> Vector3:
        """Calculates the relative scale of other scale relative to this transform.

        Args:
            other_scale (Vector3): The other scale.

        Returns:
            Vector3: The relative scale.
        """
        return Vector3(
            other_scale.x / self.scale.x,
            other_scale.y / self.scale.y,
            other_scale.z / self.scale.z,
        )

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Transform):
            return NotImplemented
        return (
            self.position == other.position
            and self.rotation == other.rotation
            and self.scale == other.scale
        )

    def __repr__(self) -> str:
        return (
            f"Transform(position={self.position}, rotation={self.rotation}, "
            f"scale={self.scale})"
        )


def calculate_vfov(hfov: float, resolution: tuple[int, int]) -> float:
    """
    Calculate the vertical field of view in degrees.

    Args:
        hfov (float): The horizontal field of view in degrees.
        resolution (tuple[int, int]): The resolution of the camera.

    Returns:
        float: The vertical field of view in degrees.
    """
    # Relation between FOVs and resolution is
    # width / height = tan(hfov / 2) / tan(vfov / 2)
    cam_width, cam_height = resolution
    hfov = math.radians(hfov)
    return math.degrees(2 * math.atan(math.tan(hfov / 2) / cam_width * cam_height))


def ray_direction_to_screen(
    direction: Vector3,
    cam_rotation: Vector3,
    hfov: float,
    vfov: float,
) -> tuple[float, float] | None:
    """
    Find where a camera ray direction projects onto the screen.

    Given a camera ray direction in world space, calculates the corresponding screen space
    coordinates where that ray would intersect the screen.

    Screen space coordinates (0-1 range):
    - (0.5, 0.5) = center
    - (0, 1) = bottom-left corner
    - (0, 0) = top-left corner
    - (1, 0) = top-right corner
    - (1, 1) = bottom-right corner

    Args:
        direction (Vector3): Camera ray direction in world space. Should be a unit vector.
        cam_rotation (Vector3): Camera rotation (roll, pitch, yaw) in degrees.
        hfov (float): Horizontal field of view in degrees.
        vfov (float): Vertical field of view in degrees.

    Returns:
        tuple[float, float] | None: The screen space coordinates where the ray intersects the screen,
        or None if the ray is behind the camera or on the camera plane.
    """
    # transform world direction to camera space (inverse rotation)
    r_camera = R.from_euler(
        "xyz", [-cam_rotation.x, -cam_rotation.y, cam_rotation.z], degrees=True
    )
    r_camera_inv = r_camera.inv()

    direction_camera = r_camera_inv.apply(direction.numpy())
    x, y, z = direction_camera

    # direction is on the camera plane
    if math.isclose(x, 0):
        return None
    # direction is behind the camera
    if x < 0:
        return None

    # pinhole projection conversion from world space to screen space
    # screen_x = 0.5 + 0.5 * (y/x) / tan(hfov/2)
    # screen_y = 0.5 - 0.5 * (z/x) / tan(vfov/2)
    tan_half_hfov = math.tan(math.radians(hfov / 2))
    tan_half_vfov = math.tan(math.radians(vfov / 2))
    screen_x = 0.5 + 0.5 * (y / x) / tan_half_hfov
    screen_y = 0.5 - 0.5 * (z / x) / tan_half_vfov

    return (screen_x, screen_y)


def world_to_screen(
    world_coords: Vector3,
    cam_position: Vector3,
    cam_rotation: Vector3,
    cam_hfov: float,
    cam_vfov: float,
    img_res: tuple[int, int],
) -> tuple[int, int] | None:
    """
    Convert world coordinates to screen coordinates.
    If world coordinates are behind the camera or on the camera plane, return None.

    Args:
        world_coords (Vector3): World coordinates to convert.
        cam_position (Vector3): Camera position.
        cam_rotation (Vector3): Camera rotation.
        cam_hfov (float): Camera horizontal field of view.
        cam_vfov (float): Camera vertical field of view.
        img_res (tuple[int, int]): Image resolution (width, height).
    """
    direction_vector = (world_coords - cam_position).normalized()
    relative_coords = ray_direction_to_screen(
        direction_vector, cam_rotation, cam_hfov, cam_vfov
    )
    if relative_coords is None:
        return None
    return round(relative_coords[0] * img_res[0]), round(
        relative_coords[1] * img_res[1]
    )