import numpy as np
from pydantic import BaseModel


class Coordinate(BaseModel):
    x: float = 0.0
    y: float = 0.0
    z: float = 0.0

    # Overload +
    def __add__(self, other: "Coordinate") -> "Coordinate":
        if isinstance(other, Coordinate):
            return Coordinate(x=self.x + other.x, y=self.y + other.y, z=self.z + other.z)
        return NotImplemented

    # Overload -
    def __sub__(self, other: "Coordinate") -> "Coordinate":
        if isinstance(other, Coordinate):
            return Coordinate(x=self.x - other.x, y=self.y - other.y, z=self.z - other.z)
        return NotImplemented

    # Overload *
    def __mul__(self, other):
        if isinstance(other, (int, float)):  # Scalar multiplication
            return Coordinate(x=self.x * other, y=self.y * other, z=self.z * other)
        elif isinstance(other, Coordinate):  # Dot product
            return self.x * other.x + self.y * other.y + self.z * other.z
        return NotImplemented

    # Overload /
    def __truediv__(self, scalar: float) -> "Coordinate":
        if isinstance(scalar, (int, float)) and scalar != 0:  # Avoid division by zero
            return Coordinate(x=self.x / scalar, y=self.y / scalar, z=self.z / scalar)
        elif scalar == 0:
            raise ValueError("Cannot divide by zero")
        return NotImplemented

    def normalize(self) -> "Coordinate":
        magnitude = np.sqrt(self.x**2 + self.y**2 + self.z**2)
        if magnitude == 0:
            raise ValueError("Cannot normalize a zero-magnitude vector.")
        return self / magnitude

    def rotate(self, angle_deg: float, axis: str) -> "Coordinate":
        """
        Rotate the coordinate around a specified axis by a given angle in degrees.

        :param angle_deg: The angle of rotation in degrees.
        :param axis: The axis to rotate around ('x', 'y', or 'z').
        :return: A new Coordinate after rotation.
        """
        angle_rad = np.radians(angle_deg)  # Convert angle to radians
        c = np.cos(angle_rad)
        s = np.sin(angle_rad)

        if axis == "x":
            # Rotate around x-axis (affects y, z)
            y_new = self.y * c - self.z * s
            z_new = self.y * s + self.z * c
            return Coordinate(x=self.x, y=y_new, z=z_new)
        elif axis == "y":
            # Rotate around y-axis (affects x, z)
            x_new = self.x * c + self.z * s
            z_new = -self.x * s + self.z * c
            return Coordinate(x=x_new, y=self.y, z=z_new)
        elif axis == "z":
            # Rotate around z-axis (affects x, y)
            x_new = self.x * c - self.y * s
            y_new = self.x * s + self.y * c
            return Coordinate(x=x_new, y=y_new, z=self.z)
        else:
            raise ValueError("Invalid axis. Choose 'x', 'y', or 'z'.")


class SmargonCoordinate(BaseModel):
    sh_mm: Coordinate | None = None  # SH coordinate in Smargon
    phi_deg: float | None = None
    chi_deg: float | None = None

    def eq(self, other: "SmargonCoordinate", tol: float) -> bool:
        return (
            abs(self.sh_mm.x - other.sh_mm.x) < tol
            and abs(self.sh_mm.y - other.sh_mm.y) < tol
            and abs(self.sh_mm.z - other.sh_mm.z) < tol
            and abs(self.phi_deg - other.phi_deg) < tol
            and abs(self.chi_deg - other.chi_deg) < tol
        )

    def __eq__(self, other: object) -> bool:
        if isinstance(other, SmargonCoordinate):
            return self.eq(other, tol=0.1)
        return NotImplemented


def positive_coords(value: Coordinate) -> Coordinate:
    if value.x <= 0 or value.y <= 0:
        raise ValueError("Coordinates must be positive")
    return value


class AerotechCoordinate(BaseModel):
    at_mm: Coordinate | None = None
    omega_deg: float | None = None

    def eq(self, other: "AerotechCoordinate", tol: float) -> bool:
        return (
            abs(self.at_mm.x - other.at_mm.x) < tol
            and abs(self.at_mm.y - other.at_mm.y) < tol
            and abs(self.at_mm.z - other.at_mm.z) < tol
            and abs(self.omega_deg - other.omega_deg) < tol
        )

    def __eq__(self, other: object) -> bool:
        if isinstance(other, AerotechCoordinate):
            return self.eq(other, tol=0.01)
        return NotImplemented