# tests/test_uniform_grid_correctness.py

import numpy as np
import pytest

import amunpy.amun as amun_mod  # to monkeypatch amun_mod.interpolate
from amunpy.amun import Amun
from amunpy.exceptions import AmunPyDatasetError


# -----------------------------------------------------------------------------
# Minimal FakeAmun for uniform-grid correctness tests
# -----------------------------------------------------------------------------
class FakeAmun(Amun):
    """
    Geometry-only fake snapshot for testing uniform-grid assembly.

    - Provides controlled chunk/block metadata (coords, levels, dims).
    - Provides deterministic raw datasets via __read_binary_data__.
    - Uses the real dataset() assembly pipeline from the base class.
    """

    def __init__(
        self,
        ndims: int,
        *,
        xblocks: int,
        yblocks: int,
        zblocks: int,
        toplev: int,
        maxlev: int,
        ncells: int = 4,
        nghosts: int = 1,
        chunks=None,
        raw_values=None,
    ):
        self.path = "fake"
        self.dataformat = "fake"
        self.initialized = True
        self.chunkname = "chunk_{:04d}"

        # Domain lengths are chosen consistent with block counts:
        # each level-1 block has unit length in each direction.
        xmin, ymin, zmin = 0.0, 0.0, 0.0
        xmax, ymax, zmax = float(xblocks), float(yblocks), float(zblocks)

        self.attributes = {
            "ndims": ndims,
            "ncells": int(ncells),
            "nghosts": int(nghosts),
            "toplev": int(toplev),
            "maxlev": int(maxlev),
            "nchunks": 1,
            "nleafs": 0,  # optional
            "xblocks": int(xblocks),
            "yblocks": int(yblocks),
            "zblocks": int(zblocks),
            "xmin": xmin,
            "xmax": xmax,
            "ymin": ymin,
            "ymax": ymax,
            "zmin": zmin,
            "zmax": zmax,
        }

        # Raw variable catalog: raw IDs exist in snapshots
        self.raw_variables = {"dens": "dens"}

        # User-facing alias catalog (Commit 3 behavior)
        self.variables = {"density": "dens"}

        # Chunk metadata
        # chunks[0]["dims"] must be (dblocks, ...) and include ghost zones.
        if chunks is None:
            raise ValueError("FakeAmun requires explicit chunks metadata.")

        self.chunks = {0: chunks}

        # Per-block constants for raw datasets, indexed by p (block index)
        # e.g. raw_values["dens"] = [1.0, 2.0, ...]
        if raw_values is None:
            raise ValueError("FakeAmun requires explicit raw_values.")
        self._raw_values = raw_values

    def __read_binary_data__(self, dataset_id, chunk_number):
        # Provide (dblocks, ...) array, with ghost zones included.
        if dataset_id not in self._raw_values:
            raise RuntimeError(f"unexpected dataset_id={dataset_id!r}")

        dims = self.chunks[chunk_number]["dims"]
        dblocks = dims[0]
        ng = int(self.attributes["nghosts"])
        nc = int(self.attributes["ncells"])
        nd = int(self.attributes["ndims"])

        out = np.zeros(dims, dtype=float)

        # Interior slice in block-local array order:
        # 2D: (y, x), 3D: (z, y, x) in your codebase.
        if nd == 2:
            interior = (slice(ng, ng + nc), slice(ng, ng + nc))
        else:
            interior = (slice(ng, ng + nc), slice(ng, ng + nc), slice(ng, ng + nc))

        vals = self._raw_values[dataset_id]
        if len(vals) != dblocks:
            raise RuntimeError("raw_values length must match dblocks")

        for p in range(dblocks):
            out[(p,) + interior] = float(vals[p])

        return out


# -----------------------------------------------------------------------------
# A spy interpolate that makes geometry tests deterministic.
#
# It returns an array of shape `newshape` filled with the mean of the interior
# (ghost-stripped) input block. This preserves constant-block behavior and
# is independent of the real interpolation implementation details.
# -----------------------------------------------------------------------------
class InterpSpy:
    def __init__(self):
        self.calls = []  # list of dicts

    def __call__(self, a, newshape, nghosts=0, method="rebin", **kwargs):
        # Record call metadata
        rec = {
            "method": method,
            "newshape": tuple(int(x) for x in newshape),
            "nghosts": int(nghosts),
            "in_shape": tuple(a.shape),
        }
        rec.update(kwargs)  # record limiter, etc., if passed
        self.calls.append(rec)

        ng = int(nghosts)
        if ng > 0:
            core = a[(slice(ng, -ng),) * a.ndim]
        else:
            core = a

        if core.size == 0:
            return np.zeros(tuple(int(x) for x in newshape), dtype=float)

        val = float(np.mean(core))
        return np.full(tuple(int(x) for x in newshape), val, dtype=float)


