import numpy as np
import pytest

from amunpy.amun import Amun
from amunpy.exceptions import AmunPyChunkError


class FakeAmun(Amun):
    def __init__(self):
        self.path = "fake"
        self.dataformat = "fake"
        self.attributes = {}
        self.chunks = {}
        self.variables = {}
        self.initialized = True

    def __read_binary_data__(self, *args, **kwargs):
        raise RuntimeError("not used")


def test_check_chunks_allows_2d_coords_with_dummy_z():
    a = FakeAmun()
    a.attributes = {
        "nchunks": 1,
        "ndims": 2,
        "maxlev": 2,
        "xblocks": 1,
        "yblocks": 1,
    }

    a.chunks = {
        0: {
            "filename": "c0",
            "dblocks": 4,
            "levels": np.array([1, 1, 1, 1], dtype=int),
            "coords": np.zeros((4, 3), dtype=float),      # (dblocks, 3) in 2D is allowed
            "bounds": np.zeros((4, 2, 3), dtype=float),   # likewise for bounds
        }
    }
    # should not raise
    a.__check_chunks__()


def test_check_chunks_rejects_too_few_coord_dims():
    a = FakeAmun()
    a.attributes = {"nchunks": 1, "ndims": 2, "maxlev": 2}
    a.chunks = {
        0: {
            "filename": "c0",
            "dblocks": 4,
            "levels": np.array([1, 1, 1, 1], dtype=int),
            "coords": np.zeros((4, 1), dtype=float),      # insufficient for ndims=2
            "bounds": np.zeros((4, 2, 2), dtype=float),
        }
    }
    with pytest.raises(AmunPyChunkError):
        a.__check_chunks__()

