import math
import numpy as np
from typing import Optional
from allytools.logger import get_logger
from allytools.units import Length
from dataclasses import dataclass

log = get_logger(__name__)
@dataclass
class PupilGrid:
    """
    Pupil grid in normalized coordinates.

    - The grid is defined in normalized pupil coordinates:
        x_norm, y_norm ∈ [-1, 1]
        rho = sqrt(x_norm^2 + y_norm^2) ≤ 1

    - For on-axis field:
        - Circular pupil: rho <= 1

    - For off-axis field (if field_x_img_mm, field_y_img_mm, distance_to_exit_pupil_mm are given):
        - Elliptical pupil oriented towards the field point.
        - Radii in normalized coordinates:
            R_sag = 1
            R_tan = 1 / cos(theta)
          where theta is the chief ray angle w.r.t. the optical axis.
    """
    grid_size: int # Grid resolution in each dimension
    field_x_img:  Optional[Length] = None
    field_y_img: Optional[Length] = None
    exit_pupil_position: Optional[Length] = None

    def build(self):
        """
        Build normalized pupil coordinates and pupil mask.

        Returns
        -------
        x_norm : np.ndarray
            Normalized x coordinate ([-1, 1]) on a square grid.
        y_norm : np.ndarray
            Normalized y coordinate ([-1, 1]) on a square grid.
        rho : np.ndarray
            Normalized radial coordinate, rho = sqrt(x_norm^2 + y_norm^2).
        phi : np.ndarray
            Azimuthal angle, phi = atan2(y_norm, x_norm).
        mask : np.ndarray
            Pupil mask (float), 1.0 inside the pupil, 0.0 outside.
            - On-axis: circular pupil (rho <= 1).
            - Off-axis: elliptical pupil in (t, s) coordinates.
        """

        # 1️⃣ Base normalized grid [-1, 1] x [-1, 1]
        lin = np.linspace(-1.0, 1.0, self.grid_size)
        x_norm, y_norm = np.meshgrid(lin, lin)
        log.debug(f"Created normalized grid {self.grid_size}x{self.grid_size}.")

        # 🔧 Align pupil coordinate system with Zemax:
        # apply a 180° rotation in the pupil plane
        x_norm = -x_norm
        y_norm = -y_norm

        # 2️⃣ Normalized polar coordinates
        rho = np.sqrt(x_norm**2 + y_norm**2)
        phi = np.arctan2(y_norm, x_norm)
        log.debug("Computed polar coordinates (rho, phi).")

        # 3️⃣ On-axis circular pupil
        if self.field_x_img is None or self.field_y_img is None or self.exit_pupil_position is None:
            log.debug("No off-axis parameters → circular pupil.")
            mask = (rho <= 1.0).astype(float)
            return x_norm, y_norm, rho, phi, mask


        # 4️⃣ Off-axis elliptical pupil
        # Tangential axis - direction toward the field point
        # Sagittal   axis - perpendicular to tangential
        fx_mm = self.field_x_img.value_mm
        fy_mm = self.field_y_img.value_mm
        l_mm = self.exit_pupil_position.value_mm
        log.debug(f"Off-axis parameters: fx={fx_mm:.4f} mm, fy={fy_mm:.4f} mm, L={l_mm:.4f} mm.")

        # 5️⃣ Chief ray angles in x and y (small-angle approximation)
        theta_x = fx_mm / l_mm
        theta_y = fy_mm / l_mm
        theta = math.sqrt(theta_x**2 + theta_y**2)
        log.debug(f"Chief-ray angles: theta_x={theta_x:.6f} theta_y={theta_y:.6f}, |theta|={theta:.6f} rad.")

        # 6️⃣If the field is essentially on-axis -> degrade to circular pupil
        if theta < 1e-9:
            log.debug("Field very small → reverting to circular pupil.")
            mask = (rho <= 1.0).astype(float)
            return x_norm, y_norm, rho, phi, mask

        # 7️⃣ Tangential direction in the image plane (and pupil):
        # Use the vector from axis to field point (fx, fy) as tangent reference.
        v_t = np.array([fx_mm, fy_mm], dtype=float)
        v_t_norm = np.linalg.norm(v_t)
        u_t = v_t / v_t_norm  # unit vector in tangential direction
        log.debug(f"Tangential direction u_t = {u_t}.")

        # 8️⃣ Sagittal direction is perpendicular to tangential:
        # rotate u_t by +90 degrees: (x, y) -> (-y, x)
        u_s = np.array([-u_t[1], u_t[0]])
        log.debug(f"Sagittal direction u_s = {u_s}.")

        # 9️⃣ Coordinates (t, s) in the (u_t, u_s) basis for each pupil point:
        # (x_norm, y_norm) -> t = x*u_t.x + y*u_t.y; s = x*u_s.x + y*u_s.y
        t = x_norm * u_t[0] + y_norm * u_t[1]
        s = x_norm * u_s[0] + y_norm * u_s[1]
        log.debug("Computed tangential/sagittal coordinates (t, s).")

        # 1️⃣0️⃣ Ellipse radii in normalized coordinates:
        # sagittal radius stays 1, tangential radius is stretched by 1/cos(theta)
        r_sag = 1.0
        cos_theta = math.cos(theta)
        cos_theta = max(cos_theta, 1e-6) # Numerical safety
        r_tan = 1.0 / cos_theta
        log.debug(f"Ellipse radii: R_tan={r_tan:.6f}, R_sag={r_sag:.6f} (cos(theta)={cos_theta:.6f})")
        ellipse = (t / r_tan) ** 2 + (s / r_sag) ** 2  # Elliptical pupil mask in (t, s) coordinates
        mask = (ellipse <= 1.0).astype(float)
        inside = np.sum(mask)
        log.debug(f"Elliptical pupil built. Mask inside count = {inside}/{mask.size}.")
        return x_norm, y_norm, rho, phi, mask

