"""Co-located equivalent pupil model, with explicit finite-FOV energy accounting.

This is intentionally NOT a thick posterior-chamber ICL model. The central
port replaces the local ICL correction with aqueous transmission; it is not
an opaque disk. All distances are SI internally. Positive equivalent power
produces OPD -D*r**2/2. Surface reflections, scatter and vault are absent.
"""

from dataclasses import asdict, dataclass
import math

import numpy as np

ARCSEC_PER_RAD = 180.0 * 3600.0 / math.pi


@dataclass(frozen=True)
class Config:
    pupil_mm: float = 4.0
    eye_power_d: float = 8.0
    icl_power_d: float = -8.0
    optic_mm: float = 6.0
    port_mm: float = 0.0
    decenter_x_mm: float = 0.0
    decenter_y_mm: float = 0.0
    cylinder_d: float = 0.0
    cylinder_axis_deg: float = 0.0
    wavelengths_nm: tuple = (550.0,)
    spectral_weights: tuple = (1.0,)
    material_amplitude: float = 1.0
    grid_n: int = 512
    padding: int = 2
    extent_mm: float = 10.24
    output_n: int = 512
    angular_pixel_arcsec: float = 8.0

    def validate(self):
        raw = asdict(self)
        for key, value in raw.items():
            values = value if isinstance(value, (list, tuple)) else [value]
            if not all(isinstance(v, (int, float)) and math.isfinite(v) for v in values):
                raise ValueError(f"{key}: finite numeric values required")
        for key in ("grid_n", "padding", "output_n"):
            if type(getattr(self, key)) is not int:
                raise ValueError(f"{key}: integer required")
        if not (32 <= self.grid_n <= 2048 and 1 <= self.padding <= 8 and 32 <= self.output_n <= 2048):
            raise ValueError("unsupported grid, padding or output size")
        if self.grid_n * self.padding > 4096:
            raise ValueError("FFT resource limit exceeded")
        if not (0 < self.pupil_mm < self.extent_mm and 0 <= self.port_mm < self.optic_mm):
            raise ValueError("invalid pupil/port/optic diameter")
        if self.optic_mm <= 0 or self.angular_pixel_arcsec <= 0:
            raise ValueError("positive optic and angular sampling required")
        if not 0 <= self.material_amplitude <= 1:
            raise ValueError("passive amplitude must be in [0,1]")
        if not self.wavelengths_nm or len(self.wavelengths_nm) != len(self.spectral_weights):
            raise ValueError("one spectral weight per wavelength required")
        if min(self.wavelengths_nm) < 380 or max(self.wavelengths_nm) > 780:
            raise ValueError("this core supports 380–780 nm only")
        if min(self.spectral_weights) < 0 or sum(self.spectral_weights) <= 0:
            raise ValueError("nonnegative spectrum with positive total required")

    @classmethod
    def from_dict(cls, data):
        return cls(**data)

    def to_dict(self):
        return asdict(self)


