import numpy as np
from amunpy.amun import Amun

class FakeAmun(Amun):
    # bypass Amun.__init__
    def __init__(self):
        self.dataformat = "FAKE"
        self.attributes = {
            "eos": "iso",
            "sound_speed": 2.0,
        }
        self.variables = {
            "dens": "dens",
            "pres": "pres",
            "velo": "velo",
            "magn": "magn",
            "ekin": "ekin",
        }
        self.chunks = {1: {"dims": (2, 2, 2), "dblocks": 1, "levels": [0]}}

        self._store = {
            "dens": np.full((2, 2, 2), 3.0),
            "pres": np.full((2, 2, 2), 5.0),
            "velx": np.full((2, 2, 2), 1.0),
            "vely": np.full((2, 2, 2), 2.0),
            "velz": np.full((2, 2, 2), 2.0),
            "magx": np.full((2, 2, 2), 2.0),
            "magy": np.full((2, 2, 2), 0.0),
            "magz": np.full((2, 2, 2), 0.0),
        }

    # IMPORTANT: this is NOT mangled because it ends with "__"
    def __read_binary_data__(self, dataset_name, chunk_number):
        return self._store[dataset_name]

def test_get_dataset_pres_isothermal_uses_dens_cs2():
    a = FakeAmun()
    out = a.__get_dataset__("pres", 1)
    assert np.allclose(out, 3.0 * (2.0**2))  # dens * cs^2 = 12

def test_get_dataset_velo_is_speed_magnitude():
    a = FakeAmun()
    out = a.__get_dataset__("velo", 1)
    assert np.allclose(out, 3.0)  # sqrt(1 + 4 + 4)

def test_get_dataset_magn_is_field_magnitude():
    a = FakeAmun()
    out = a.__get_dataset__("magn", 1)
    assert np.allclose(out, 2.0)

def test_get_dataset_ekin_is_half_rho_v2():
    a = FakeAmun()
    out = a.__get_dataset__("ekin", 1)
    assert np.allclose(out, 13.5)  # 0.5 * 3 * 9
