"""
AmunPy: AMR interpolation utilities (NumPy-only)
===============================================

This module provides internal interpolation operators used by AmunPy to
assemble uniform-resolution datasets from AMR blocks.

AmunPy contract
---------------
- Input blocks may include ghost zones of width `nghosts` on every axis.
- The output of `interpolate()` is always **interior-only** (no ghosts),
  with shape exactly `newshape`.
- Higher-order prolongation methods (TVD / monotone cubic PCHIP-style)
  use ghost zones as stencil support. Therefore, ghost values must be
  physically consistent (as stored in valid AMUN snapshots).

Supported methods
-----------------
- method=None or "rebin"
    Conservative restriction (block averaging) and zeroth-order
    prolongation (nearest-neighbor replication). Does not require ghosts.

- "tvd"
    TVD piecewise-linear prolongation with slope limiting (MC/minmod).
    Restriction is still conservative via "rebin".
    Requires `nghosts >= 1` (stencil support).

- "pchip" / "monotonic" / "mono3"
    Monotonicity-preserving cubic (PCHIP-style) prolongation based on
    cubic Hermite interpolation with slope limiting.
    Restriction is still conservative via "rebin".
    `nghosts >= 2` is recommended (ng=1 works but boundary behavior is
    slightly more diffusive).

Notes on backwards compatibility
--------------------------------
- SciPy-based interpolation backends have been removed.
  Interpolation in AmunPy is now fully NumPy-based and no longer supports
  order-driven interpolation parameters.
"""

from __future__ import annotations

import numpy as np

__all__ = ["interpolate"]


# -----------------------------------------------------------------------------
# Core helpers
# -----------------------------------------------------------------------------

def _as_shape(shape) -> tuple[int, ...]:
    return tuple(int(n) for n in shape)


def _core_view(a: np.ndarray, ng: int) -> np.ndarray:
    """Return a view of the interior region (strip `ng` ghost cells on all axes)."""
    if ng == 0:
        return a
    if ng < 0:
        raise ValueError(f"interpolate(): nghosts must be >= 0, got {ng}")
    for s in a.shape:
        if s <= 2 * ng:
            raise ValueError(f"interpolate(): nghosts={ng} too large for shape {a.shape}")
    sl = tuple(slice(ng, -ng) for _ in range(a.ndim))
    return a[sl]


def _strip_ghosts_along_axis(a: np.ndarray, ng: int, axis: int) -> np.ndarray:
    """Strip ghosts only along a single axis (safe when ng == 0)."""
    if ng == 0:
        return a
    sl = [slice(None)] * a.ndim
    sl[axis] = slice(ng, -ng)
    return a[tuple(sl)]


# -----------------------------------------------------------------------------
# Conservative AMR resize (restriction/prolongation)
# -----------------------------------------------------------------------------

def _rebin(a: np.ndarray, newshape) -> np.ndarray:
    """
    Resize a cell-centered array to `newshape` using AMR-style integer ratios:
      - restriction (shrinking): block average (mean)
      - prolongation (growing): nearest-neighbor repeat

    Requirements:
      - a.ndim == len(newshape)
      - per-axis sizes must change by an integer factor
    """
    newshape = _as_shape(newshape)
    if a.ndim != len(newshape):
        raise ValueError(f"_rebin(): ndim mismatch: a.ndim={a.ndim}, newshape={newshape}")

    out = a
    for ax, n_new in enumerate(newshape):
        n_old = out.shape[ax]
        if n_new == n_old:
            continue

        if n_new > n_old:
            if n_new % n_old != 0:
                raise ValueError(f"_rebin(): axis {ax}: new ({n_new}) not divisible by old ({n_old})")
            r = n_new // n_old
            out = np.repeat(out, r, axis=ax)
        else:
            if n_old % n_new != 0:
                raise ValueError(f"_rebin(): axis {ax}: old ({n_old}) not divisible by new ({n_new})")
            r = n_old // n_new
            # reshape to (..., n_new, r, ...) and average over the r-axis
            shp = list(out.shape)
            shp[ax] = n_new
            shp.insert(ax + 1, r)
            out = out.reshape(shp).mean(axis=ax + 1)

    return out