def pupil_opd(config):
    """Return pupil amplitude, OPD [m], and incident sampled pupil area count."""
    config.validate()
    dx = config.extent_mm * 1e-3 / config.grid_n
    coord = (np.arange(config.grid_n) - config.grid_n // 2) * dx
    x, y = np.meshgrid(coord, coord)
    pupil = x*x + y*y <= (config.pupil_mm * 0.5e-3)**2
    ix = x - config.decenter_x_mm * 1e-3
    iy = y - config.decenter_y_mm * 1e-3
    local_r2 = ix*ix + iy*iy
    material = local_r2 <= (config.optic_mm * 0.5e-3)**2
    if config.port_mm > 0:
        material &= local_r2 >= (config.port_mm * 0.5e-3)**2
    # Base power is equivalent residual myopia in the chosen reference plane.
    # Outside ICL material, including the transparent port, only this remains.
    opd = -0.5 * config.eye_power_d * (x*x + y*y)
    opd += np.where(material, -0.5 * config.icl_power_d * local_r2, 0.0)
    axis = math.radians(config.cylinder_axis_deg)
    meridian = -x * math.sin(axis) + y * math.cos(axis)
    opd -= 0.5 * config.cylinder_d * meridian**2
    # Sample the smooth phase gradients inside each material region. Edges
    # still require separate convergence checks; this catches bulk aliasing.
    gx = -config.eye_power_d*x - material*config.icl_power_d*ix + config.cylinder_d*meridian*math.sin(axis)
    gy = -config.eye_power_d*y - material*config.icl_power_d*iy - config.cylinder_d*meridian*math.cos(axis)
    phase_cycles_per_pixel = max(float(np.max(np.abs(gx[pupil]))),float(np.max(np.abs(gy[pupil]))))*dx/(min(config.wavelengths_nm)*1e-9)
    if phase_cycles_per_pixel >= .45:
        raise ValueError("bulk pupil phase is undersampled: increase grid_n, not FFT padding")
    amplitude = pupil * np.where(material, config.material_amplitude, 1.0)
    return amplitude, opd, int(pupil.sum())


def _resample_square(data, source_axis, target_axis):
    """Conservative pixel-energy remap via separable bin-overlap integration.

    Each source bin has constant angular energy density. Interpolating its
    cumulative integral at target edges preserves energy and handles finite
    FOV loss without wrapping or an artificial normalization correction.
    """
    source_step = source_axis[1]-source_axis[0]
    target_step = target_axis[1]-target_axis[0]
    source_edges = np.r_[source_axis-source_step/2, source_axis[-1]+source_step/2]
    target_edges = np.r_[target_axis-target_step/2, target_axis[-1]+target_step/2]
    def remap(row):
        cumulative = np.r_[0., np.cumsum(row)]
        return np.diff(np.interp(target_edges, source_edges, cumulative, left=0., right=cumulative[-1]))
    rows = np.array([remap(row) for row in data])
    return np.array([remap(col) for col in rows.T]).T


def simulate(config):
    """PSF is normalized conditional on the documented detector field of view.

    throughput: integrated pupil output / incident pupil energy.
    captured_energy: energy on the common detector grid / incident pupil energy.
    Conservative resampling preserves total energy; finite FOV crops it.
    """
    amplitude, opd, incident = pupil_opd(config)
    size = config.grid_n * config.padding
    center = (size - config.grid_n) // 2
    dx = config.extent_mm * 1e-3 / config.grid_n
    target = (np.arange(config.output_n) - config.output_n // 2) * config.angular_pixel_arcsec
    energy = np.zeros((config.output_n, config.output_n), dtype=np.float64)
    weights = np.asarray(config.spectral_weights, dtype=float)
    weights /= weights.sum()
    for wavelength_nm, weight in zip(config.wavelengths_nm, weights):
        wavelength = wavelength_nm * 1e-9
        field = amplitude * np.exp(2j * np.pi * opd / wavelength)
        padded = np.zeros((size, size), dtype=complex)
        padded[center:center+config.grid_n, center:center+config.grid_n] = field
        image = np.abs(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(padded))))**2
        # Parseval: sum(image)/(size**2 * incident) == pupil throughput.
        image /= size**2 * incident
        source_step = wavelength / (size * dx) * ARCSEC_PER_RAD
        source_axis = (np.arange(size) - size // 2) * source_step
        mapped = _resample_square(image, source_axis, target)
        energy += weight * mapped
    captured = float(energy.sum())
    if captured <= 0 or not math.isfinite(captured):
        raise ValueError("no finite energy captured")
    return {
        "psf": energy / captured,
        "throughput": float(np.sum(amplitude**2) / incident),
        "captured_energy": captured,
        "angular_pixel_arcsec": config.angular_pixel_arcsec,
    }


def metrics(psf, pixel_arcsec):
    """All metrics are calculated from raw linear arrays, never display images."""
    n = psf.shape[0]
    axis = (np.arange(n) - n // 2) * pixel_arcsec
    x, y = np.meshgrid(axis, axis)
    cx, cy = float(np.sum(psf*x)), float(np.sum(psf*y))
    radius = np.sqrt(x*x + y*y)  # optical axis fixed, do not align away prism
    radii = np.array([16, 32, 64, 128, 256, 512, 1024], dtype=float)
    ee = np.array([psf[radius <= r].sum() for r in radii])
    otf = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(psf)))
    mtf = np.abs(otf) / abs(otf[n//2, n//2])
    return {"centroid_arcsec": [cx, cy], "ee_radii_arcsec": radii.tolist(),
            "encircled_energy": ee.tolist(),
            "mtf_x": mtf[n//2, n//2:n//2+32].tolist(),
            "mtf_y": mtf[n//2:n//2+32, n//2].tolist()}