"""Model builder and numpyro model function for spectral line fitting.

The :class:`ModelBuilder` assembles a :class:`~unite.line.config.LineConfiguration`,
an optional :class:`~unite.continuum.config.ContinuumConfiguration`, and a
:class:`~unite.spectrum.spectrum.Spectra` collection into a numpyro model
function that can be passed to any numpyro inference algorithm (NUTS, SVI, etc.).
"""

from __future__ import annotations

import warnings
from collections.abc import Callable
from dataclasses import dataclass

import jax.numpy as jnp
import numpyro
from numpyro import distributions as dist

from unite._utils import C_KMS, _get_conversion_factor
from unite.continuum.config import ContinuumConfiguration
from unite.line.config import ConfigMatrices, LineConfiguration
from unite.line.profiles import integrate_lines
from unite.prior import Fixed, Parameter, Prior, topological_sort
from unite.spectrum import Spectra, Spectrum

# ------------------------------------------------------------------
# ModelArgs — data bundle for the numpyro model function
# ------------------------------------------------------------------


@dataclass
class ModelArgs:
    """Bundle of arguments passed to :func:`unite_model`.

    Created by :meth:`ModelBuilder.build`; not intended for direct
    construction by users.
    """

    #: Precomputed parameter matrices and line metadata.
    matrices: ConfigMatrices
    #: Individual spectra.
    spectra: list[Spectrum]
    #: Systemic redshift.
    redshift: float
    #: Continuum configuration, or ``None`` if not used.
    cont_config: ContinuumConfiguration | None
    #: Resolved ``{param_name: ContinuumParam}`` mappings per region, from :attr:`ContinuumConfiguration.resolved_params`.
    cont_resolved_params: list[dict[str, Parameter]] | None
    #: All parameters with their priors (line, calibration, continuum).
    all_priors: dict[str, Prior]
    #: Topological sampling order for all parameters.
    dependency_order: list[str]
    name_to_token: dict[str, object]
    # --- Wavelength unit conversion ---
    spec_to_canonical: list[float]
    # --- Pre-converted continuum bounds (rest-frame, canonical unit) ---
    cont_low: list[float] | None
    cont_high: list[float] | None
    cont_center: list[float] | None
    # --- Unit conversion factors: region._unit → canonical unit, per region ---
    cont_nw_conv: list[float] | None
    #: Pre-converted continuum forms (static wavelength config in canonical unit).
    cont_forms: list | None
    # --- Flux normalization ---
    norm_factors: list[float]
    #: Per-spectrum line flux scale (in each spectrum's flux_unit * canonical_wl_unit).
    line_flux_scales: list[float]
    #: Per-spectrum continuum scale (in each spectrum's flux_unit).
    continuum_scales: list[float]
    #: Wavelength unit of canonical frame (first spectrum's disperser unit).
    canonical_unit: object
    #: Per-spectrum flux density units.
    flux_units: list
    #: The Quantity line_scale and continuum_scale from Spectra (for results output).
    line_scale_quantity: object
    continuum_scale_quantity: object
    #: Human-readable column labels for each line, parallel to ``matrices.wavelengths``.
    #: Derived from user-supplied line names and rest-frame wavelengths.
    line_labels: list[str]
    #: Human-readable column labels for each continuum region, parallel to ``cont_config``.
    #: Derived from form type and wavelength bounds.
    continuum_labels: list[str]

    def __len__(self) -> int:
        """Return the number of spectra in the model."""
        return len(self.spectra)

    def __bool__(self) -> bool:
        """Return True if the model has at least one spectrum."""
        return len(self.spectra) > 0


# ------------------------------------------------------------------
# Numpyro model function
# ------------------------------------------------------------------


