"""Data containers for he11lib: measurement inputs and reconstruction outputs."""

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np


@dataclass
class MeasurementPlane:
    """A single thermal (flux) image at a nominal distance from the output window.

    Parameters
    ----------
    flux : 2D array of flux values (already dead-pixel/background/saturation
        corrected upstream).
    z : nominal distance from the output window, in meters. Must be positive.
    z_tolerance : +/- bound, in meters, around the nominal `z` within which
        the true distance is refined during fitting. Must be `>= 0`; `0`
        means `z` is trusted exactly and held fixed.
    label : optional human-readable label (e.g. "plane_40cm").
    """

    flux: np.ndarray
    z: float
    z_tolerance: float = 0.0
    label: str | None = None

    def __post_init__(self) -> None:
        if self.flux.ndim != 2:
            raise ValueError(
                f"MeasurementPlane.flux must be a 2D array, got shape {self.flux.shape}"
            )
        if self.z <= 0:
            raise ValueError(f"MeasurementPlane.z must be positive, got {self.z}")
        if self.z_tolerance < 0:
            raise ValueError(
                f"MeasurementPlane.z_tolerance must be >= 0, got {self.z_tolerance}"
            )


def validate_planes(planes: list[MeasurementPlane]) -> None:
    """Validate a list of MeasurementPlanes for use in reconstruction.

    Raises ValueError if there are fewer than 3 planes, shapes mismatch
    across planes, or z distances are not distinct.
    """
    if len(planes) < 3:
        raise ValueError(
            f"At least 3 measurement planes are required, got {len(planes)}"
        )

    shapes = {p.flux.shape for p in planes}
    if len(shapes) > 1:
        raise ValueError(f"All MeasurementPlanes must have the same shape, got {shapes}")

    z_values = [p.z for p in planes]
    if len(set(z_values)) != len(z_values):
        raise ValueError(f"MeasurementPlane z distances must be distinct, got {z_values}")


@dataclass
class ReconstructionResult:
    """Output of a full mode-purity reconstruction.

    Parameters
    ----------
    purity : mapping from (p, l) mode index to (power_fraction, phase_rad).
    reconstructed_field : reconstructed complex field (at the reference
        waist, or as configured).
    centers : fitted beam transverse center (x, y) in meters, per plane.
    pointing_angle_horizontal_deg, pointing_angle_vertical_deg : fitted
        shared beam pointing (tilt) angles, in degrees.
    geometry : fitted/held geometry parameters, keyed by name (the 9
        `CameraModel` field names from `he11lib.geometry.CAMERA_FIELD_NAMES`,
        plus `z_{i}` per plane index `i`).
    residuals : per-plane residual maps (measured - modeled flux).
    coefficient_uncertainty : mapping from (p, l) mode index to the
        1-sigma uncertainty on its fitted power fraction.
    used_phase_retrieval : whether the phase-retrieval fallback was used
        instead of (or to seed) the modal fit.
    """

    purity: dict[tuple[int, int], tuple[float, float]]
    reconstructed_field: np.ndarray
    centers: list[tuple[float, float]]
    pointing_angle_horizontal_deg: float
    pointing_angle_vertical_deg: float
    geometry: dict[str, float] = field(default_factory=dict)
    residuals: list[np.ndarray] = field(default_factory=list)
    coefficient_uncertainty: dict[tuple[int, int], float] = field(default_factory=dict)
    used_phase_retrieval: bool = False