import numpy as np
import pytest

from amunpy.amun import Amun
from amunpy.exceptions import AmunPyVariableError


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

        # minimal attributes used by resolvers
        self.attributes = {
            "ndims": 3,
            "adiabatic_index": 5.0 / 3.0,
            "eos": "adi",
            "eqsys": "mhd",
        }

        # minimal chunk metadata
        self.chunks = {
            0: {
                "dblocks": 1,
                "dims": (1, 2, 2, 2),
                "levels": np.array([1]),
            }
        }

        # raw snapshot variables
        self.raw_variables = {
            "dens": "dens",
            "velx": "velx",
            "vely": "vely",
            "velz": "velz",
        }

    # ---- raw readers -------------------------------------------------

    def __read_binary_data__(self, dataset_id, chunk_number):
        if dataset_id == "dens":
            return np.ones(self.chunks[chunk_number]["dims"])
        if dataset_id == "velx":
            return np.ones(self.chunks[chunk_number]["dims"])
        if dataset_id == "vely":
            return 2.0 * np.ones(self.chunks[chunk_number]["dims"])
        if dataset_id == "velz":
            return 3.0 * np.ones(self.chunks[chunk_number]["dims"])
        raise RuntimeError(f"unexpected dataset_id={dataset_id}")

    def __read_vec3__(self, prefix, chunk_number):
        if prefix == "vel":
            return (
                self.__read_binary_data__("velx", chunk_number),
                self.__read_binary_data__("vely", chunk_number),
                self.__read_binary_data__("velz", chunk_number),
            )
        raise RuntimeError(f"unexpected prefix={prefix}")

    # ---- vector calculus helpers ------------------------------------

    def __curl_chunk__(self, chunk_number, ax, ay, az):
        # deterministic dummy curl
        return (
            np.ones_like(ax),
            2.0 * np.ones_like(ax),
            3.0 * np.ones_like(ax),
        )

def test_get_dataset_raw_scalar():
    a = FakeAmun()
    d = a.__get_dataset__("dens", 0)
    assert d.shape == (1, 2, 2, 2)
    assert np.all(d == 1.0)

def test_get_dataset_raw_vector():
    a = FakeAmun()
    vx, vy, vz = a.__get_dataset__("vvec", 0)

    assert np.all(vx == 1.0)
    assert np.all(vy == 2.0)
    assert np.all(vz == 3.0)

def test_get_dataset_derived_scalar_velo():
    a = FakeAmun()
    v = a.__get_dataset__("velo", 0)

    expected = np.sqrt(1.0**2 + 2.0**2 + 3.0**2)
    assert np.allclose(v, expected)

def test_get_dataset_derived_vector_wvec():
    a = FakeAmun()
    wx, wy, wz = a.__get_dataset__("wvec", 0)

    assert np.all(wx == 1.0)
    assert np.all(wy == 2.0)
    assert np.all(wz == 3.0)

def test_get_dataset_unknown_raises():
    a = FakeAmun()
    with pytest.raises(AmunPyVariableError):
        a.__get_dataset__("no_such_dataset", 0)

def test_resolver_registry_built_once():
    a = FakeAmun()
    a.__get_dataset__("dens", 0)
    assert hasattr(a, "_dataset_resolvers")