def unite_model(args: ModelArgs) -> None:
    """Numpyro model function for multi-spectrum emission-line fitting.

    All lines are integrated simultaneously via :func:`jax.vmap` with
    ``lax.switch`` dispatching to the correct profile kernel per line.
    Parameter broadcasting from unique tokens to per-line arrays is done
    with precomputed indicator matrices.

    Wavelength unit conversion is handled via pre-computed scalar factors
    stored in ``args.spec_to_canonical`` (one per spectrum).  Flux is
    normalized per spectrum so that the likelihood operates on O(1) values.

    Parameters
    ----------
    args : ModelArgs
        Pre-built data bundle from :meth:`ModelBuilder.build`.
    """
    cm = args.matrices
    z_sys = args.redshift
    n_lines = cm.wavelengths.shape[0]

    # --- 1. Sample all parameters in dependency order ---
    # Two parallel dicts are maintained:
    #   context    — str → value, used for all downstream name-based lookups.
    #   obj_ctx    — token_object → value, passed to prior.to_dist() so that
    #                parameter expressions can look up dependency values by
    #                token identity.
    context: dict[str, jnp.ndarray] = {}
    obj_ctx: dict[object, jnp.ndarray] = {}
    for pname in args.dependency_order:
        prior = args.all_priors[pname]
        if isinstance(prior, Fixed):
            val = jnp.asarray(prior.resolved_value(obj_ctx))
        else:
            val = numpyro.sample(pname, prior.to_dist(obj_ctx))
        context[pname] = val
        tok = args.name_to_token.get(pname)
        if tok is not None:
            obj_ctx[tok] = val

    # --- 2. Per-line parameter arrays via matrix products ---
    # Flux: include multiplet strengths.
    flux_vec = jnp.stack([context[n] for n in cm.flux_names])
    flux_per_line = flux_vec @ cm.flux_matrix * cm.strengths  # (n_lines,)

    # Redshift (delta from systemic).
    z_vec = jnp.stack([context[n] for n in cm.z_names])
    z_per_line = z_vec @ cm.z_matrix  # (n_lines,)

    # Observed-frame centers (in canonical wavelength unit).
    centers = cm.wavelengths * (1.0 + z_sys + z_per_line)  # (n_lines,)

    # Slot 0: primary velocity FWHM → canonical wavelength units.
    if cm.p0_names:
        p0_vec = jnp.stack([context[n] for n in cm.p0_names])
        p0_kms = p0_vec @ cm.p0_matrix  # (n_lines,)
    else:
        p0_kms = jnp.zeros(n_lines)
    p0 = centers * p0_kms / C_KMS  # canonical wavelength units

    # Slot 1: velocity FWHMs (p1v) + dimensionless params (p1d), summed.
    if cm.p1v_names:
        p1v_vec = jnp.stack([context[n] for n in cm.p1v_names])
        p1v_kms = p1v_vec @ cm.p1v_matrix  # (n_lines,)
    else:
        p1v_kms = jnp.zeros(n_lines)
    p1v = centers * p1v_kms / C_KMS

    if cm.p1d_names:
        p1d_vec = jnp.stack([context[n] for n in cm.p1d_names])
        p1d = p1d_vec @ cm.p1d_matrix  # (n_lines,), dimensionless
    else:
        p1d = jnp.zeros(n_lines)
    p1 = p1v + p1d  # for any line, only one sub-matrix is nonzero

    # Slot 2: dimensionless params only.
    if cm.p2_names:
        p2_vec = jnp.stack([context[n] for n in cm.p2_names])
        p2 = p2_vec @ cm.p2_matrix  # (n_lines,)
    else:
        p2 = jnp.zeros(n_lines)

    # --- 3. Per-spectrum likelihood ---
    for i, spectrum in enumerate(args.spectra):
        disp = spectrum.disperser
        wl_scale = args.spec_to_canonical[i]
        inv_wl_scale = 1.0 / wl_scale
        norm = args.norm_factors[i]
        lfs = args.line_flux_scales[i]
        cs = args.continuum_scales[i]

        # Calibration values (fall back to identity when no token is attached).
        r_scale = context[disp.r_scale.name] if disp.r_scale is not None else 1.0
        flux_scale = (
            context[disp.flux_scale.name] if disp.flux_scale is not None else 1.0
        )
        pix_offset = (
            context[disp.pix_offset.name] if disp.pix_offset is not None else 0.0
        )

        # Pixel edges in canonical wavelength unit.
        low = spectrum.low * wl_scale
        high = spectrum.high * wl_scale
        if disp.pix_offset is not None:
            mid_disp = (spectrum.low + spectrum.high) / 2.0  # disperser unit
            shift = pix_offset * disp.dlam_dpix(mid_disp) * wl_scale
            low = low + shift
            high = high + shift

        wavelength = (low + high) / 2.0

        # LSF FWHM at each line center (n_lines,).
        # R() expects disperser-unit wavelengths, result is in canonical unit.
        lsf_fwhm = centers / (disp.R(centers * inv_wl_scale) * r_scale)

        # Integrate all lines simultaneously and divide by pixel width
        # to get average flux density per pixel.
        pixints = integrate_lines(
            low, high, centers, lsf_fwhm, p0, p1, p2, cm.profile_codes
        ) / (high - low)

        # Sum over lines weighted by flux, scaled by per-spectrum
        # line_flux_scale and divided by norm to bring to O(1) space.
        line_model = ((flux_per_line * lfs / norm)[:, None] * pixints).sum(
            axis=0
        )  # (n_pixels,)

        # Continuum (evaluated in canonical-unit wavelengths, then normalized).
        continuum = jnp.zeros(spectrum.npix)
        if args.cont_config is not None:
            for k in range(len(args.cont_config)):
                obs_low = args.cont_low[k] * (1.0 + z_sys)
                obs_high = args.cont_high[k] * (1.0 + z_sys)
                obs_center = args.cont_center[k] * (1.0 + z_sys)
                in_region = (wavelength >= obs_low) & (wavelength <= obs_high)
                cont_params = {
                    pn: (
                        context[tok.name] * args.cont_nw_conv[k] * (1.0 + z_sys)
                        if pn == 'norm_wav'
                        else context[tok.name]
                    )
                    for pn, tok in args.cont_resolved_params[k].items()
                }
                form = args.cont_forms[k]
                region_cont = form.evaluate(
                    wavelength, obs_center, cont_params, obs_low, obs_high
                )
                continuum = continuum + jnp.where(in_region, region_cont, 0.0)

        # Likelihood (normalized flux space).
        # Continuum is scaled by per-spectrum continuum_scale so that
        # sampled 'scale' parameters are O(1).
        model = flux_scale * (line_model + continuum * cs / norm)
        obs_name = f'obs_{spectrum.name}' if spectrum.name else f'obs_{i}'
        numpyro.sample(
            obs_name,
            dist.Normal(model, spectrum.scaled_error / norm),
            obs=spectrum.flux / norm,
        )


