from __future__ import annotations

from collections import deque
from typing import Callable, Literal, Optional, Union

import autograd.numpy as np
import pydantic.v1 as pd
from autograd import value_and_grad
from numpy.typing import NDArray
from scipy.optimize import minimize

import tidy3d as td
from tidy3d.components.base import Tidy3dBaseModel
from tidy3d.components.grid.grid import Coords
from tidy3d.plugins.autograd.constants import BETA_DEFAULT, ETA_DEFAULT
from tidy3d.plugins.autograd.types import KernelType, PaddingType

from .filters import make_filter
from .projections import tanh_projection


class FilterAndProject(Tidy3dBaseModel):
    """A class that combines filtering and projection operations."""

    radius: Union[float, tuple[float, ...]] = pd.Field(
        ..., title="Radius", description="The radius of the kernel."
    )
    dl: Union[float, tuple[float, ...]] = pd.Field(
        ..., title="Grid Spacing", description="The grid spacing."
    )
    size_px: Union[int, tuple[int, ...]] = pd.Field(
        None, title="Size in Pixels", description="The size of the kernel in pixels."
    )
    beta: pd.NonNegativeFloat = pd.Field(
        BETA_DEFAULT, title="Beta", description="The beta parameter for the tanh projection."
    )
    eta: pd.NonNegativeFloat = pd.Field(
        ETA_DEFAULT, title="Eta", description="The eta parameter for the tanh projection."
    )
    filter_type: KernelType = pd.Field(
        "conic", title="Filter Type", description="The type of filter to create."
    )
    padding: PaddingType = pd.Field(
        "reflect", title="Padding", description="The padding mode to use."
    )

    def __call__(
        self, array: NDArray, beta: Optional[float] = None, eta: Optional[float] = None
    ) -> NDArray:
        """Apply the filter and projection to an input array.

        Parameters
        ----------
        array : np.ndarray
            The input array to filter and project.
        beta : float = None
            The beta parameter for the tanh projection. If None, uses the instance's beta.
        eta : float = None
            The eta parameter for the tanh projection. If None, uses the instance's eta.

        Returns
        -------
        np.ndarray
            The filtered and projected array.
        """
        filter_instance = make_filter(
            radius=self.radius,
            dl=self.dl,
            size_px=self.size_px,
            filter_type=self.filter_type,
            padding=self.padding,
        )
        filtered = filter_instance(array)
        beta = beta if beta is not None else self.beta
        eta = eta if eta is not None else self.eta
        projected = tanh_projection(filtered, beta, eta)
        return projected


def make_filter_and_project(
    radius: Optional[Union[float, tuple[float, ...]]] = None,
    dl: Optional[Union[float, tuple[float, ...]]] = None,
    *,
    size_px: Optional[Union[int, tuple[int, ...]]] = None,
    beta: float = BETA_DEFAULT,
    eta: float = ETA_DEFAULT,
    filter_type: KernelType = "conic",
    padding: PaddingType = "reflect",
) -> Callable:
    """Create a function that filters and projects an array.

    See Also
    --------
    :func:`~parametrizations.FilterAndProject`.
    """
    return FilterAndProject(
        radius=radius,
        dl=dl,
        size_px=size_px,
        beta=beta,
        eta=eta,
        filter_type=filter_type,
        padding=padding,
    )