# -----------------------------------------------------------------------------
# Tests
# -----------------------------------------------------------------------------

def test_uniform_grid_2d_two_blocks_tile_correctly(monkeypatch):
    """
    Two adjacent blocks in x should tile into the global array with correct offsets.
    """
    spy = InterpSpy()
    monkeypatch.setattr(amun_mod, "interpolate", spy)

    nc, ng = 4, 1
    ncellg = nc + 2 * ng

    # Two leaf blocks at level 1 (toplev=1), coords in xyz => (x,y)
    chunks = {
        "dblocks": 2,
        "dims": (2, ncellg, ncellg),
        "levels": np.array([1, 1], dtype=int),
        "coords": np.array([[0, 0], [1, 0]], dtype=int),  # x=0 and x=1
    }

    a = FakeAmun(
        2,
        xblocks=2,
        yblocks=1,
        zblocks=1,
        toplev=1,
        maxlev=1,
        ncells=nc,
        nghosts=ng,
        chunks=chunks,
        raw_values={"dens": [1.0, 2.0]},
    )

    out = a.dataset("density")  # default target_level=toplev=1

    assert out.shape == (nc, 2 * nc)
    assert np.all(out[:, :nc] == 1.0)
    assert np.all(out[:, nc:] == 2.0)


def test_uniform_grid_2d_extent_cropping(monkeypatch):
    """
    Extent cropping should select the correct sub-region (here: the right block).
    """
    spy = InterpSpy()
    monkeypatch.setattr(amun_mod, "interpolate", spy)

    nc, ng = 4, 1
    ncellg = nc + 2 * ng

    chunks = {
        "dblocks": 2,
        "dims": (2, ncellg, ncellg),
        "levels": np.array([1, 1], dtype=int),
        "coords": np.array([[0, 0], [1, 0]], dtype=int),
    }

    a = FakeAmun(
        2,
        xblocks=2,
        yblocks=1,
        zblocks=1,
        toplev=1,
        maxlev=1,
        ncells=nc,
        nghosts=ng,
        chunks=chunks,
        raw_values={"dens": [1.0, 2.0]},
    )

    # Domain: x in [0,2], y in [0,1]. Crop right half: x in [1,2].
    out = a.dataset("density", extent=[1.0, 2.0, 0.0, 1.0])

    assert out.shape == (nc, nc)
    assert np.all(out == 2.0)


def test_uniform_grid_target_level_downsamples_consistently(monkeypatch):
    """
    target_level should deterministically control output resolution.

    We construct a consistent toplev=2 domain (xblocks=2 at level-1 => 4 leaf blocks at level-2).
    Then check:
      - target_level=2 => full resolution
      - target_level=1 => downsampled by factor 2 in each direction
    """
    spy = InterpSpy()
    monkeypatch.setattr(amun_mod, "interpolate", spy)

    nc, ng = 4, 1
    ncellg = nc + 2 * ng

    # toplev=2, xblocks(level1)=2 => effective leaf blocks in x at level2 is 4
    # coords are leaf-block coords at their level (here all at level 2)
    chunks = {
        "dblocks": 4,
        "dims": (4, ncellg, ncellg),
        "levels": np.array([2, 2, 2, 2], dtype=int),
        "coords": np.array([[0, 0], [1, 0], [2, 0], [3, 0]], dtype=int),
    }

    a = FakeAmun(
        2,
        xblocks=2,
        yblocks=1,
        zblocks=1,
        toplev=2,
        maxlev=2,
        ncells=nc,
        nghosts=ng,
        chunks=chunks,
        raw_values={"dens": [1.0, 2.0, 3.0, 4.0]},
    )

    out_lvl2 = a.dataset("density", target_level=2)
    out_lvl1 = a.dataset("density", target_level=1)

    # Full resolution at target_level=2: xcells = xblocks*nc*2**(2-1) = 2*4*2 = 16
    assert out_lvl2.shape == (nc * 2, nc * 4)  # y: 8, x: 16
    # Downsample at target_level=1: xcells = 2*4*2**0 = 8
    assert out_lvl1.shape == (nc, nc * 2)      # y: 4, x: 8

    # Coarse result should contain the same block-constant structure (just fewer cells):
    # left-to-right 4 leaf blocks at level2 each contribute 2 columns at target_level=1
    assert np.all(out_lvl1[:, 0:2] == 1.0)
    assert np.all(out_lvl1[:, 2:4] == 2.0)
    assert np.all(out_lvl1[:, 4:6] == 3.0)
    assert np.all(out_lvl1[:, 6:8] == 4.0)


