"""Spectrum collection and scale diagnostics."""

from __future__ import annotations

from collections.abc import Iterator, Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, cast

import jax.numpy as jnp
from astropy import units as u

from unite._utils import (
    C_KMS,
    _ensure_flux,
    _ensure_flux_density,
    _ensure_velocity,
    _get_conversion_factor,
)
from unite.instrument.base import Disperser
from unite.spectrum.spectrum import Spectrum

if TYPE_CHECKING:
    from unite.continuum.config import ContinuumConfiguration
    from unite.line.config import LineConfiguration


# ---------------------------------------------------------------------------
# Scale diagnostics dataclasses
# ---------------------------------------------------------------------------


@dataclass
class RegionDiagnostic:
    """Diagnostics for a single continuum region fit.

    Attributes
    ----------
    obs_low : float
        Observed-frame lower bound of the region (disperser's unit).
    obs_high : float
        Observed-frame upper bound of the region (disperser's unit).
    in_region : jnp.ndarray
        Boolean mask of shape ``(npix,)`` selecting all pixels inside the
        region bounds.
    good_mask : jnp.ndarray
        Boolean mask of shape ``(npix,)`` selecting pixels used for
        fitting (``in_region & ~line_mask``).
    model_on_region : jnp.ndarray or None
        Best-fit continuum model evaluated at the ``in_region`` pixels.
        ``None`` if the fit failed (too few unmasked pixels).
    chi2_red : float or None
        Reduced chi-squared of the fit, or ``None`` if the fit failed.
    fit_params : dict of str to float
        Best-fit parameter dict returned by :func:`~unite.continuum.fit.fit_continuum_form`.
        Empty if the fit failed.
    """

    obs_low: float
    obs_high: float
    in_region: jnp.ndarray
    good_mask: jnp.ndarray
    model_on_region: jnp.ndarray | None
    chi2_red: float | None
    fit_params: dict = field(default_factory=dict)


class ScaleDiagnosticList(list):
    """A list of :class:`SpectrumScaleDiagnostic` objects that also supports lookup by spectrum name.

    Inherits all standard :class:`list` behaviour.  Additionally,
    ``diags["my_spectrum"]`` returns the first diagnostic whose
    :attr:`~SpectrumScaleDiagnostic.name` matches the given string,
    raising :exc:`KeyError` if no match is found.
    """

    def __getitem__(self, idx):
        if isinstance(idx, str):
            for diag in self:
                if diag.name == idx:
                    return diag
            msg = f'No scale diagnostic found for spectrum name {idx!r}.'
            raise KeyError(msg)
        result = super().__getitem__(idx)
        if isinstance(idx, slice):
            result = ScaleDiagnosticList(result)
        return result


@dataclass
class SpectrumScaleDiagnostic:
    """Diagnostics for one spectrum produced by :meth:`Spectra.compute_scales`.

    Stored directly on the spectrum via
    :attr:`~unite.spectrum.Spectrum.scale_diagnostic` so
    callers always have the spectrum's own ``wavelength``, ``flux``, and
    ``error`` to hand — those fields are not duplicated here.

    Attributes
    ----------
    name : str
        Spectrum name (from :attr:`~unite.spectrum.Spectrum.name`).
    line_mask : jnp.ndarray
        Boolean mask of shape ``(npix,)``; ``True`` where a pixel was
        excluded because it lies near an emission line.
    continuum_model : jnp.ndarray
        Full-spectrum continuum model array of shape ``(npix,)``.
        Pixels not covered by any continuum region are ``NaN``.
    regions : list of RegionDiagnostic
        Per-region fit diagnostics (one entry per region that overlaps
        this spectrum).
    """

    name: str
    line_mask: jnp.ndarray
    continuum_model: jnp.ndarray
    regions: list


# ---------------------------------------------------------------------------
# Spectra collection
# ---------------------------------------------------------------------------


