"""Interface layer: user-facing function for beam parameter estimation.

This module is the **only** entry-point external code should call.  It
orchestrates the engine helpers and returns a structured result.
"""

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

import jax.numpy as jnp
import numpy as np

from . import engine

ArrayLike = np.ndarray | jnp.ndarray


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


@dataclass
class Metadata:
    """Known or measured quantities for beam estimation.

    Attributes:
        pulse_duration_ms: Pulse duration in milliseconds.
        fx: Horizontal scan frequency in kHz (optional; fixed if given).
        fy: Vertical scan frequency in kHz (optional; fixed if given).
        sampling_rate: Temporal sampling rate in Hz.
    """

    pulse_duration_ms: float = 0.05
    fx: Optional[float] = None
    fy: Optional[float] = None
    sampling_rate: int = 2_000_000


@dataclass
class OptimizerConfig:
    """User choices for the optimiser.

    Attributes:
        debug: Return all intermediate results.
        verbose: Print iteration details.
        max_iterations: Maximum L-BFGS-B iterations per run.
        multistart: Multistart spec, e.g.
            ``{"n_starts": 5, "type": "uniform"}``.
        regularization: Regularization spec, e.g.
            ``{"params": {"sigx": 0.1}}``.
    """

    debug: bool = False
    verbose: bool = False
    max_iterations: int = 100
    multistart: Optional[Dict[str, Any]] = None
    regularization: Optional[Dict[str, Any]] = None
    random_seed: Optional[int] = None


@dataclass
class OptimizerResult:
    """Result from :func:`estimate_with_optimizer`.

    Attributes:
        estimated_parameters: Best estimated parameter vector (10-D).
        objective_value: Final loss value.
        success: Whether the optimiser converged.
        message: Optimiser status message.
        iterations: Number of iterations used.
    """

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

    all_results: Optional[List[Tuple[float, np.ndarray]]] = None
    """All ``(loss, k)`` pairs sorted by loss (debug mode only)."""
    cropped_image: Optional[np.ndarray] = None
    crop_bounds: Optional[Tuple[int, int, int, int]] = None
    x_grid: Optional[jnp.ndarray] = None
    y_grid: Optional[jnp.ndarray] = None
    t_vals: Optional[jnp.ndarray] = None


# ---------------------------------------------------------------------------
# Input parsing
# ---------------------------------------------------------------------------


def parse_input(
    image: ArrayLike,
    meta: Optional[Metadata] = None,
) -> Tuple[jnp.ndarray, float]:
    """Validate the raw image and extract the pulse duration.

    Args:
        image: Raw 2D camera image.
        meta: Metadata with pulse duration.

    Returns:
        Tuple ``(image_jnp, pulse_duration_ms)``.

    Raises:
        ValueError: If the image is not 2D.
    """
    image_jnp = jnp.asarray(image, dtype=jnp.float32)
    if image_jnp.ndim != 2:
        raise ValueError(f"Image must be 2D, got shape {image_jnp.shape}")

    pulse_duration_ms = 0.05
    if meta is not None:
        pulse_duration_ms = meta.pulse_duration_ms

    return image_jnp, pulse_duration_ms


def _build_engine_metadata(meta: Metadata) -> Dict[str, float]:
    """Build engine meta dict from Metadata (fx/fy only when set)."""
    engine_meta: Dict[str, float] = {}
    if meta.fx is not None:
        engine_meta["fx"] = meta.fx
    if meta.fy is not None:
        engine_meta["fy"] = meta.fy
    return engine_meta


def _prepare_problem(
    image_jnp: jnp.ndarray,
    meta: Metadata,
) -> Tuple[
    jnp.ndarray,
    Tuple[int, int, int, int],
    jnp.ndarray,
    jnp.ndarray,
    jnp.ndarray,
    jnp.ndarray,
    jnp.ndarray,
    jnp.ndarray,
    Dict[int, float],
]:
    """Preprocess image, build grids, and set initial guess and bounds."""
    I_obs, bbox = engine.preprocess_image(image_jnp)
    H, W = I_obs.shape
    T = meta.pulse_duration_ms
    X, Y, t_vals = engine.build_grid_and_time(H, W, T, meta.sampling_rate)
    engine_meta = _build_engine_metadata(meta)
    k0, lower, upper, fixed_params = engine.init_guess_and_bounds(
        H, W, meta=engine_meta
    )
    return I_obs, bbox, X, Y, t_vals, k0, lower, upper, fixed_params


def _build_loss_function(
    config: OptimizerConfig,
    k0: jnp.ndarray,
) -> Tuple[
    Callable[
        [jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray],
        jnp.ndarray,
    ],
    Dict[str, Any],
]:
    """Build loss function and internal config from config and initial guess."""
    internal_config: Dict[str, Any] = {}
    if config.regularization is not None:
        internal_config["regularization"] = config.regularization
        internal_config = engine.config_regularization(k0, internal_config)

        def loss_fn(
            k: jnp.ndarray,
            I_obs: jnp.ndarray,
            X: jnp.ndarray,
            Y: jnp.ndarray,
            t_vals: jnp.ndarray,
        ) -> jnp.ndarray:
            return engine.loss_regularized(
                k, I_obs, X, Y, t_vals, internal_config
            )

        return loss_fn, internal_config
    return engine.loss_function, internal_config