# ------------------------------------------------------------------
# ModelBuilder
# ------------------------------------------------------------------


class ModelBuilder:
    """Assemble configuration objects into a numpyro model.

    Collects all unique parameter tokens (line, calibration, continuum),
    builds precomputed indicator matrices, performs a topological sort for
    dependency resolution, and packages everything into a
    ``(model_fn, model_args)`` pair.

    Parameters
    ----------
    line_config : LineConfiguration
        Emission/absorption line configuration.
    continuum_config : ContinuumConfiguration or None
        Continuum configuration.  ``None`` for a lines-only model.
    spectra : Spectra
        Spectrum collection with systemic redshift.

    Examples
    --------
    >>> model_fn, args = ModelBuilder(line_config, cont, spectra).build()
    >>> kernel = numpyro.infer.NUTS(model_fn)
    >>> mcmc = numpyro.infer.MCMC(kernel, num_warmup=500, num_samples=1000)
    >>> mcmc.run(jax.random.PRNGKey(0), args)
    """

    def __init__(
        self,
        line_config: LineConfiguration,
        continuum_config: ContinuumConfiguration | None,
        spectra: Spectra,
    ) -> None:
        self._spectra = spectra

        # --- Auto-prepare if needed ---
        if not spectra.is_prepared:
            warnings.warn(
                'Spectra not prepared; calling spectra.prepare() with defaults. '
                'Call spectra.prepare(line_config, continuum_config) explicitly '
                'for full control.',
                UserWarning,
                stacklevel=2,
            )
            line_config, continuum_config = spectra.prepare(
                line_config, continuum_config
            )
        else:
            # Use the prepared configs.
            line_config = spectra.prepared_line_config
            continuum_config = spectra.prepared_cont_config

        # --- Auto-compute scales if needed ---
        if spectra.line_scale is None:
            warnings.warn(
                'Line scale not set; calling spectra.compute_scales() with '
                'defaults. Call spectra.compute_scales(line_config, '
                'continuum_config) explicitly for full control.',
                UserWarning,
                stacklevel=2,
            )
            spectra.compute_scales(line_config, continuum_config)

        self._cont_config = continuum_config

        # --- Canonical wavelength unit: use Spectra's canonical_unit ---
        self._canonical_unit = spectra.canonical_unit

        # Build precomputed matrices from line entries.
        self._matrices = line_config.build_matrices()

        # Convert line wavelengths to canonical unit.
        # Each line may have a different Quantity unit, so convert per-line.
        if len(line_config) > 0:
            canon_wls = jnp.array(
                [
                    float(e.wavelength.to(self._canonical_unit).value)
                    for e in line_config._entries
                ]
            )
            self._matrices.wavelengths = canon_wls

        # --- Collect all unique parameter tokens for prior / topo-sort ---
        all_priors: dict[str, Prior] = dict(self._matrices.priors)
        param_to_name: dict[object, str] = {
            # We can reconstruct token→name from the matrices' name lists and
            # the original entries since tokens carry their .name attribute.
            tok: tok.name
            for entry in line_config._entries
            for tok in (entry.flux, entry.redshift, *entry.fwhms.values())
        }

        # Calibration tokens from each unique disperser.
        seen_dispersers: set[int] = set()
        seen_tok_ids: set[int] = set(id(t) for t in param_to_name)
        for spectrum in spectra:
            disp = spectrum.disperser
            if id(disp) not in seen_dispersers:
                seen_dispersers.add(id(disp))
                for tok in (disp.r_scale, disp.flux_scale, disp.pix_offset):
                    if tok is not None and id(tok) not in seen_tok_ids:
                        seen_tok_ids.add(id(tok))
                        all_priors[tok.name] = tok.prior
                        param_to_name[tok] = tok.name

        # Continuum parameters: collect unique ContinuumParam tokens by identity.
        # Shared tokens (same object in multiple regions) produce one numpyro site.
        if continuum_config is not None:
            seen_cont_ids: set[int] = set()
            for resolved in continuum_config.resolved_params:
                for tok in resolved.values():
                    if id(tok) not in seen_cont_ids:
                        seen_cont_ids.add(id(tok))
                        all_priors[tok.name] = tok.prior
                        param_to_name[tok] = tok.name

        self._all_priors = all_priors
        self._dep_order = (
            topological_sort(all_priors, param_to_name) if all_priors else []
        )
        # Reverse mapping: site name → token object, for obj_ctx in unite_model.
        # Continuum params have no token objects and are intentionally absent.
        self._line_config = line_config
        self._name_to_token: dict[str, object] = {
            name: tok for tok, name in param_to_name.items()
        }

    @property
    def matrices(self) -> ConfigMatrices:
        """Precomputed matrices (after coverage filtering)."""
        return self._matrices

    def build(self) -> tuple[Callable, ModelArgs]:
        """Build the numpyro model function and its arguments.

        Returns
        -------
        model_fn : callable
            The numpyro model function (signature: ``model_fn(args)``).
        model_args : ModelArgs
            Pre-built data bundle to pass to the model function.
        """
        # Trim spectra to union of continuum regions (observed frame).
        # Pixels outside all regions have model = 0 and would corrupt the
        # likelihood if observed flux is nonzero.  Trimming also reduces
        # array sizes passed to JAX.
        z = self._spectra.redshift
        if self._cont_config is not None:
            trimmed_spectra: list[Spectrum] = []
            for s in self._spectra:
                mask = jnp.zeros(s.npix, dtype=bool)
                for region in self._cont_config:
                    conv = _get_conversion_factor(region._unit, s.unit)
                    obs_low = region.low * conv * (1.0 + z)
                    obs_high = region.high * conv * (1.0 + z)
                    mask = mask | s.pixel_mask(obs_low, obs_high)
                trimmed = s._sliced(mask)
                if trimmed.npix > 0:
                    trimmed_spectra.append(trimmed)
                else:
                    warnings.warn(
                        f'Spectrum {s.name!r} has no pixels overlapping '
                        f'any continuum region and will be excluded from '
                        f'the fit.',
                        UserWarning,
                        stacklevel=2,
                    )
            if not trimmed_spectra:
                warnings.warn(
                    'All spectra are fully masked after trimming to '
                    'continuum regions. The resulting model has no spectra '
                    'and should not be used for fitting. Check that your '
                    'continuum configuration covers the observed wavelength '
                    'range.',
                    UserWarning,
                    stacklevel=2,
                )
        else:
            trimmed_spectra = list(self._spectra)

        # Per-spectrum wavelength conversion factors.
        spec_to_canonical = [
            _get_conversion_factor(s.unit, self._canonical_unit)
            for s in trimmed_spectra
        ]

        # Per-spectrum flux normalization.
        norm_factors = [_compute_norm_factor(s) for s in trimmed_spectra]

        # Line flux scale (Quantity from Spectra.compute_scales).
        line_scale_qty = self._spectra.line_scale
        # Continuum scale (Quantity from Spectra.compute_scales; fallback).
        cont_scale_qty = self._spectra.continuum_scale

        # Convert Quantity scales to per-spectrum float values.
        # line_flux_scale needs units of [flux_density * canonical_wl_unit]
        # for each spectrum.
        # continuum_scale needs units of [flux_density] for each spectrum.
        canonical_unit = self._canonical_unit
        line_flux_scales: list[float] = []
        continuum_scales: list[float] = []
        for s in trimmed_spectra:
            target_line_unit = s.flux_unit * canonical_unit
            lfs = float(line_scale_qty.to(target_line_unit).value)
            line_flux_scales.append(lfs)

            if cont_scale_qty is not None:
                cs = float(cont_scale_qty.to(s.flux_unit).value)
            else:
                cs = 1.0
            continuum_scales.append(cs)

        flux_units = [s.flux_unit for s in trimmed_spectra]

        # Pre-convert continuum region bounds to canonical unit.
        if self._cont_config is not None:
            cont_low = []
            cont_high = []
            cont_center = []
            cont_nw_conv = []
            cont_forms = []
            for region in self._cont_config:
                conv = _get_conversion_factor(region._unit, self._canonical_unit)
                cont_low.append(region.low * conv)
                cont_high.append(region.high * conv)
                cont_center.append(region.center * conv)
                cont_nw_conv.append(conv)
                cont_forms.append(region.form)
        else:
            cont_low = None
            cont_high = None
            cont_center = None
            cont_nw_conv = None
            cont_forms = None

        line_labels = _make_line_labels(self._line_config)
        continuum_labels = (
            _make_continuum_labels(self._cont_config)
            if self._cont_config is not None
            else []
        )

        args = ModelArgs(
            matrices=self._matrices,
            spectra=trimmed_spectra,
            redshift=self._spectra.redshift,
            cont_config=self._cont_config,
            cont_resolved_params=(
                self._cont_config.resolved_params
                if self._cont_config is not None
                else None
            ),
            all_priors=self._all_priors,
            dependency_order=self._dep_order,
            name_to_token=self._name_to_token,
            spec_to_canonical=spec_to_canonical,
            cont_low=cont_low,
            cont_high=cont_high,
            cont_center=cont_center,
            cont_nw_conv=cont_nw_conv,
            cont_forms=cont_forms,
            norm_factors=norm_factors,
            line_flux_scales=line_flux_scales,
            continuum_scales=continuum_scales,
            canonical_unit=canonical_unit,
            flux_units=flux_units,
            line_scale_quantity=line_scale_qty,
            continuum_scale_quantity=cont_scale_qty,
            line_labels=line_labels,
            continuum_labels=continuum_labels,
        )
        return unite_model, args

    def fit(
        self,
        num_warmup: int = 250,
        num_samples: int = 1000,
        num_chains: int = 1,
        seed: int = 0,
        progress_bar: bool = True,
    ) -> dict:
        """Fit the model using NUTS sampling (convenience wrapper).

        This method builds the model, runs MCMC with the NUTS kernel, and
        returns the posterior samples. For more control over the sampler
        (e.g., custom kernel, SVI, nested sampling), call :meth:`build`
        directly and use numpyro's inference APIs.

        Parameters
        ----------
        num_warmup : int, optional
            Number of warmup iterations per chain (default: 1000).
        num_samples : int, optional
            Number of posterior samples per chain (default: 1000).
        num_chains : int, optional
            Number of MCMC chains to run in parallel (default: 1).
        seed : int, optional
            Random seed for JAX's PRNG (default: 0).
        progress_bar : bool, optional
            Whether to display a progress bar (default: True).

        Returns
        -------
        dict
            Posterior samples as a dictionary with parameter names as keys.
            Shape is ``(num_chains, num_samples)`` per parameter.

        Examples
        --------
        >>> samples = builder.fit(num_warmup=200, num_samples=500, num_chains=4)
        """
        import jax
        from numpyro import infer

        model_fn, model_args = self.build()
        mcmc = infer.MCMC(
            infer.NUTS(model_fn, dense_mass=True),
            num_warmup=num_warmup,
            num_samples=num_samples,
            num_chains=num_chains,
            progress_bar=progress_bar,
        )
        mcmc.run(jax.random.PRNGKey(seed), model_args)
        samples = mcmc.get_samples()
        return samples, model_args