class Spectra:
    """Collection of spectra with coverage filtering.

    Wraps one or more :class:`~unite.spectrum.Spectrum` objects
    together with a systemic redshift estimate.  The main role of this class is
    :meth:`filter_config`, which drops lines and continuum regions not covered
    by any spectrum.

    Parameters
    ----------
    spectra : sequence of Spectrum
        Individual spectrum objects.  Must not be empty.
    redshift : float
        Systemic redshift estimate used for rest-frame → observed-frame
        conversion during coverage checks.  Default ``0.0``.
    canonical_unit : astropy.units.UnitBase, optional
        Wavelength unit used for all internal model computations (line
        centers, continuum bounds, pixel edges).  Defaults to the first
        spectrum's disperser unit.  Override this if you want a specific
        unit regardless of spectrum order.

    Raises
    ------
    ValueError
        If *spectra* is empty or contains non-Spectrum objects.
    """

    def __init__(
        self,
        spectra: Sequence[Spectrum],
        redshift: float = 0.0,
        canonical_unit: u.UnitBase | None = None,
    ) -> None:
        if not spectra:
            msg = 'Spectra collection must contain at least one spectrum.'
            raise ValueError(msg)
        for i, s in enumerate(spectra):
            if not isinstance(s, Spectrum):
                msg = f'All elements must be Spectrum instances; element {i} is {type(s).__name__}.'
                raise TypeError(msg)
        self._spectra: list[Spectrum] = list(spectra)
        self._redshift: float = float(redshift)
        self._line_scale: u.Quantity | None = None
        self._continuum_scale: u.Quantity | None = None
        self._is_prepared: bool = False
        self._prepared_line_config: LineConfiguration | None = None
        self._prepared_cont_config: ContinuumConfiguration | None = None
        self._scale_diagnostics: ScaleDiagnosticList | None = None

        # Assign names to any anonymous calibration tokens on the dispersers.
        # Local import avoids circular dependency at module load time.
        from unite.instrument.config import InstrumentConfig

        seen_dispersers: dict[int, Disperser] = {}
        for s in self._spectra:
            if id(s.disperser) not in seen_dispersers:
                seen_dispersers[id(s.disperser)] = s.disperser
        InstrumentConfig(list(seen_dispersers.values()))

        # Canonical wavelength unit: default to the first spectrum's unit.
        if canonical_unit is not None:
            _cu = cast(u.UnitBase, u.Unit(str(canonical_unit)))
            if not _cu.is_equivalent(u.m):
                msg = f'canonical_unit must have wavelength dimensions, got {_cu!r}.'
                raise u.UnitConversionError(msg)
            self._canonical_unit: u.UnitBase = _cu
        else:
            self._canonical_unit = self._spectra[0].unit

    # -- properties -----------------------------------------------------------

    @property
    def redshift(self) -> float:
        """Systemic redshift estimate."""
        return self._redshift

    @property
    def canonical_unit(self) -> u.UnitBase:
        """Canonical wavelength unit used for all internal model computations.

        Defaults to the first spectrum's wavelength unit.  Can be overridden
        via the *canonical_unit* constructor parameter.
        """
        return self._canonical_unit

    @property
    def line_scale(self) -> u.Quantity | None:
        """Characteristic line flux scale (integrated flux), or ``None``."""
        return self._line_scale

    @line_scale.setter
    def line_scale(self, value: u.Quantity) -> None:
        value = _ensure_flux(value, 'line_scale', ndim=0)
        if value.value <= 0:
            msg = f'line_scale must be > 0, got {value}'
            raise ValueError(msg)
        self._line_scale = value

    @property
    def continuum_scale(self) -> u.Quantity | None:
        """Characteristic continuum flux density scale, or ``None``."""
        return self._continuum_scale

    @continuum_scale.setter
    def continuum_scale(self, value: u.Quantity) -> None:
        value = _ensure_flux_density(value, 'continuum_scale', ndim=0)
        if value.value <= 0:
            msg = f'continuum_scale must be > 0, got {value}'
            raise ValueError(msg)
        self._continuum_scale = value

    @property
    def scale_diagnostics(self) -> ScaleDiagnosticList | None:
        """Per-spectrum diagnostics from the most recent :meth:`compute_scales` call.

        Returns a :class:`ScaleDiagnosticList` of
        :class:`SpectrumScaleDiagnostic` objects (one per spectrum).
        Supports both integer indexing and lookup by spectrum name
        (e.g. ``diags["my_spectrum"]``).  ``None`` if
        :meth:`compute_scales` has not been called yet.

        The diagnostic for an individual spectrum is also available via
        :attr:`~unite.spectrum.Spectrum.scale_diagnostic`,
        which is the preferred access pattern when iterating over spectra::

            for s in spectra:
                diag = s.scale_diagnostic
                # s.wavelength / s.flux / s.error available directly
        """
        return self._scale_diagnostics

    def compute_scales(
        self,
        line_config: LineConfiguration,
        continuum_config: ContinuumConfiguration | None = None,
        *,
        line_mask_width: u.Quantity = 1000.0 * u.km / u.s,
        box_width: u.Quantity = 1000.0 * u.km / u.s,
        error_scale: bool = False,
    ) -> None:
        """Estimate characteristic flux scales for lines and continuum.

        The algorithm:

        1. Mask pixels near emission lines (using *line_mask_fwhm*, a full
           FWHM convolved in quadrature with the local LSF).
        2. Fit a low-order polynomial to the unmasked pixels in each
           continuum region.
        3. Estimate the **line scale** as ``peak_above_continuum * line_width``
           where the peak is measured relative to the fitted continuum.
        4. Estimate the **continuum scale** as the maximum median ``|flux|``
           across continuum regions (after masking lines).
        5. Optionally compute per-region **error scale** factors
           (``sqrt(chi2_red)``) and store them as per-pixel
           arrays on each :class:`~unite.spectrum.Spectrum`.

        Both flux scales are stored as :class:`~astropy.units.Quantity`
        objects: integrated flux for ``line_scale`` and flux density for
        ``continuum_scale``.

        Parameters
        ----------
        line_config : LineConfiguration
            Line configuration.
        continuum_config : ContinuumConfiguration, optional
            Continuum configuration.  If ``None``, only line scale is
            computed (without continuum subtraction).
        max_fwhm : astropy.units.Quantity
            Maximum expected intrinsic line FWHM (full width).  Must have
            velocity units.  Default ``1000 km/s``.
        line_mask_fwhm : astropy.units.Quantity
            Full FWHM used for masking lines when fitting continuum.  Must
            have velocity units.  Default ``1000 km/s``.
        error_scale : bool
            If ``True``, compute per-region error scale factors from the
            reduced chi-squared of the continuum fit residuals and store
            them as per-pixel arrays on each spectrum via
            :attr:`~unite.spectrum.Spectrum.error_scale`.
            Default ``False``.
        """
        box_width = _ensure_velocity(box_width, 'max_fwhm', ndim=0)
        line_mask_width = _ensure_velocity(line_mask_width, 'line_mask_fwhm', ndim=0)
        box_width_kms = float(box_width.to(u.km / u.s).value)
        mask_width_kms = float(line_mask_width.to(u.km / u.s).value)

        z = self._redshift
        ref_flux_unit = self._spectra[0].flux_unit
        ref_wl_unit = self._canonical_unit

        # --- Helper: build a per-pixel line exclusion mask for a spectrum ---
        def _build_line_mask(spectrum, width_kms):
            wl = spectrum.wavelength
            mask = jnp.zeros(spectrum.npix, dtype=bool)
            for lam_rest in line_config.wavelengths:
                lam_obs = float(lam_rest.to(spectrum.unit).value) * (1.0 + z)
                lsf_fwhm = lam_obs / float(spectrum.disperser.R(lam_obs))
                user_width = lam_obs * width_kms / C_KMS
                # half_width = float(jnp.sqrt(user_width**2 + lsf_fwhm**2)) / 2
                half_width = (user_width + lsf_fwhm) / 2
                mask = mask | (
                    (wl > lam_obs - half_width) & (wl < lam_obs + half_width)
                )
            return mask

        # --- Helper: fit continuum form in a region ---
        def _fit_continuum_region(wl, flux, error, obs_low, obs_high, line_mask, form):
            """Fit the region's continuum form to unmasked pixels.

            Returns (model_on_region, in_region, good, chi2_red, fit_params).
            model_on_region covers all pixels in_region; good is the
            unmasked pixel mask used for fitting.
            """
            from unite.continuum.fit import fit_continuum_form

            in_region = (wl >= obs_low) & (wl <= obs_high)
            good = in_region & ~line_mask
            n_good = int(jnp.sum(good))

            min_params = form.n_params  # includes normalization_wavelength
            if n_good < max(min_params + 1, 3):
                return None, in_region, good, None, {}

            center = float((obs_low + obs_high) / 2.0)
            result = fit_continuum_form(
                form, wl[good], flux[good], error[good], center, obs_low, obs_high
            )

            # Evaluate on all in-region pixels.
            model_region = form.evaluate(
                wl[in_region], center, result.params, obs_low, obs_high
            )

            return model_region, in_region, good, result.chi2_red, result.params

        # --- Single pass: fit continuum once per region per spectrum, then
        #     derive line scale, continuum scale, error scale, and diagnostics.
        max_line_scale = 0.0
        max_cont_scale = 0.0
        any_good_continuum_pixels = False
        any_successful_fit = False
        diag_list = []

        for spectrum in self._spectra:
            wl = spectrum.wavelength
            flux_conv = _get_conversion_factor(spectrum.flux_unit, ref_flux_unit)
            wl_conv = _get_conversion_factor(spectrum.unit, ref_wl_unit)
            line_mask = _build_line_mask(spectrum, mask_width_kms)

            continuum_model = jnp.full(spectrum.npix, jnp.nan)
            per_pixel_scale = jnp.ones(spectrum.npix)
            region_diags: list[RegionDiagnostic] = []

            if continuum_config is not None:
                for region in continuum_config:
                    conv = _get_conversion_factor(region._unit, spectrum.unit)
                    obs_low = region.low * conv * (1.0 + z)
                    obs_high = region.high * conv * (1.0 + z)

                    in_region = (wl >= obs_low) & (wl <= obs_high)
                    if not jnp.any(in_region):
                        continue

                    model_region, in_region, good, chi2_red, fit_params = (
                        _fit_continuum_region(
                            wl,
                            spectrum.flux,
                            spectrum.error,
                            obs_low,
                            obs_high,
                            line_mask,
                            region.form,
                        )
                    )

                    # Continuum estimate for line-scale peak subtraction.
                    if model_region is not None:
                        continuum_model = continuum_model.at[in_region].set(
                            model_region
                        )
                        any_successful_fit = True

                    # Continuum scale: median |flux| of unmasked pixels.
                    if jnp.any(good):
                        any_good_continuum_pixels = True
                        median_flux = float(jnp.median(jnp.abs(spectrum.flux[good])))
                        max_cont_scale = max(max_cont_scale, median_flux * flux_conv)

                    # Error scale: per-pixel sqrt(chi2_red) from fit residuals.
                    if error_scale and chi2_red is not None:
                        region_scale = float(jnp.sqrt(chi2_red))
                        per_pixel_scale = jnp.where(
                            in_region, region_scale, per_pixel_scale
                        )

                    region_diags.append(
                        RegionDiagnostic(
                            obs_low=float(obs_low),
                            obs_high=float(obs_high),
                            in_region=in_region,
                            good_mask=good,
                            model_on_region=model_region,
                            chi2_red=chi2_red,
                            fit_params=fit_params,
                        )
                    )

            if error_scale and (continuum_config is not None):
                spectrum.error_scale = per_pixel_scale

            # Line scale: peak above continuum within ±total_fwhm of each line.
            # Where continuum is not fitted (NaN), assume zero baseline.
            baseline = jnp.where(jnp.isnan(continuum_model), 0.0, continuum_model)
            for lam_rest in line_config.wavelengths:
                lam_obs = float(lam_rest.to(spectrum.unit).value) * (1.0 + z)
                lsf_fwhm = lam_obs / float(jnp.asarray(spectrum.disperser.R(lam_obs)))
                mask_width_lam = lam_obs * mask_width_kms / C_KMS
                half_width = float(jnp.sqrt(mask_width_lam**2 + lsf_fwhm**2)) / 2
                in_window = (wl >= lam_obs - half_width) & (wl <= lam_obs + half_width)
                if not jnp.any(in_window):
                    continue
                peak_above = float(
                    jnp.max(spectrum.flux[in_window] - baseline[in_window])
                )
                box_width_lam = box_width_kms * lam_obs / C_KMS
                flux_est = peak_above * flux_conv * box_width_lam * wl_conv
                max_line_scale = max(max_line_scale, flux_est)

            diag = SpectrumScaleDiagnostic(
                name=spectrum.name,
                line_mask=line_mask,
                continuum_model=continuum_model,
                regions=region_diags,
            )
            spectrum._scale_diagnostic = diag
            diag_list.append(diag)

        if max_line_scale <= 0:
            msg = (
                'Could not estimate line flux scale: no emission line peak '
                'was found above the continuum in any spectrum. This can '
                'happen if (a) the configured lines fall outside the spectral '
                'coverage, (b) the continuum model over-subtracts the flux, '
                'or (c) the lines are too faint relative to noise. '
                'Check that line_config wavelengths match the observed data, '
                'or set spectra.line_scale manually.'
            )
            raise ValueError(msg)
        self._line_scale = max_line_scale * ref_flux_unit * ref_wl_unit

        if continuum_config is not None:
            if not any_good_continuum_pixels:
                msg = (
                    'Could not estimate continuum scale: the line mask covers '
                    'all pixels in every continuum region, leaving no unmasked '
                    'pixels for fitting. Try reducing line_mask_width (currently '
                    f'{line_mask_width}) or widening the continuum regions.'
                )
                raise ValueError(msg)
            if not any_successful_fit:
                msg = (
                    'Could not estimate continuum scale: all continuum region '
                    'fits failed (too few unmasked pixels in every region after '
                    'line masking). Try reducing line_mask_width (currently '
                    f'{line_mask_width}) or widening the continuum regions.'
                )
                raise ValueError(msg)
            if max_cont_scale <= 0:
                msg = (
                    'Could not estimate continuum scale: the median |flux| in '
                    'all continuum regions is zero. Check that the input spectra '
                    'contain non-zero flux values, or set '
                    'spectra.continuum_scale manually.'
                )
                raise ValueError(msg)
            self._continuum_scale = max_cont_scale * ref_flux_unit

        self._scale_diagnostics = ScaleDiagnosticList(diag_list)

    # -- preparation ----------------------------------------------------------

    def prepare(
        self,
        line_config: LineConfiguration,
        continuum_config: ContinuumConfiguration | None = None,
        *,
        linedet_width: u.Quantity = 1000.0 * u.km / u.s,
        drop_empty_regions: bool = True,
    ) -> tuple[LineConfiguration, ContinuumConfiguration | None]:
        """Filter configs for coverage and optionally drop empty continuum regions.

        Calls :meth:`filter_config` to remove lines and continuum regions not
        covered by any spectrum, then optionally drops continuum regions that
        contain no emission lines.  Stores the filtered configs on this
        object.

        Parameters
        ----------
        line_config : LineConfiguration
            Line configuration.
        continuum_config : ContinuumConfiguration, optional
            Continuum configuration.
        linedet_width : astropy.units.Quantity
            Detection width.  Must have velocity units.
            Default ``1000 km/s``.
        drop_empty_regions : bool
            If ``True`` (default), drop continuum regions that do not contain
            any filtered line (in rest frame).

        Returns
        -------
        filtered_lines : LineConfiguration
        filtered_continuum : ContinuumConfiguration or None
        """
        from unite.continuum.config import ContinuumConfiguration

        filtered_lines, filtered_cont = self.filter_config(
            line_config, continuum_config, linedet_width=linedet_width
        )

        if drop_empty_regions and filtered_cont is not None and len(filtered_lines) > 0:
            kept = []
            for region in filtered_cont:
                # Check if any filtered line falls in this region (rest frame).
                has_line = False
                for lam_rest in filtered_lines.wavelengths:
                    lam_val = float(lam_rest.to(region._unit).value)
                    if region.low <= lam_val <= region.high:
                        has_line = True
                        break
                if has_line:
                    kept.append(region)
            filtered_cont = ContinuumConfiguration(kept) if kept else None

        self._prepared_line_config = filtered_lines
        self._prepared_cont_config = filtered_cont
        self._is_prepared = True

        return filtered_lines, filtered_cont

    @property
    def is_prepared(self) -> bool:
        """``True`` if :meth:`prepare` has been called."""
        return self._is_prepared

    @property
    def prepared_line_config(self) -> LineConfiguration | None:
        """Filtered line configuration from :meth:`prepare`, or ``None``."""
        return self._prepared_line_config

    @property
    def prepared_cont_config(self) -> ContinuumConfiguration | None:
        """Filtered continuum configuration from :meth:`prepare`, or ``None``."""
        return self._prepared_cont_config

    # -- coverage filtering ---------------------------------------------------

    def filter_config(
        self,
        line_config: LineConfiguration,
        continuum_config: ContinuumConfiguration | None = None,
        *,
        linedet_width: u.Quantity = 1000.0 * u.km / u.s,
    ) -> tuple[LineConfiguration, ContinuumConfiguration | None]:
        """Drop lines and continuum regions not covered by any spectrum.

        Each line's rest-frame wavelength is shifted to the observed frame
        using :attr:`redshift`, then padded by *linedet_width* to form a
        detection window.  A line is kept if **any** spectrum partially
        overlaps that window.  Continuum regions are checked the same way.

        Parameters
        ----------
        line_config : LineConfiguration
            Line configuration to filter.
        continuum_config : ContinuumConfiguration, optional
            Continuum configuration to filter.  ``None`` is passed through.
        linedet_width : astropy.units.Quantity
            Detection width.  Must have velocity units.
            Default ``1000 km/s``.

        Returns
        -------
        filtered_lines : LineConfiguration
        filtered_continuum : ContinuumConfiguration or None
        """
        from unite.continuum.config import ContinuumConfiguration

        linedet_width = _ensure_velocity(linedet_width, 'linedet_width', ndim=0)
        eps = float(linedet_width.to(u.km / u.s).value) / (2 * C_KMS)
        z = self._redshift

        # --- filter lines ---
        mask = []
        for wavelength in line_config.wavelengths:
            # wavelength is a Quantity; convert to each spectrum's disperser unit
            lam_obs = cast(u.Quantity, wavelength * (1.0 + z))
            covered = any(
                s.covers(
                    float(lam_obs.to(s.unit).value) * (1.0 - eps),
                    float(lam_obs.to(s.unit).value) * (1.0 + eps),
                )
                for s in self._spectra
            )
            mask.append(covered)
        filtered_lines = line_config._filter(mask)

        # --- filter continuum regions ---
        if continuum_config is not None:
            kept = []
            for region in continuum_config:
                # Convert region bounds (stored as bare floats in region._unit)
                # to each spectrum's disperser unit before checking coverage.
                region_covered = False
                for s in self._spectra:
                    conv = _get_conversion_factor(region._unit, s.unit)
                    obs_low = region.low * conv * (1.0 + z)
                    obs_high = region.high * conv * (1.0 + z)
                    if s.covers(obs_low, obs_high):
                        region_covered = True
                        break
                if region_covered:
                    kept.append(region)
            filtered_cont: ContinuumConfiguration | None = ContinuumConfiguration(kept)
        else:
            filtered_cont = None

        return filtered_lines, filtered_cont

    # -- container interface --------------------------------------------------

    def __len__(self) -> int:
        return len(self._spectra)

    def __iter__(self) -> Iterator[Spectrum]:
        return iter(self._spectra)

    def __getitem__(self, idx: int | str) -> Spectrum:
        """Return a spectrum by integer index or by name.

        Parameters
        ----------
        idx : int or str
            Integer index into the internal spectrum list, or a string
            matching :attr:`~unite.spectrum.Spectrum.name`
            of the desired spectrum.

        Raises
        ------
        KeyError
            If *idx* is a string and no spectrum with that name exists.
        IndexError
            If *idx* is an integer and is out of range.
        """
        if isinstance(idx, str):
            for spectrum in self._spectra:
                if spectrum.name == idx:
                    return spectrum
            msg = f'No spectrum found with name {idx!r}.'
            raise KeyError(msg)
        return self._spectra[idx]

    def __repr__(self) -> str:
        lines = [f'Spectra: {len(self._spectra)} spectrum/a, z={self._redshift:.4f}']
        for i, s in enumerate(self._spectra):
            lo, hi = s.wavelength_range
            unit_str = s.unit.to_string()
            label = s.name or f'#{i}'
            cal = ' [calibrated]' if s.has_calibration_priors else ''
            lines.append(
                f'  [{i}] {label}: {s.npix} px, λ ∈ [{lo:.4g}, {hi:.4g}] {unit_str}{cal}'
            )
        return '\n'.join(lines)
