# ============================
# camera.py — Orbit Camera System
# ============================
from __future__ import annotations
from dataclasses import dataclass, field
import math
import numpy as np
from typing import Optional

def normalize(v: np.ndarray) -> np.ndarray:
    """Returns the normalized vector v."""
    n = float(np.linalg.norm(v))
    return v / (n if n > 1e-12 else 1.0)

def rodrigues(v: np.ndarray, axis: np.ndarray, angle: float) -> np.ndarray:
    """Rotates vector v around axis by angle (radians) using Rodrigues' formula."""
    mag = np.linalg.norm(axis)
    if mag < 1e-10:
        return v
    k = axis / mag
    c = math.cos(angle)
    s = math.sin(angle)
    return v * c + np.cross(k, v) * s + k * np.dot(k, v) * (1.0 - c)

WORLD_UP = np.array([0, 0, 1], dtype=np.float32)

@dataclass
class OrbitCamera:
    """
    Standard Orbit Camera for molecular visualization.
    Handles yaw, pitch, panning, and isometric snapping.
    """
    center: np.ndarray
    offset: np.ndarray # Vector from center to eye
    fov: float = 45.0
    ortho_scale: float = 20.0
    near: float = 0.1
    far: float = 5000.0
    
    # Camera Mode (ORBIT or FLY)
    mode: str = "ORBIT"
    fly_speed: float = 10.0
    
    # Physics/Momentum
    yaw_v: float = 0.0
    pitch_v: float = 0.0
    damping: float = 0.92
    sensitivity: float = 0.005
    momentum: bool = True
    
    # Camera Shift (Pan)
    pan: np.ndarray = field(default_factory=lambda: np.zeros(2, dtype=np.float32))
    invert_x: bool = False
    invert_y: bool = False
    clamped_pitch: bool = True
    up: np.ndarray = field(default_factory=lambda: np.array([0, 0, 1], dtype=np.float32))

    def move_local(self, dx: float, dy: float, dz: float):
        """Moves camera in its own local coordinate system (right, screen-up, forward)."""
        fwd = self.forward()
        right = self.right()
        # Screenspace Up (orthogonal to forward and right)
        up = np.cross(right, fwd)
        
        delta = right * dx + up * dy + fwd * dz
        self.center += delta

    def update(self, dt: float):
        """Processes momentum and applies damping to camera velocities."""
        if not self.momentum:
            self.yaw_v = self.pitch_v = 0.0
            return

        # Apply velocities
        if abs(self.yaw_v) > 1e-5:
            self.rotate_yaw(self.yaw_v * dt * 60.0)
        if abs(self.pitch_v) > 1e-5:
            self.rotate_pitch(self.pitch_v * dt * 60.0)

        # Decay
        d = math.pow(self.damping, dt * 60.0)
        self.yaw_v *= d
        self.pitch_v *= d
        
        if abs(self.yaw_v) < 1e-5: self.yaw_v = 0.0
        if abs(self.pitch_v) < 1e-5: self.pitch_v = 0.0

    def get_pan_vector(self) -> np.ndarray:
        # Convert 2D pan (X, Y) to 3D vector in camera basis
        right = self.right()
        # Up vector relative to camera view (Screen Up)
        up = np.cross(right, self.forward())
        return right * self.pan[0] + up * self.pan[1]

    def eye(self) -> np.ndarray:
        # Eye is center + offset + pan (in camera plane)
        return self.center + self.offset + self.get_pan_vector()

    def target(self) -> np.ndarray:
        # Target is center + pan
        return self.center + self.get_pan_vector()

    def forward(self) -> np.ndarray:
        return normalize(-self.offset)

    def right(self) -> np.ndarray:
        """Returns the world-space 'right' vector (perpendicular to forward and up)."""
        return normalize(np.cross(self.forward(), self.up))

    def rotate_yaw(self, angle: float):
        if self.invert_x: angle = -angle
        if self.mode == "ORBIT":
            self.offset = rodrigues(self.offset, self.up, angle)
        else:
            # FLY mode: pivot around eye
            eye_fixed = self.eye()
            self.offset = rodrigues(self.offset, self.up, -angle)
            self.center = eye_fixed - (self.offset + self.get_pan_vector())

    def rotate_pitch(self, angle: float):
        if self.invert_y: angle = -angle
        
        forward = normalize(-self.offset)
        right = self.right()
        
        if np.linalg.norm(right) < 1e-10:
            # Fallback if at a pole: use a global axis
            right = np.array([1, 0, 0], dtype=np.float32)

        if self.clamped_pitch:
            # Traditional Turntable Clamping
            up_dot = np.dot(forward, self.up)
            max_elev = math.radians(89.0)
            cur_elev = math.asin(np.clip(up_dot, -1.0, 1.0))
            
            new_elev = cur_elev + angle
            if new_elev > max_elev:
                angle = max_elev - cur_elev
            elif new_elev < -max_elev:
                angle = -max_elev - cur_elev
            
        if self.mode == "ORBIT":
            self.offset = rodrigues(self.offset, right, angle)
            # If unclamped, rotate the UP vector too to enable smooth "Trackball" movement
            if not self.clamped_pitch:
                self.up = normalize(rodrigues(self.up, right, angle))
        else:
            # FLY mode: pivot around eye
            eye_fixed = self.eye()
            self.offset = rodrigues(self.offset, right, -angle)
            self.center = eye_fixed - (self.offset + self.get_pan_vector())
            if not self.clamped_pitch:
                self.up = normalize(rodrigues(self.up, right, -angle))

    def rotate_roll(self, angle: float):
        """Rotates the camera's UP vector around the view direction."""
        # Rotate UP around forward axis
        self.up = rodrigues(self.up, self.forward(), angle)
        # To keep right() consistent, we might also want to rotate offset? 
        # No, roll only affects the orientation of the 'top' of the screen.

    def get_view_matrix(self) -> np.ndarray:
        from .renderer.utils import look_at
        return look_at(self.eye(), self.target(), self.up)
        
    def get_projection_matrix(self, aspect: float, is_ortho: bool = False) -> np.ndarray:
        from .renderer.utils import perspective, ortho
        if is_ortho:
            s = self.ortho_scale / 2.0
            return ortho(-s * aspect, s * aspect, -s, s, self.near, self.far)
        else:
            return perspective(self.fov, aspect, self.near, self.far)

    def snap_isometric(self):
        """Snaps camera to a perfect isometric view (45 deg yaw, 35.264 deg pitch)."""
        dist = np.linalg.norm(self.offset)
        iso_pitch = math.radians(35.264)
        c, s = math.cos(iso_pitch), math.sin(iso_pitch)
        v = np.array([0, -dist*c, dist*s], dtype=np.float32)
        rot_z = math.radians(-45.0)
        v = rodrigues(v, np.array([0, 0, 1], dtype=np.float32), rot_z)
        self.offset = v
        self.up = np.array([0, 0, 1], dtype=np.float32)
