"""Shared utilities and constants for model interfaces.

Grid construction (``make_grid_configuration``) is kept here because it is
used by the MCMC interface and example scripts.  Preprocessing functions
now live in :mod:`beam_infer.engine`.
"""

import jax.numpy as jnp
import numpy as np
from typing import Tuple

# Standard parameter names (10-D model)
PARAM_NAMES_10D = [
    "Ax", "Ay", "sigx", "sigy", "cx", "cy", "fx", "fy", "phix", "phiy"
]


# Reference parameter values used by the MCMC interface and examples.
# The optimization pipeline (engine.py) derives bounds from the cropped
# image dimensions instead of using these constants.
DEFAULT_NOMINAL_PARAMS_10D = np.array(
    [60.0, 20.0, 13.5, 5.05, 0.0, 0.0, 39.55, 29.05, 0.0, 0.0],
    dtype=np.float64,
)
DEFAULT_LOWER_BOUNDS_10D = np.array(
    [0.0, 0.0, 2.0, 2.0, -20.0, -20.0, 20.0, 20.0, -np.pi, -np.pi],
    dtype=np.float64,
)
DEFAULT_UPPER_BOUNDS_10D = np.array(
    [65.0, 25.0, 20.0, 20.0, 20.0, 20.0, 50.0, 50.0, np.pi, np.pi],
    dtype=np.float64,
)


def make_grid_configuration(
    pulse_duration_ms: float = 0.05,
    sampling_rate: int = 2_000_000,
    field_x_mm: float = 400.0,
    field_y_mm: float = 350.0,
    pixel_size_mm: float = 1.0,
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, float]:
    """Create spatial and temporal grids for beam simulation.

    Args:
        pulse_duration_ms: Pulse duration in milliseconds.
        sampling_rate: Sampling rate in Hz (samples per second).
        field_x_mm: Field width in mm (x-direction).
        field_y_mm: Field height in mm (y-direction).
        pixel_size_mm: Pixel size in mm.

    Returns:
        Tuple ``(X, Y, t_vals, dt)`` where *X*, *Y* are 2-D coordinate
        grids, *t_vals* is a 1-D time array, and *dt* is the time step.
    """
    nx = int(field_x_mm / pixel_size_mm)
    ny = int(field_y_mm / pixel_size_mm)

    x = jnp.linspace(-field_x_mm / 2, field_x_mm / 2, nx)
    y = jnp.linspace(-field_y_mm / 2, field_y_mm / 2, ny)
    X, Y = jnp.meshgrid(x, y, indexing="xy")

    dt = 1.0 / sampling_rate
    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 X, Y, t_vals, dt
