"""Engine layer: all numerical steps for beam parameter estimation.

This module contains every helper used by estimate_with_optimizer():
preprocessing, grid construction, parameter initialization, loss functions,
multistart point generation, and the L-BFGS-B optimizer wrapper.

The implementation matches the original ``ess-beam-imaging-inverse-problems``
codebase exactly, serving as the single source of numerical truth.
"""

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

import jax.numpy as jnp
import numpy as np
from jax import grad
from scipy.ndimage import gaussian_filter
from scipy.optimize import Bounds, minimize

from beam_infer.base_model import BeamForwardModel

# Module-level forward model instance (stateless, used by loss functions).
_forward_model = BeamForwardModel()


# ---------------------------------------------------------------------------
# Preprocessing
# ---------------------------------------------------------------------------


def remove_background(image: np.ndarray, sigma: float = 30.0) -> np.ndarray:
    """Remove smooth background via Gaussian-filter subtraction.

    Args:
        image: 2D image array.
        sigma: Gaussian kernel standard deviation.

    Returns:
        Background-subtracted image, clipped to non-negative values.
    """
    bg = gaussian_filter(image.astype(np.float64), sigma=sigma)
    return np.clip(image - bg, 0, None)


def tight_crop_bbox(
    img: np.ndarray,
    frac: float = 0.02,
    pad: int = 10,
    min_h: int = 32,
    min_w: int = 32,
) -> Tuple[int, int, int, int]:
    """Find tight bounding box around the signal region.

    Algorithm:
        1. Threshold at ``frac * max_intensity``.
        2. Find rows/columns with any pixel above threshold.
        3. Add ``pad`` pixels on each side.
        4. Enforce minimum dimensions ``min_h x min_w``.

    Args:
        img: 2D image (background-subtracted).
        frac: Fraction of max intensity used as threshold.
        pad: Padding in pixels around detected region.
        min_h: Minimum crop height.
        min_w: Minimum crop width.

    Returns:
        Tuple ``(y1, y2, x1, x2)`` defining the crop region.
    """
    H, W = img.shape
    m = float(np.max(img))
    if m <= 0:
        return 0, H, 0, W

    thr = frac * m
    mask = img > thr
    if not np.any(mask):
        return 0, H, 0, W

    ys = np.where(np.any(mask, axis=1))[0]
    xs = np.where(np.any(mask, axis=0))[0]

    y1, y2 = int(ys[0]), int(ys[-1] + 1)
    x1, x2 = int(xs[0]), int(xs[-1] + 1)

    y1 = max(0, y1 - pad)
    y2 = min(H, y2 + pad)
    x1 = max(0, x1 - pad)
    x2 = min(W, x2 + pad)

    if (y2 - y1) < min_h:
        c = (y1 + y2) // 2
        y1 = max(0, c - min_h // 2)
        y2 = min(H, y1 + min_h)
        y1 = max(0, y2 - min_h)

    if (x2 - x1) < min_w:
        c = (x1 + x2) // 2
        x1 = max(0, c - min_w // 2)
        x2 = min(W, x1 + min_w)
        x1 = max(0, x2 - min_w)

    return y1, y2, x1, x2


def crop_image(
    img: np.ndarray, bbox: Tuple[int, int, int, int]
) -> np.ndarray:
    """Crop image to bounding box.

    Args:
        img: 2D image.
        bbox: Tuple ``(y1, y2, x1, x2)``.

    Returns:
        Cropped sub-image.
    """
    y1, y2, x1, x2 = bbox
    return img[y1:y2, x1:x2]


def preprocess_image(
    image: np.ndarray | jnp.ndarray,
) -> Tuple[jnp.ndarray, Tuple[int, int, int, int]]:
    """Full preprocessing pipeline: background removal, crop, L2-normalize.

    Matches the original ``_preprocess_image()`` exactly:
        1. Gaussian background subtraction (sigma=30).
        2. Tight crop (2 % threshold, 10 px padding, min 32x32).
        3. L2 normalization.

    Args:
        image: Raw 2D image as JAX or NumPy array.

    Returns:
        Tuple of ``(preprocessed JAX image, bounding box)``.
    """
    I_np = np.asarray(image)

    I_clean = remove_background(I_np)
    bbox = tight_crop_bbox(I_clean, frac=0.02, pad=10, min_h=32, min_w=32)
    I_crop = crop_image(I_clean, bbox)

    I_crop = I_crop / (np.linalg.norm(I_crop) + 1e-12)

    I_obs = jnp.asarray(I_crop, dtype=jnp.float32)
    return I_obs, bbox


# ---------------------------------------------------------------------------
# Grid construction
# ---------------------------------------------------------------------------


def build_grid_and_time(
    H: int,
    W: int,
    pulse_duration_ms: float,
    sampling_rate: int = 2_000_000,
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
    """Build spatial and temporal grids from cropped image dimensions.

    Spatial grids are centered at origin with ``pixel_size = 1``::

        X in [-W/2, W/2],  Y in [-H/2, H/2]

    This matches the original ``_build_grid_and_time()`` exactly.

    Args:
        H: Cropped image height (pixels).
        W: Cropped image width (pixels).
        pulse_duration_ms: Pulse duration in milliseconds.
        sampling_rate: Temporal sampling rate in Hz (default 2 MHz).

    Returns:
        Tuple ``(X, Y, t_vals)`` of JAX float32 arrays.
    """
    x = jnp.linspace(-W / 2.0, W / 2.0, W)
    y = jnp.linspace(-H / 2.0, H / 2.0, H)
    X, Y = jnp.meshgrid(x, y, indexing="xy")

    pulse_s = pulse_duration_ms / 1000.0
    n_time = int(pulse_s * sampling_rate)
    t_vals = jnp.linspace(0.0, pulse_s, n_time)

    return (
        jnp.asarray(X, dtype=jnp.float32),
        jnp.asarray(Y, dtype=jnp.float32),
        jnp.asarray(t_vals, dtype=jnp.float32),
    )


# ---------------------------------------------------------------------------
# Initial guess and bounds
# ---------------------------------------------------------------------------


def init_guess_and_bounds(
    H: int,
    W: int,
    meta: Optional[Dict[str, float]] = None,
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, Dict[int, float]]:
    """Initial guess and bounds derived from cropped image dimensions.

    Matches the original ``_init_guess_and_bounds()``::

        Ax0 = 0.7 * W / 2,  Ay0 = 0.7 * H / 2
        sigx0 = sigy0 = 5.0,  cx0 = cy0 = 0.0
        Bounds scale with cropped image size.

    Args:
        H: Cropped image height.
        W: Cropped image width.
        meta: Dict with optional ``'fx'``, ``'fy'`` keys.

    Returns:
        Tuple ``(k0, lower, upper, fixed_params)``.
    """
    meta = meta or {}

    Ax0 = 0.7 * W / 2
    Ay0 = 0.7 * H / 2

    fx0 = meta.get("fx", 39.55)
    fy0 = meta.get("fy", 29.05)

    k0 = np.array(
        [Ax0, Ay0, 5.0, 5.0, 0.0, 0.0, fx0, fy0, 0.0, 0.0],
        dtype=np.float32,
    )

    lower = np.array(
        [
            0.0, 0.0,
            1.0, 1.0,
            -W / 2.0, -H / 2.0,
            0.9 * fx0, 0.9 * fy0,
            -np.pi, -np.pi,
        ],
        dtype=np.float32,
    )
    upper = np.array(
        [
            float(W), float(H),
            10.0, 10.0,
            W / 2.0, H / 2.0,
            1.1 * fx0, 1.1 * fy0,
            np.pi, np.pi,
        ],
        dtype=np.float32,
    )

    fixed_params: Dict[int, float] = {}
    if "fx" in meta:
        lower[6] = upper[6] = fx0
        fixed_params[6] = fx0
    if "fy" in meta:
        lower[7] = upper[7] = fy0
        fixed_params[7] = fy0

    return (
        jnp.asarray(k0),
        jnp.asarray(lower),
        jnp.asarray(upper),
        fixed_params,
    )


# ---------------------------------------------------------------------------
# Loss functions
# ---------------------------------------------------------------------------


def loss_function(
    k: jnp.ndarray,
    I_obs: jnp.ndarray,
    X: jnp.ndarray,
    Y: jnp.ndarray,
    t_vals: jnp.ndarray,
) -> jnp.ndarray:
    """Loss comparing simulated and observed images.

    Both images are L2-normalized before comparison so that only their
    *shapes* are compared, not their absolute scales.  This matches the
    original ``loss_function()`` exactly.

    Args:
        k: Parameter vector (10-D).
        I_obs: Observed image (already L2-normalized by preprocessing).
        X, Y: 2D coordinate grids.
        t_vals: 1D time values.

    Returns:
        Scalar loss value.
    """
    I_sim = _forward_model.simulate_image(X, Y, t_vals, k)

    I_sim = I_sim / jnp.linalg.norm(I_sim)
    I_obs_n = I_obs / jnp.linalg.norm(I_obs)

    return jnp.sum((I_sim - I_obs_n) ** 2)


def config_regularization(
    k0: jnp.ndarray, config: Dict[str, Any]
) -> Dict[str, Any]:
    """Convert a high-level regularization spec into arrays for the loss.

    ``config["regularization"]["params"]`` maps parameter names to
    per-parameter penalty strengths.  Reference values come from ``k0``.

    Args:
        k0: Initial guess (reference values for the penalty).
        config: Dict with ``'regularization'`` key.

    Returns:
        Updated config with ``'reg_mask'``, ``'reg_ref'``, ``'reg_lambda'``.
    """
    reg_cfg = config.get("regularization", None)
    if reg_cfg is None:
        return config

    param_index = {
        "Ax": 0, "Ay": 1, "sigx": 2, "sigy": 3,
        "cx": 4, "cy": 5, "fx": 6, "fy": 7,
        "phix": 8, "phiy": 9,
    }

    D = k0.shape[0]
    reg_mask = np.zeros(D, dtype=np.float32)
    reg_ref = np.zeros(D, dtype=np.float32)
    reg_lambda = np.zeros(D, dtype=np.float32)

    params = reg_cfg.get("params", {})
    for name, lam in params.items():
        if name not in param_index:
            raise ValueError(f"Unknown regularized parameter '{name}'")
        idx = param_index[name]
        reg_mask[idx] = 1.0
        reg_ref[idx] = float(k0[idx])
        reg_lambda[idx] = float(lam)

    config = dict(config)
    config["reg_mask"] = reg_mask
    config["reg_ref"] = reg_ref
    config["reg_lambda"] = reg_lambda
    return config


def loss_regularized(
    k: jnp.ndarray,
    I_obs: jnp.ndarray,
    X: jnp.ndarray,
    Y: jnp.ndarray,
    t_vals: jnp.ndarray,
    config: Dict[str, Any],
) -> jnp.ndarray:
    """Loss = data mismatch + quadratic regularization around *k0*.

    Penalty term is ``reg_lambda * (k - k_ref)^2``, matching the original
    ``loss_regularized()``.

    Args:
        k: Parameter vector.
        I_obs: Observed image.
        X, Y: Coordinate grids.
        t_vals: Time values.
        config: Dict with ``'reg_mask'``, ``'reg_ref'``, ``'reg_lambda'``.

    Returns:
        Scalar loss value.
    """
    L_data = loss_function(k, I_obs, X, Y, t_vals)

    reg_mask = config.get("reg_mask", None)
    if reg_mask is None:
        return L_data

    reg_ref = jnp.asarray(config["reg_ref"])
    reg_mask_jnp = jnp.asarray(config["reg_mask"])
    reg_lambda = jnp.asarray(config["reg_lambda"])

    diff = (k - reg_ref) * reg_mask_jnp
    L_reg = jnp.sum(reg_lambda * diff**2)

    return L_data + L_reg


# ---------------------------------------------------------------------------
# Multistart point generation
# ---------------------------------------------------------------------------


def generate_initial_points(
    k0: jnp.ndarray,
    lower: jnp.ndarray,
    upper: jnp.ndarray,
    fixed_params: Dict[int, float],
    config: Dict[str, Any],
) -> List[jnp.ndarray]:
    """Generate multiple starting points for multi-start optimization.

    Fixed and regularized parameters are never randomized.

    Args:
        k0: Base initial guess.
        lower: Lower bounds.
        upper: Upper bounds.
        fixed_params: Dict mapping parameter index -> fixed value.
        config: Dict with optional ``'multistart'`` and ``'reg_mask'`` keys.

    Returns:
        List of initial parameter vectors.
    """
    ms_cfg = config.get("multistart", None)
    if ms_cfg is None:
        return [k0]

    n = int(ms_cfg.get("n_starts", 1))
    mode = ms_cfg.get("type", "uniform")
    D = k0.shape[0]

    fixed_idx = set(fixed_params.keys())
    reg_mask = config.get("reg_mask", None)
    reg_idx = set(np.where(reg_mask)[0]) if reg_mask is not None else set()
    free_idx = [i for i in range(D) if i not in fixed_idx and i not in reg_idx]

    k0_np = np.asarray(k0)
    lower_np = np.asarray(lower)
    upper_np = np.asarray(upper)

    starts: List[jnp.ndarray] = []

    if mode == "uniform":
        for _ in range(n):
            k = k0_np.copy()
            for i in free_idx:
                k[i] = lower_np[i] + (upper_np[i] - lower_np[i]) * np.random.rand()
            starts.append(jnp.asarray(k, dtype=jnp.float32))

    elif mode == "gaussian":
        std_scale = 0.5
        sigma = std_scale * (upper_np - lower_np)
        for _ in range(n):
            k = k0_np.copy()
            for i in free_idx:
                k[i] = k0_np[i] + sigma[i] * np.random.randn()
                k[i] = np.clip(k[i], lower_np[i], upper_np[i])
            starts.append(jnp.asarray(k, dtype=jnp.float32))

    else:
        raise ValueError(f"Unknown multistart type '{mode}'")

    return starts


# ---------------------------------------------------------------------------
# L-BFGS-B optimizer wrapper
# ---------------------------------------------------------------------------


def estimate_parameters_bfgs(
    I_obs: jnp.ndarray,
    X: jnp.ndarray,
    Y: jnp.ndarray,
    t_vals: jnp.ndarray,
    k0: jnp.ndarray,
    loss_fn: Callable = loss_function,
    bounds: Optional[Tuple[jnp.ndarray, jnp.ndarray]] = None,
    maxiter: int = 100,
    verbose: bool = True,
) -> Any:
    """Run L-BFGS-B optimization.

    Matches the original ``estimate_parameters_BFGS()`` signature and
    behavior.

    Args:
        I_obs: Observed image.
        X, Y: Coordinate grids.
        t_vals: Time values.
        k0: Initial guess.
        loss_fn: Loss callable ``(k, I_obs, X, Y, t_vals) -> scalar``.
        bounds: Tuple ``(lower, upper)`` arrays, or ``None``.
        maxiter: Maximum L-BFGS-B iterations.
        verbose: Print iteration details.

    Returns:
        ``scipy.optimize.OptimizeResult`` from ``minimize``.
    """
    if verbose:
        print("=== Starting parameter estimation ===")

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

    loss_grad = grad(loss)
    scipy_bounds = Bounds(*bounds) if bounds is not None else None
    iteration: Dict[str, int] = {"count": 0}

    def callback(k: np.ndarray) -> None:
        if not verbose:
            return
        iteration["count"] += 1
        l_val = loss(k)
        g = loss_grad(k)
        g_norm = jnp.linalg.norm(g)
        print(f"--- Iteration {iteration['count']} ---")
        print(f"Loss:         {l_val:.6e}")
        print(f"||grad||:     {g_norm:.6e}")
        print(f"Current k:    {np.array2string(np.asarray(k), precision=3)}")
        print(f"Gradient:     {np.array2string(np.asarray(g), precision=3)}\n")

    return minimize(
        fun=loss,
        x0=jnp.array(k0),
        jac=loss_grad,
        method="L-BFGS-B",
        bounds=scipy_bounds,
        callback=callback,
        options={"maxiter": maxiter},
    )
