import numpy as np
import pytest
from numpy.testing import assert_allclose

import topovec as tv
from topovec.io import load_lcsim_npz, save_npz_lcsim


def test_save_lcsim_npz_roundtrip(tmp_path):
    nn = np.arange(2 * 3 * 4 * 1 * 3, dtype=np.float32).reshape(2, 3, 4, 1, 3)
    path = tmp_path / "sample.npz"

    save_npz_lcsim(path, nn, settings=dict(L=33, I0toIs=0, p0=-36, pinf=-36))

    images, system, settings = load_lcsim_npz(path)

    assert len(images) == 1
    assert_allclose(images[0], nn)
    assert tuple(system.size) == (2, 3, 4)
    assert settings["L"] == 33
    assert settings["I0toIs"] == 0
    assert settings["p0"] == -36
    assert settings["pinf"] == -36
    assert settings["sx"] == 2
    assert settings["sy"] == 3
    assert settings["sz"] == 4
    assert settings["fixed_dir"] == [0, 0, 1]


def test_save_lcsim_npz_periodic_boundary_condition(tmp_path):
    nn = np.zeros((2, 2, 2, 1, 3), dtype=np.float32)

    periodic_path = tmp_path / "periodic.npz"
    save_npz_lcsim(periodic_path, nn, settings={}, periodic=True)
    _, _, periodic_settings = load_lcsim_npz(periodic_path)
    assert periodic_settings["bc"] == [0, 1, 0]

    nonperiodic_path = tmp_path / "nonperiodic.npz"
    save_npz_lcsim(nonperiodic_path, nn, settings={}, periodic=False)
    _, _, nonperiodic_settings = load_lcsim_npz(nonperiodic_path)
    assert nonperiodic_settings["bc"] == [0, 0, 0]


@pytest.mark.parametrize(
    "shape",
    [
        (2, 3, 4, 3),
        (2, 3, 4, 1, 2),
    ],
)
def test_save_lcsim_npz_rejects_invalid_shapes(tmp_path, shape):
    nn = np.zeros(shape, dtype=np.float32)

    with pytest.raises(ValueError, match="`nn` must have shape"):
        save_npz_lcsim(tmp_path / "invalid.npz", nn, settings={})


def test_top_level_io_namespace():
    assert tv.io.load_lcsim_npz is load_lcsim_npz
    assert tv.io.save_npz_lcsim is save_npz_lcsim