# -----------------------------------------------------------------------------
# TVD linear prolongation
# -----------------------------------------------------------------------------

def _minmod(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """Elementwise minmod limiter."""
    return 0.5 * (np.sign(a) + np.sign(b)) * np.minimum(np.abs(a), np.abs(b))


def _mc_limiter(dl: np.ndarray, dr: np.ndarray) -> np.ndarray:
    """Monotonized central (MC) limiter: minmod(0.5*(dl+dr), 2*dl, 2*dr)."""
    return _minmod(_minmod(0.5 * (dl + dr), 2.0 * dl), 2.0 * dr)


def _tvd_prolong_axis(q: np.ndarray, n_fine: int, ng: int, axis: int = 0, limiter: str = "mc") -> np.ndarray:
    """
    TVD linear AMR prolongation along one axis.

    Contract:
      - q includes ghosts of width ng along `axis`
      - output is interior-only along `axis` with length n_fine
      - other axes are preserved (may still include ghosts)

    Requirements:
      - ng >= 1 (stencil)
      - integer refinement ratio: n_fine % (n_tot - 2*ng) == 0
    """
    if ng < 1:
        raise ValueError("tvd prolongation requires nghosts >= 1")

    q = np.asarray(q)
    n_tot = q.shape[axis]
    n_coarse = n_tot - 2 * ng
    if n_coarse <= 0:
        raise ValueError(f"invalid nghosts={ng} for axis size {n_tot}")

    if n_fine % n_coarse != 0:
        raise ValueError(
            f"tvd prolongation requires integer ratio: n_fine={n_fine} not divisible by "
            f"coarse interior n_coarse={n_coarse} (axis={axis})"
        )
    r = n_fine // n_coarse
    if r < 1:
        raise ValueError("tvd prolongation: refinement ratio must be >= 1")

    # Move target axis to front: shape (n_tot, ...)
    if axis != 0:
        perm = (axis,) + tuple(i for i in range(q.ndim) if i != axis)
        inv = np.argsort(perm)
        q0 = np.transpose(q, perm)
    else:
        inv = None
        q0 = q

    # Slopes using neighbor differences; defined for i=1..n_tot-2
    dl = q0[1:-1] - q0[:-2]
    dr = q0[2:] - q0[1:-1]

    if limiter == "mc":
        s_inner = _mc_limiter(dl, dr)
    elif limiter == "minmod":
        s_inner = _minmod(dl, dr)
    else:
        raise ValueError(f"unknown limiter '{limiter}' (use 'mc' or 'minmod')")

    s = np.zeros_like(q0)
    s[1:-1] = s_inner

    # Coarse interior cells
    qc = q0[ng : ng + n_coarse]
    sc = s[ng : ng + n_coarse]

    # Fine offsets inside each coarse cell: t_k in [-0.5, 0.5]
    t = ((np.arange(r) + 0.5) / r) - 0.5  # (r,)
    t = t.reshape((1, r) + (1,) * (q0.ndim - 1))  # broadcast-safe for 2D/3D

    # Reconstruct then reshape to (n_fine, ...)
    qf = (qc[:, None, ...] + sc[:, None, ...] * t).reshape((n_coarse * r,) + q0.shape[1:])

    # Restore original axis ordering
    if inv is not None:
        qf = np.transpose(qf, inv)

    return qf


# -----------------------------------------------------------------------------
# Monotone cubic (PCHIP-style) prolongation
# -----------------------------------------------------------------------------

def _pchip_slopes_uniform(y: np.ndarray) -> np.ndarray:
    """
    Vectorized monotone cubic (PCHIP-style) slope construction on a uniform grid.

    y: array of shape (n, ...)
    returns m: slopes with same shape as y
    """
    y = np.asarray(y)
    n = y.shape[0]
    m = np.zeros_like(y)
    if n < 2:
        return m

    d = y[1:] - y[:-1]  # (n-1,...)

    if n == 2:
        m[0] = d[0]
        m[1] = d[0]
        return m

    d0 = d[:-1]
    d1 = d[1:]
    same_sign = (d0 * d1) > 0

    # Interior slopes: weighted harmonic mean where monotone
    if np.any(same_sign):
        w1 = 2.0 * d1 + d0
        w2 = d1 + 2.0 * d0

        denom = np.empty_like(d0)
        denom[same_sign] = (w1[same_sign] / d0[same_sign]) + (w2[same_sign] / d1[same_sign])

        mi = np.empty_like(d0)
        mi[same_sign] = (w1[same_sign] + w2[same_sign]) / denom[same_sign]

        m[1:-1][same_sign] = mi[same_sign]

    # Endpoint slopes + monotonicity clamping (Fritsch–Carlson style)
    m0 = 2.0 * d[0] - d[1]
    mn = 2.0 * d[-1] - d[-2]

    def _clamp_end(mend, dend):
        mend = np.where(mend * dend <= 0, 0.0, mend)
        mend = np.where(np.abs(mend) > 3.0 * np.abs(dend), 3.0 * dend, mend)
        return mend

    m[0] = _clamp_end(m0, d[0])
    m[-1] = _clamp_end(mn, d[-1])

    return m


def _mono3_prolong_axis(q: np.ndarray, n_fine: int, ng: int, axis: int = 0) -> np.ndarray:
    """
    Monotone cubic prolongation along one axis (PCHIP-style, uniform grid).

    Contract:
      - q includes ghosts of width ng along `axis`
      - output is interior-only along `axis` with length n_fine
      - integer refinement ratio required

    Recommended: ng >= 2 (ng>=1 works).
    """
    q = np.asarray(q)
    n_tot = q.shape[axis]
    n_coarse = n_tot - 2 * ng
    if n_coarse <= 0:
        raise ValueError(f"invalid nghosts={ng} for axis size {n_tot}")

    if n_fine % n_coarse != 0:
        raise ValueError(
            f"monotone cubic prolongation requires integer ratio: n_fine={n_fine} "
            f"not divisible by coarse interior n_coarse={n_coarse} (axis={axis})"
        )
    r = n_fine // n_coarse
    if r < 1:
        raise ValueError("refinement ratio must be >= 1")

    # Move axis to front
    if axis != 0:
        perm = (axis,) + tuple(i for i in range(q.ndim) if i != axis)
        inv = np.argsort(perm)
        q0 = np.transpose(q, perm)
    else:
        inv = None
        q0 = q

    m = _pchip_slopes_uniform(q0)

    # Fine offsets inside a coarse cell
    tcell = ((np.arange(r) + 0.5) / r) - 0.5  # (r,)
    ic = np.arange(ng, ng + n_coarse)  # coarse interior indices in full array
    x_new = (ic[:, None] + tcell[None, :]).reshape(-1)  # (n_fine,)

    # Interval index j such that x_new in [j, j+1]
    j = np.floor(x_new).astype(np.int64)
    j = np.clip(j, 0, n_tot - 2)
    t = (x_new - j).astype(q0.dtype)

    # Hermite basis (dx = 1)
    t2 = t * t
    t3 = t2 * t
    h00 = (2.0 * t3 - 3.0 * t2 + 1.0)
    h10 = (t3 - 2.0 * t2 + t)
    h01 = (-2.0 * t3 + 3.0 * t2)
    h11 = (t3 - t2)

    bshape = (t.shape[0],) + (1,) * (q0.ndim - 1)
    h00 = h00.reshape(bshape)
    h10 = h10.reshape(bshape)
    h01 = h01.reshape(bshape)
    h11 = h11.reshape(bshape)

    y0 = np.take(q0, j, axis=0)
    y1 = np.take(q0, j + 1, axis=0)
    m0 = np.take(m, j, axis=0)
    m1 = np.take(m, j + 1, axis=0)

    qf = h00 * y0 + h10 * m0 + h01 * y1 + h11 * m1

    if inv is not None:
        qf = np.transpose(qf, inv)
    return qf


# -----------------------------------------------------------------------------
# Public API
# -----------------------------------------------------------------------------

def interpolate(
    a: np.ndarray,
    newshape,
    nghosts: int = 0,
    method: str | None = None,
    limiter: str = "mc",
) -> np.ndarray:
    """
    Resample an AMR block to an interior-only array of shape `newshape`.

    Parameters
    ----------
    a : ndarray
        Input array, possibly including ghost zones.
    newshape : sequence[int]
        Desired interior-only output shape (no ghosts).
    nghosts : int
        Ghost-zone width in `a`.
    method : {None, "rebin", "tvd", "pchip", "monotonic", "mono3"}
        Interpolation method. None defaults to "rebin" for backwards
        compatibility.
    limiter : {"mc","minmod"}
        Limiter for TVD prolongation.

    Returns
    -------
    ndarray
        Interior-only array of shape `newshape`.
    """
    a = np.asarray(a)
    newshape = np.asarray(newshape, dtype=int)
    target_ndim = int(len(newshape))

    # Allow extra singleton axes (e.g., 2D stored as (1,ny,nx))
    if a.ndim != target_ndim:
        if a.ndim > target_ndim:
            while a.ndim > target_ndim:
                singleton_axes = [i for i, n in enumerate(a.shape) if n == 1]
                if not singleton_axes:
                    raise ValueError(
                        f"interpolate(): input rank {a.ndim} cannot be reduced to {target_ndim} "
                        f"without losing non-singleton information; shape={a.shape}, newshape={tuple(newshape)}"
                    )
                a = np.squeeze(a, axis=singleton_axes[0])
        else:
            raise ValueError(
                f"interpolate(): input rank {a.ndim} < target rank {target_ndim}; "
                f"shape={a.shape}, newshape={tuple(newshape)}"
            )

    ng = int(nghosts)
    if ng < 0:
        raise ValueError(f"interpolate(): nghosts must be >= 0, got {ng}")

    # Normalize method name (keep permissive behavior: unknown -> rebin)
    if method is None:
        meth = "rebin"
    else:
        m = method.lower()
        if m in ("rebin",):
            meth = "rebin"
        elif m in ("tvd",):
            meth = "tvd"
        elif m in ("pchip", "monotonic", "mono3"):
            meth = "pchip"
        else:
            meth = "rebin"

    # Conservative restriction if any axis shrinks (independent of method)
    coarse_interior = np.array([s - 2 * ng for s in a.shape], dtype=int)
    if np.any(newshape < coarse_interior):
        return _rebin(_core_view(a, ng), newshape)

    # Base conservative / nearest method
    if meth == "rebin":
        return _rebin(_core_view(a, ng), newshape)

    # Higher-order prolongation paths
    q = a
    if meth == "tvd":
        if ng < 1:
            raise ValueError("interpolate(tvd): nghosts must be >= 1")
        for axis in range(q.ndim):
            n_coarse = q.shape[axis] - 2 * ng
            n_target = int(newshape[axis])

            if n_target == n_coarse:
                q = _strip_ghosts_along_axis(q, ng, axis)
            elif n_target > n_coarse:
                q = _tvd_prolong_axis(q, n_target, ng, axis=axis, limiter=limiter)
            else:
                raise ValueError("interpolate(tvd): unexpected shrink along an axis (should be handled by rebin)")
        if tuple(q.shape) != tuple(newshape):
            raise RuntimeError(f"interpolate(tvd): shape mismatch: got {q.shape}, expected {tuple(newshape)}")
        return q

    # meth == "pchip"
    for axis in range(q.ndim):
        n_coarse = q.shape[axis] - 2 * ng
        n_target = int(newshape[axis])

        if n_target == n_coarse:
            q = _strip_ghosts_along_axis(q, ng, axis)
        elif n_target > n_coarse:
            q = _mono3_prolong_axis(q, n_target, ng, axis=axis)
        else:
            raise ValueError("interpolate(pchip): unexpected shrink along an axis (should be handled by rebin)")

    if tuple(q.shape) != tuple(newshape):
        raise RuntimeError(f"interpolate(pchip): shape mismatch: got {q.shape}, expected {tuple(newshape)}")
    return q
