import numpy as np
import pytest

from amunpy.amun import Amun


class FakeAmun(Amun):
    def __init__(self, ndims=2):
        self.path = "fake"
        self.dataformat = "fake"
        self.initialized = True

        nc = 4
        ng = 1  # IMPORTANT: dataset()/interpolate assumes ghost zones exist

        self.attributes = {
            "ndims": ndims,
            "ncells": nc,
            "nghosts": ng,
            "toplev": 1,
            "maxlev": 1,
            "nchunks": 1,
            "nleafs": 1,
            "xblocks": 1,
            "yblocks": 1,
            "zblocks": 1,
            "xmin": 0.0,
            "xmax": 1.0,
            "ymin": 0.0,
            "ymax": 1.0,
            "zmin": 0.0,
            "zmax": 1.0,
        }

        # raw variable catalog
        self.raw_variables = {"dens": "dens", "velx": "velx", "vely": "vely", "velz": "velz"}

        # alias catalog
        self.variables = {
            "density": "dens",
            "velocity": "vvec",
        }

        # one chunk, one block (include ghost zones in dims)
        ncellg = nc + 2 * ng
        if ndims == 2:
            dims = (1, ncellg, ncellg)
            coords = np.array([[0, 0]], dtype=int)
        else:
            dims = (1, ncellg, ncellg, ncellg)
            coords = np.array([[0, 0, 0]], dtype=int)

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

    # --- raw readers -------------------------------------------------
    def __read_binary_data__(self, dataset_id, chunk_number):
        shape = self.chunks[chunk_number]["dims"]
        if dataset_id == "dens":
            return np.ones(shape)
        if dataset_id == "velx":
            return np.ones(shape)
        if dataset_id == "vely":
            return 2.0 * np.ones(shape)
        if dataset_id == "velz":
            return 3.0 * np.ones(shape)
        raise RuntimeError(dataset_id)

    def __read_vec3__(self, prefix, chunk_number):
        return (
            self.__read_binary_data__("velx", chunk_number),
            self.__read_binary_data__("vely", chunk_number),
            self.__read_binary_data__("velz", chunk_number),
        )

def test_dataset_pipeline_scalar_2d():
    a = FakeAmun(ndims=2)
    arr = a.dataset("density")

    assert isinstance(arr, np.ndarray)
    assert arr.shape == (4, 4)
    assert np.all(arr == 1.0)

def test_dataset_pipeline_vector_2d():
    a = FakeAmun(ndims=2)
    vec = a.dataset("velocity")

    assert isinstance(vec, list)
    assert len(vec) == 3
    assert all(isinstance(v, np.ndarray) for v in vec)

    vx, vy, vz = vec
    assert vx.shape == (4, 4)
    assert np.all(vx == 1.0)
    assert np.all(vy == 2.0)
    assert np.all(vz == 3.0)

def test_dataset_pipeline_scalar_3d():
    a = FakeAmun(ndims=3)
    arr = a.dataset("density")

    assert arr.shape == (4, 4, 4)

