# tests/test_interpolation.py
import numpy as np
import pytest

from amunpy.interpolation import interpolate


def _ghosted_shape(shape_interior, ng: int):
    return tuple(n + 2 * ng for n in shape_interior)


def _make_ghosted_with_sentinel(shape_interior, ng: int, interior_value=1.0, ghost_value=99.0):
    """Interior set to interior_value, all ghosts set to ghost_value."""
    a = np.full(_ghosted_shape(shape_interior, ng), ghost_value, dtype=float)
    core = tuple(slice(ng, ng + n) for n in shape_interior)
    a[core] = float(interior_value)
    return a


def _make_ghosted_from_func_full(shape_interior, ng: int, func):
    """
    Fill the FULL ghosted array using func evaluated on cell-center coordinates.

    Coordinates are in "index space" with cell centers at:
      x_i = (i - ng) + 0.5
    so the coarse interior runs from 0.5 .. (n-0.5).
    """
    full_shape = _ghosted_shape(shape_interior, ng)

    # cell-center coords in index space including ghosts
    coords_1d = [(np.arange(s, dtype=float) - ng + 0.5) for s in full_shape]
    grids = np.meshgrid(*coords_1d, indexing="ij")

    vals = func(*grids)
    # Ensure we always return an array of the correct shape
    out = np.asarray(vals, dtype=float)
    if out.shape == ():  # scalar -> broadcast
        out = np.full(full_shape, float(out), dtype=float)
    else:
        assert out.shape == full_shape, f"func returned shape {out.shape}, expected {full_shape}"
    return out


@pytest.mark.parametrize("ndim", [2, 3])
def test_interpolate_rebin_respects_nghosts_crop(ndim):
    ng = 2
    shape_interior = (6,) * ndim

    a = _make_ghosted_with_sentinel(shape_interior, ng, interior_value=1.0, ghost_value=99.0)
    out = interpolate(a, newshape=(3,) * ndim, nghosts=ng, method="rebin")

    assert out.shape == (3,) * ndim
    assert np.allclose(out, 1.0), "ghost sentinel leaked into rebin result"


@pytest.mark.parametrize("method", ["rebin", "tvd", "pchip"])
@pytest.mark.parametrize("ndim", [2, 3])
def test_prolongation_preserves_constant(method, ndim):
    ng = 2
    shape_interior = (4,) * ndim
    newshape = tuple(2 * n for n in shape_interior)

    if method == "rebin":
        # For rebin we *want* to be sure ghosts are ignored
        a = _make_ghosted_with_sentinel(shape_interior, ng, interior_value=3.14159, ghost_value=99.0)
    else:
        # For TVD/PCHIP, fill everything with the constant
        a = _make_ghosted_from_func_full(shape_interior, ng, lambda *c: 3.14159)

    out = interpolate(a, newshape=newshape, nghosts=ng, method=method)

    assert out.shape == newshape
    assert np.allclose(out, 3.14159)


@pytest.mark.parametrize("method", ["tvd", "pchip"])
@pytest.mark.parametrize("ndim", [2, 3])
def test_prolongation_linear_exact(method, ndim):
    """
    For a globally linear field, TVD (with limiter) should be exact
    (no limiting triggers), and monotone cubic should also be exact.
    """
    ng = 2
    shape_interior = (5,) * ndim
    newshape = tuple(2 * n for n in shape_interior)

    coeffs = np.arange(1, ndim + 1, dtype=float)  # [1,2] or [1,2,3]

    def linear(*coords):
        acc = 0.0
        for c, a in zip(coeffs, coords):
            acc = acc + c * a
        return acc

    a = _make_ghosted_from_func_full(shape_interior, ng, linear)
    out = interpolate(a, newshape=newshape, nghosts=ng, method=method)

    # Build the expected fine grid values at fine cell centers.
    # Fine interior has extent [0.5 .. n-0.5] but with twice resolution:
    # x_f = (i + 0.5) / 2  in coarse-cell units?  Our prolongation kernel is AMR-style:
    # it reconstructs inside each coarse cell with offsets t in [-0.5,0.5], so the fine
    # cell centers correspond to (coarse index + 0.5 + t).
    #
    # Easiest: just evaluate the same "index-space" coordinate convention directly
    # on the fine interior grid with spacing 0.5.
    coords_1d = [np.arange(n, dtype=float) * 0.5 + 0.25 for n in newshape]  # centers at 0.25,0.75,...
    grids = np.meshgrid(*coords_1d, indexing="ij")
    expected = linear(*grids)

    assert out.shape == newshape
    assert np.allclose(out, expected, rtol=0, atol=1e-12)


@pytest.mark.parametrize("method", ["tvd", "pchip"])
@pytest.mark.parametrize("ndim", [2, 3])
def test_prolongation_step_no_overshoot(method, ndim):
    """
    A step function should not overshoot min/max of the coarse interior.
    """
    ng = 2
    shape_interior = (6,) * ndim
    newshape = tuple(2 * n for n in shape_interior)

    # Step along the last axis at x = 3.0 (in coarse index-space coordinate)
    def step(*coords):
        x = coords[-1]
        return np.where(x < 3.0, 0.0, 1.0)

    a = _make_ghosted_from_func_full(shape_interior, ng, step)
    out = interpolate(a, newshape=newshape, nghosts=ng, method=method)

    core = tuple(slice(ng, ng + n) for n in shape_interior)
    a_core = a[core]
    vmin, vmax = float(np.min(a_core)), float(np.max(a_core))

    assert out.shape == newshape
    assert np.min(out) >= vmin - 1e-12
    assert np.max(out) <= vmax + 1e-12

