"""Minimization interface for beam-parameter inference.

Provides structured entry-points (``run_minimization_inference`` and
``run_multistart_minimization``) that wrap the numerical engine.

``BeamMinimizationModel`` is kept for backward compatibility but delegates
all loss computation to :mod:`beam_infer.engine`.
"""

from dataclasses import dataclass
from typing import Any, List, Optional, Tuple, Dict

import jax.numpy as jnp
import jax.random as jr
import numpy as np

from beam_infer.base_model import BeamForwardModel
from .engine import (
    estimate_parameters_bfgs,
    loss_function as engine_loss_function,
)

ArrayLike = np.ndarray | jnp.ndarray


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class MinimizationConfig:
    """Configuration for minimization-based parameter estimation."""

    initial_guess: ArrayLike
    bounds: Optional[Tuple[ArrayLike, ArrayLike]] = None
    max_iterations: int = 100
    verbose: bool = False
    regularization: Optional[Dict[str, float]] = None
    """Regularization strengths by parameter name. Keys are parameter names
    (e.g. 'Ax', 'sigx'), values are regularization strengths (quadratic
    penalty around the initial guess)."""


@dataclass(frozen=True)
class MinimizationResult:
    """Normalized result object returned by minimization entry-point."""

    estimated_parameters: np.ndarray
    objective_value: float
    success: bool
    message: str
    iterations: int


@dataclass(frozen=True)
class MultistartResult:
    """Result from multistart optimization."""

    results: List[MinimizationResult]
    """All optimization results, sorted by objective_value (best first)."""
    best_result: MinimizationResult
    """Best result (lowest objective value)."""


# ---------------------------------------------------------------------------
# Backward-compatible model class
# ---------------------------------------------------------------------------


class BeamMinimizationModel(BeamForwardModel):
    """Self-contained minimization model (forward + loss + L-BFGS-B).

    The loss function delegates to :func:`beam_infer.engine.loss_function`
    (L2-normalised on both simulated and observed images), matching the
    original implementation.
    """

    def loss_function(
        self,
        k: jnp.ndarray,
        i_obs: jnp.ndarray,
        x: jnp.ndarray,
        y: jnp.ndarray,
        t_vals: jnp.ndarray,
        regularization: Optional[Dict[str, float]] = None,
        param_names: Optional[List[str]] = None,
    ) -> jnp.ndarray:
        """Compute loss with optional regularization.

        Uses L2 normalisation on both images (shape comparison only),
        matching the original ``loss_function``.

        When *regularization* is supplied the penalty is
        ``lambda * (k[i] - k0[i])^2`` where *k0* is the **initial guess
        stored at the time config_regularization was called**.  For the
        simplified backward-compatible API the reference is not available,
        so regularisation here penalises absolute value (``k[i]^2``).
        Prefer using :func:`beam_infer.engine.loss_regularized` directly for
        the full original behaviour.
        """
        base = engine_loss_function(k, i_obs, x, y, t_vals)

        reg_loss = jnp.float32(0.0)
        if regularization is not None and param_names is not None:
            for param_name, strength in regularization.items():
                if strength > 0:
                    try:
                        idx = param_names.index(param_name)
                        reg_loss = reg_loss + strength * (k[idx] ** 2)
                    except ValueError:
                        pass

        return base + reg_loss

    def estimate_parameters_bfgs(
        self,
        i_obs: jnp.ndarray,
        x: jnp.ndarray,
        y: jnp.ndarray,
        t_vals: jnp.ndarray,
        k0: jnp.ndarray,
        max_iterations: int = 100,
        verbose: bool = False,
        bounds: Optional[Tuple[np.ndarray, np.ndarray]] = None,
        regularization: Optional[Dict[str, float]] = None,
        param_names: Optional[List[str]] = None,
    ) -> Any:
        """Run L-BFGS-B optimisation, delegating to the engine."""

        def _loss_fn(k, I_obs, X, Y, tv):
            return self.loss_function(
                k, I_obs, X, Y, tv, regularization, param_names
            )

        return estimate_parameters_bfgs(
            I_obs=i_obs,
            X=x,
            Y=y,
            t_vals=t_vals,
            k0=k0,
            loss_fn=_loss_fn,
            bounds=bounds,
            maxiter=max_iterations,
            verbose=verbose,
        )


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _as_jax_array(value: ArrayLike) -> jnp.ndarray:
    return jnp.asarray(value, dtype=jnp.float32)


def _normalize_bounds(
    bounds: Optional[Tuple[ArrayLike, ArrayLike]],
) -> Optional[Tuple[np.ndarray, np.ndarray]]:
    if bounds is None:
        return None

    lower, upper = bounds
    lower_np = np.asarray(lower, dtype=np.float64)
    upper_np = np.asarray(upper, dtype=np.float64)

    if lower_np.shape != upper_np.shape:
        raise ValueError(
            "Bounds must have matching shapes for lower and upper arrays."
        )
    if np.any(lower_np > upper_np):
        raise ValueError(
            "Bounds are invalid: lower bound exceeds upper bound."
        )

    return lower_np, upper_np