def test_uniform_grid_interpolation_dispatch_policy(monkeypatch):
    """
    Verify dispatch policy: method='rebin' when nl <= shrink_equiv, else user interpolation.

    We set toplev=2 and target_level=2 => shrink_equiv=1.
    Then:
      - level=2 block => nl=1  => rebin
      - level=1 block => nl=2  => user interpolation (e.g., 'linear')
    """
    spy = InterpSpy()
    monkeypatch.setattr(amun_mod, "interpolate", spy)

    nc, ng = 4, 1
    ncellg = nc + 2 * ng

    # Two blocks stacked in y so they don't overlap:
    # coords are xyz (x,y). xblocks=1, yblocks=2, toplev=2 => domain ycells=16 at target_level=2.
    chunks = {
        "dblocks": 2,
        "dims": (2, ncellg, ncellg),
        "levels": np.array([2, 1], dtype=int),
        "coords": np.array([[0, 0], [0, 1]], dtype=int),
    }

    a = FakeAmun(
        2,
        xblocks=1,
        yblocks=2,
        zblocks=1,
        toplev=2,
        maxlev=2,
        ncells=nc,
        nghosts=ng,
        chunks=chunks,
        raw_values={"dens": [1.0, 2.0]},
    )

    _ = a.dataset("density", target_level=2, interpolation="linear")

    methods = [c["method"] for c in spy.calls]
    # We expect both policies to appear:
    assert "rebin" in methods
    assert "linear" in methods


def test_uniform_grid_target_level_downsamples_consistently(monkeypatch):
    """
    target_level should deterministically control output resolution.

    We construct a consistent toplev=2 domain with full y-coverage:
      - xblocks(level1)=2 => 4 leaf blocks in x at level2
      - yblocks(level1)=1 => 2 leaf blocks in y at level2

    Then check:
      - target_level=2 => full resolution
      - target_level=1 => downsampled by factor 2 in each direction
    """
    spy = InterpSpy()
    monkeypatch.setattr(amun_mod, "interpolate", spy)

    nc, ng = 4, 1
    ncellg = nc + 2 * ng

    # Leaf blocks at level 2: coords cover x=0..3 and y=0..1 (full domain)
    coords = []
    vals = []
    for y in (0, 1):
        for x in (0, 1, 2, 3):
            coords.append([x, y])
            # Make both y-blocks share the same x-pattern so coarse-grid
            # assertions can be made column-wise across all rows.
            vals.append(float(x + 1))  # 1,2,3,4 repeated for y=0 and y=1

    chunks = {
        "dblocks": 8,
        "dims": (8, ncellg, ncellg),
        "levels": np.array([2] * 8, dtype=int),
        "coords": np.array(coords, dtype=int),
    }

    a = FakeAmun(
        2,
        xblocks=2,
        yblocks=1,
        zblocks=1,
        toplev=2,
        maxlev=2,
        ncells=nc,
        nghosts=ng,
        chunks=chunks,
        raw_values={"dens": vals},
    )

    out_lvl2 = a.dataset("density", target_level=2)
    out_lvl1 = a.dataset("density", target_level=1)

    # Full resolution at target_level=2:
    # xcells = xblocks*nc*2**(2-1) = 2*4*2 = 16
    # ycells = yblocks*nc*2**(2-1) = 1*4*2 = 8
    assert out_lvl2.shape == (nc * 2, nc * 4)  # (8, 16)

    # Downsample at target_level=1:
    # xcells = 2*4*2**0 = 8
    # ycells = 1*4*2**0 = 4
    assert out_lvl1.shape == (nc, nc * 2)      # (4, 8)

    # Coarse result should keep the left-to-right block structure:
    # 4 leaf blocks in x at level2 each contribute 2 columns at target_level=1.
    assert np.all(out_lvl1[:, 0:2] == 1.0)
    assert np.all(out_lvl1[:, 2:4] == 2.0)
    assert np.all(out_lvl1[:, 4:6] == 3.0)
    assert np.all(out_lvl1[:, 6:8] == 4.0)