def _run_optimization(
    I_obs: jnp.ndarray,
    X: jnp.ndarray,
    Y: jnp.ndarray,
    t_vals: jnp.ndarray,
    k0: jnp.ndarray,
    lower: jnp.ndarray,
    upper: jnp.ndarray,
    fixed_params: Dict[int, float],
    loss_fn: Callable[
        [jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray],
        jnp.ndarray,
    ],
    internal_config: Dict[str, Any],
    config: OptimizerConfig,
) -> List[Tuple[float, np.ndarray]]:
    """Run L-BFGS-B from k0 and optionally from multistart points.

    Returns a list of (loss, parameter_vector) pairs sorted by loss (best first).
    """
    debug = config.debug
    results: List[Tuple[float, np.ndarray]] = []

    result = engine.estimate_parameters_bfgs(
        I_obs=I_obs,
        X=X,
        Y=Y,
        t_vals=t_vals,
        k0=k0,
        loss_fn=loss_fn,
        bounds=(lower, upper),
        maxiter=config.max_iterations,
        verbose=config.verbose,
    )
    results.append((float(result.fun), np.asarray(result.x, dtype=np.float64)))

    if config.multistart is not None:
        if config.random_seed is not None:
            np.random.seed(config.random_seed)
        internal_config["multistart"] = config.multistart
        k_starts = engine.generate_initial_points(
            k0, lower, upper, fixed_params, internal_config
        )
        for i, k_start in enumerate(k_starts):
            if debug:
                print(f"\n--- Multistart {i + 1}/{len(k_starts)} ---")
            ms_result = engine.estimate_parameters_bfgs(
                I_obs=I_obs,
                X=X,
                Y=Y,
                t_vals=t_vals,
                k0=k_start,
                loss_fn=loss_fn,
                bounds=(lower, upper),
                maxiter=config.max_iterations,
                verbose=debug,
            )
            results.append(
                (
                    float(ms_result.fun),
                    np.asarray(ms_result.x, dtype=np.float64),
                )
            )

    results.sort(key=lambda x: x[0])
    return results


def _build_optimizer_result(
    best_k: np.ndarray,
    best_loss: float,
    results: List[Tuple[float, np.ndarray]],
    I_obs: jnp.ndarray,
    bbox: Tuple[int, int, int, int],
    X: jnp.ndarray,
    Y: jnp.ndarray,
    t_vals: jnp.ndarray,
    config: OptimizerConfig,
) -> OptimizerResult:
    """Build OptimizerResult from best solution and optional debug data."""
    return OptimizerResult(
        estimated_parameters=best_k,
        objective_value=best_loss,
        success=True,
        message="Optimization completed",
        iterations=len(results),
        all_results=results if config.debug else None,
        cropped_image=np.asarray(I_obs),
        crop_bounds=bbox,
        x_grid=X if config.debug else None,
        y_grid=Y if config.debug else None,
        t_vals=t_vals if config.debug else None,
    )


# ---------------------------------------------------------------------------
# Main entry-point
# ---------------------------------------------------------------------------


def estimate_with_optimizer(
    image: ArrayLike,
    meta: Optional[Metadata] = None,
    config: Optional[OptimizerConfig] = None,
) -> OptimizerResult:
    """Complete solver for beam parameter estimation.

    Pipeline (matches the original ``estimate_with_optimizer`` exactly):

    1. Parse inputs and metadata.
    2. Preprocess the image (bg removal, crop, L2-normalize).
    3. Build spatial and temporal grids from *cropped* dimensions.
    4. Set initial guesses and parameter bounds.
    5. Optionally configure regularisation.
    6. Run bounded L-BFGS optimisation (with optional multistart).
    7. Rank and return solutions.

    Args:
        image: Raw 2D camera image (NumPy array).
        meta: Known or measured quantities (scan frequencies, pulse
            duration, sampling rate).
        config: User choices (debug, regularisation, multistart).

    Returns:
        :class:`OptimizerResult` with the best estimated parameters.
    """
    if config is None:
        config = OptimizerConfig()
    if meta is None:
        meta = Metadata()

    # 1) Parse inputs
    image_jnp, _ = parse_input(image, meta)

    # 2–4) Preprocess, build grids, initial guess and bounds
    I_obs, bbox, X, Y, t_vals, k0, lower, upper, fixed_params = _prepare_problem(
        image_jnp, meta
    )

    # 5) Loss function (with optional regularisation)
    loss_fn, internal_config = _build_loss_function(config, k0)

    # 6) Optimization (with optional multistart)
    results = _run_optimization(
        I_obs=I_obs,
        X=X,
        Y=Y,
        t_vals=t_vals,
        k0=k0,
        lower=lower,
        upper=upper,
        fixed_params=fixed_params,
        loss_fn=loss_fn,
        internal_config=internal_config,
        config=config,
    )
    best_loss, best_k = results[0]

    # 7) Build and return result
    return _build_optimizer_result(
        best_k=best_k,
        best_loss=best_loss,
        results=results,
        I_obs=I_obs,
        bbox=bbox,
        X=X,
        Y=Y,
        t_vals=t_vals,
        config=config,
    )