# ---------------------------------------------------------------------------
# Public entry-points
# ---------------------------------------------------------------------------


def run_minimization_inference(
    image: ArrayLike,
    x_grid: ArrayLike,
    y_grid: ArrayLike,
    t_vals: ArrayLike,
    config: MinimizationConfig,
    param_names: Optional[List[str]] = None,
) -> MinimizationResult:
    """Entry-point for deterministic (L-BFGS-B) inference.

    Parameters
    ----------
    image :
        Observed 2D beam image.
    x_grid, y_grid :
        2D coordinate grids matching image shape.
    t_vals :
        1D time discretization used by the forward model.
    config :
        Minimization hyperparameters and initialisation.
    param_names :
        Parameter names for regularisation (required when regularisation
        is used).
    """
    i_obs = _as_jax_array(image)
    x = _as_jax_array(x_grid)
    y = _as_jax_array(y_grid)
    t = _as_jax_array(t_vals)
    k0 = _as_jax_array(config.initial_guess)
    bounds = _normalize_bounds(config.bounds)

    if i_obs.shape != x.shape or i_obs.shape != y.shape:
        raise ValueError("image, x_grid, and y_grid must have the same 2D shape.")
    if t.ndim != 1:
        raise ValueError("t_vals must be a 1D array.")
    if k0.ndim != 1:
        raise ValueError("initial_guess must be a 1D vector.")
    if bounds is not None and bounds[0].shape[0] != k0.shape[0]:
        raise ValueError("Bounds dimension must match initial_guess dimension.")
    if config.regularization is not None and param_names is None:
        raise ValueError("param_names must be provided when using regularization.")

    model = BeamMinimizationModel()
    raw_result = model.estimate_parameters_bfgs(
        i_obs=i_obs,
        x=x,
        y=y,
        t_vals=t,
        k0=k0,
        max_iterations=config.max_iterations,
        verbose=config.verbose,
        bounds=bounds,
        regularization=config.regularization,
        param_names=param_names,
    )

    return MinimizationResult(
        estimated_parameters=np.asarray(raw_result.x, dtype=np.float64),
        objective_value=float(raw_result.fun),
        success=bool(raw_result.success),
        message=str(raw_result.message),
        iterations=int(getattr(raw_result, "nit", 0)),
    )


def run_multistart_minimization(
    image: ArrayLike,
    x_grid: ArrayLike,
    y_grid: ArrayLike,
    t_vals: ArrayLike,
    config: MinimizationConfig,
    num_starts: int = 5,
    param_names: Optional[List[str]] = None,
    random_seed: int = 42,
    free_param_indices: Optional[List[int]] = None,
) -> MultistartResult:
    """Run optimisation from multiple starting points.

    Parameters
    ----------
    image :
        Observed 2D beam image.
    x_grid, y_grid :
        2D coordinate grids matching image shape.
    t_vals :
        1D time discretization used by the forward model.
    config :
        Minimization hyperparameters and initialisation.
    num_starts :
        Number of different starting points.
    param_names :
        Parameter names for regularisation.
    random_seed :
        Random seed for generating starting points.
    free_param_indices :
        Indices of free (non-fixed) parameters to randomize.
        If ``None``, all parameters with different bounds are randomized.
    """
    k0_base = np.asarray(config.initial_guess)
    bounds = _normalize_bounds(config.bounds)

    if free_param_indices is None and bounds is not None:
        lower, upper = bounds
        free_param_indices = [
            i for i in range(len(k0_base)) if abs(lower[i] - upper[i]) > 1e-10
        ]
    elif free_param_indices is None:
        free_param_indices = list(range(len(k0_base)))

    key = jr.PRNGKey(random_seed)
    results: List[MinimizationResult] = []

    for i in range(num_starts):
        if i == 0:
            k0 = k0_base.copy()
        else:
            key, subkey = jr.split(key)
            k0 = k0_base.copy()
            if bounds is not None:
                lower, upper = bounds
                key, subkey = jr.split(key)
                u_vec = jr.uniform(subkey, shape=(len(free_param_indices),))
                for j, idx in enumerate(free_param_indices):
                    k0[idx] = lower[idx] + float(u_vec[j]) * (upper[idx] - lower[idx])
            else:
                key, subkey = jr.split(key)
                noise = jr.normal(subkey, shape=k0.shape) * 0.1 * np.abs(k0)
                k0 = k0 + np.asarray(noise)

        start_config = MinimizationConfig(
            initial_guess=k0,
            bounds=config.bounds,
            max_iterations=config.max_iterations,
            verbose=config.verbose and (i == 0),
            regularization=config.regularization,
        )

        result = run_minimization_inference(
            image=image,
            x_grid=x_grid,
            y_grid=y_grid,
            t_vals=t_vals,
            config=start_config,
            param_names=param_names,
        )
        results.append(result)

    results.sort(key=lambda r: r.objective_value)
    return MultistartResult(results=results, best_result=results[0])