# ------------------------------------------------------------------
# Normalization helpers
# ------------------------------------------------------------------


def _make_line_labels(line_config: LineConfiguration) -> list[str]:
    """Return the unique name for every line entry.

    Since line names are required to be unique within a
    :class:`~unite.line.config.LineConfiguration`, the label for each entry
    is simply its name.  Use ``add_lines`` (which auto-generates names as
    ``'{name}_{center.value:g}'``) or supply explicit unique names to
    ``add_line`` for multiplets and multi-component lines.
    """
    return [entry.name for entry in line_config._entries]


def _make_continuum_labels(cont_config: ContinuumConfiguration) -> list[str]:
    """Build a human-readable label for every continuum region.

    Format: ``'{form_type}_{low:.4g}_{high:.4g}'`` where the wavelength
    values are in the region's native unit and ``form_type`` is the
    lower-cased class name of the functional form (e.g. ``'linear'``,
    ``'powerlaw'``, ``'polynomial'``).

    Examples: ``'linear_6400_6700'``, ``'powerlaw_0.95_2.5'``.
    """
    labels: list[str] = []
    for region in cont_config:
        form_type = type(region.form).__name__.lower()
        low_str = f'{region.low:.4g}'
        high_str = f'{region.high:.4g}'
        labels.append(f'{form_type}_{low_str}_{high_str}')
    return labels


def _compute_norm_factor(s: Spectrum) -> float:
    """Robust scale factor to bring a spectrum's flux to ~O(1).

    Uses the median of the absolute non-zero flux values.

    Parameters
    ----------
    s : Spectrum

    Returns
    -------
    float
        Positive normalization factor.
    """
    absflux = jnp.abs(s.flux)
    positive = absflux[absflux > 0]
    if positive.size == 0:
        fallback = float(jnp.max(s.error))
        return fallback if fallback > 0 else 1.0
    return float(jnp.median(positive))