def initialize_params_from_simulation(
    sim: td.Simulation,
    param_to_structure: Callable[..., td.Structure],
    params0: np.ndarray,
    *,
    freq: Optional[float] = None,
    outside_handling: Literal["extrapolate", "mask", "nan"] = "mask",
    maxiter: int = 100,
    bounds: tuple[Optional[float], Optional[float]] = (0.0, 1.0),
    rel_improve_tol: float = 1e-3,
    verbose: bool = False,
    **param_kwargs,
) -> np.ndarray:
    """Initialize design parameters to match base simulation permittivity in a region.

    Builds an objective that compares the base simulation's permittivity against
    the permittivity generated by a user‑supplied
    parameterization and minimizes the relative L2 error using L‑BFGS‑B.

    Notes
    -----
    - The simulation is not modified. If you need finer sampling inside the design
      region, add a :class:`.MeshOverrideStructure` to ``sim.grid_spec`` first.
    - The callable ``param_to_structure`` controls the parameterization and must return
      a :class:`.Structure` (typically with a :class:`.CustomMedium`). Its permittivity
      dataset (and coordinates) defines the grid used for comparison.
    - The base simulation permittivity is sampled once on an extended subgrid covering
      the design geometry and then interpolated onto the design region coordinates
      using coordinate-aware interpolation (no per-iteration interpolation).
    - Early stopping uses a single knob ``rel_improve_tol``; optimization stops once the
      relative improvement over a small fixed window falls below this value.
    - Points outside the base‑epsilon coverage can be handled via ``outside_handling``:
      ``'extrapolate'`` (use nearest extrapolation, default in earlier versions),
      ``'mask'`` (ignore outside points using a coverage mask; default), or
      ``'nan'`` (non‑extended grid; treat outside as NaN and ignore in the loss).

    Parameters
    ----------
    sim : :class:`.Simulation`
        Base simulation without the design structure.
    param_to_structure : Callable[..., :class:`.Structure`]
        Function mapping parameters to a :class:`.Structure` whose medium permittivity is used
        for comparison. Any extra keyword arguments passed to this initializer are forwarded to
        ``param_to_structure``.
    params0 : :class:`numpy.ndarray`
        Initial parameter array (2D or 3D). Values may take any range, consistent with ``bounds``
        and the expectations of ``param_to_structure``.
    maxiter : int = 100
        Maximum number of L‑BFGS‑B iterations.
    freq : float, optional
        Frequency at which permittivity is evaluated. If ``None``, uses infinite frequency.
    outside_handling : {"extrapolate", "mask", "nan"} = "mask"
        Strategy for points where design coordinates fall outside the sampled base epsilon:
        - "extrapolate": include them using nearest extrapolation.
        - "mask": include only points within the coverage bounds of ``sim.epsilon`` on the
          extended subgrid.
        - "nan": sample on a non‑extended subgrid and ignore points where interpolation returns NaN.
    bounds : tuple[float | None, float | None] = (0.0, 1.0)
        Element‑wise parameter bounds, e.g. ``(0.0, 1.0)`` or ``(-1.0, 1.0)``. Use ``None`` to
        indicate an unbounded side.
    rel_improve_tol : float = 1e-3
        Early‑stop tolerance on relative improvement over a small window.
    verbose : bool = False
        If ``True``, prints SciPy optimizer messages.

    Returns
    -------
    :class:`numpy.ndarray`
        Optimized parameters with the same shape as ``params0``.

    Examples
    --------
    Initialize a small parameter array in a 2D simulation using a minimal,
    differentiable parameterization that maps parameters to a permittivity field.

    >>> import autograd.numpy as np
    >>> import tidy3d as td
    >>> from tidy3d.plugins.autograd.invdes import initialize_params_from_simulation
    >>> sim = td.Simulation(
    ...     size=(2.0, 2.0, 0.0),
    ...     grid_spec=td.GridSpec.uniform(dl=1.0),
    ...     boundary_spec=td.BoundarySpec.pml(x=True, y=True),
    ...     run_time=1e-12,
    ... )
    >>> box = td.Box(center=(0, 0, 0), size=(1.0, 1.0, 1.0))
    >>> params0 = np.zeros((3, 3))
    >>> def param_to_structure(p):
    ...     # Map params -> density in [0, 1], then to eps in [1, 4].
    ...     density = 0.5 * (1 + np.tanh(p))
    ...     eps = 1.0 + 3.0 * density
    ...     eps3d = eps.reshape((*p.shape, 1))
    ...     return td.Structure.from_permittivity_array(geometry=box, eps_data=eps3d)
    >>> params = initialize_params_from_simulation(
    ...     sim=sim,
    ...     param_to_structure=param_to_structure,
    ...     params0=params0,
    ...     maxiter=3,
    ... )
    >>> params.shape
    (3, 3)
    """
    structure_init = param_to_structure(params0, **param_kwargs)

    if outside_handling not in ("extrapolate", "mask", "nan"):
        raise ValueError("'outside_handling' must be one of {'extrapolate', 'mask', 'nan'}.")

    # build base epsilon and interpolate onto design coords
    extend_flag = outside_handling != "nan"
    subgrid = sim.discretize(structure_init.geometry, extend=extend_flag)
    eps_base_da = sim.epsilon_on_grid(grid=subgrid, coord_key="centers", freq=freq)

    design_eps_da = structure_init.medium.permittivity
    design_coords = Coords(
        x=np.array(design_eps_da.coords["x"]),
        y=np.array(design_eps_da.coords["y"]),
        z=np.array(design_eps_da.coords["z"]),
    )

    if outside_handling == "nan":
        eps_base_interp = design_coords.spatial_interp(
            array=eps_base_da, interp_method="linear", fill_value=np.nan
        ).data
        mask = np.isfinite(eps_base_interp)
    else:
        eps_base_interp = design_coords.spatial_interp(
            array=eps_base_da, interp_method="linear", fill_value="extrapolate"
        ).data
        if outside_handling == "mask":
            # build mask from coverage bounds of base epsilon coordinates
            xb, yb, zb = [np.array(eps_base_da.coords[d]) for d in ("x", "y", "z")]
            xd, yd, zd = [np.array(design_eps_da.coords[d]) for d in ("x", "y", "z")]
            mask_x = (xd >= np.min(xb)) & (xd <= np.max(xb))
            mask_y = (yd >= np.min(yb)) & (yd <= np.max(yb))
            mask_z = (zd >= np.min(zb)) & (zd <= np.max(zb))
            mask = (mask_x[:, None, None]) & (mask_y[None, :, None]) & (mask_z[None, None, :])
        else:
            mask = None

    if mask is not None:
        covered = int(np.sum(mask))
        total = int(np.prod(mask.shape))
        frac = covered / max(total, 1)
        if frac < 0.9:
            td.log.warning(
                f"Only {frac:.1%} of design points are covered by base epsilon sampling. "
                "Consider adding a 'MeshOverrideStructure' or adjusting design coordinates."
            )

    if mask is None:
        denom = np.sqrt(np.sum(eps_base_interp.real**2 + eps_base_interp.imag**2))
    else:
        denom = np.sqrt(
            np.sum((eps_base_interp.real[mask]) ** 2 + (eps_base_interp.imag[mask]) ** 2)
        )
    denom = 1.0 if denom == 0 else denom

    def loss_fn(params_vec: np.ndarray) -> float:
        params = params_vec.reshape(params0.shape)
        structure = param_to_structure(params, **param_kwargs)
        eps_design = structure.medium.permittivity.data
        if mask is None:
            res = eps_base_interp - eps_design
            return 0.5 * np.sum(res.real**2 + res.imag**2) / denom
        res = (eps_base_interp - eps_design)[mask]
        return 0.5 * np.sum(res.real**2 + res.imag**2) / denom

    val_and_grad = value_and_grad(loss_fn)

    WINDOW = 3
    state = {
        "vals": deque(maxlen=WINDOW),
        "best_val": np.inf,
        "best_x": params0.ravel().copy(),
    }

    def callback(xk: np.ndarray):
        val = loss_fn(xk)
        if val < state["best_val"]:
            state["best_val"] = val
            state["best_x"] = xk.copy()
        state["vals"].append(val)
        if len(state["vals"]) == WINDOW:
            v_then = state["vals"][0]
            v_now = state["vals"][-1]
            rel_improve = (v_then - v_now) / max(abs(v_then), 1e-12)
            if rel_improve < rel_improve_tol:
                raise StopIteration

    bounds_list = [bounds] * params0.size
    try:
        res = minimize(
            fun=val_and_grad,
            x0=params0.ravel(),
            method="L-BFGS-B",
            jac=True,
            bounds=bounds_list,
            callback=callback,
            options={"maxiter": maxiter, "disp": verbose},
        )
        x_final = res.x
    except StopIteration:
        x_final = state["best_x"]
    return x_final.reshape(params0.shape)
