import os

import numpy as np
import pytest
import h5py


def test_shell_attrs(datadir, thickshell):
    shell_type, shell = thickshell

    assert shell.geom_type == "thick"

    # shape
    np.testing.assert_allclose(shell.shape, (8, 10, 4))

    filename = os.path.join(datadir, f"{shell_type}.h5")
    with h5py.File(filename) as file:
        # mesh
        for i, var in enumerate(["x", "y", "z"]):
            print("var", var)
            np.testing.assert_allclose(
                np.take(shell.xyz, i, -1).ravel(), file[f"mesh/{var}"][()], atol=1.0e-15
            )

        # vars
        np.testing.assert_allclose(shell.dz, file["variables/dz"], atol=1.0e-15)


def test_dump(tmpdir, thickshell):
    _, shell = thickshell

    filename = os.path.join(tmpdir, "shell")

    # create a random (external) field
    data = np.random.random(shell.shape)
    data_dict = {"random_external": data}

    # create a random (internal) field
    data = np.random.random(shell.shape)
    shell.fields["random_internal"] = data

    # TODO: delete after deprecation
    shell.dump_shell(f"{filename}_bc", fields=data_dict)
    shell.dump(filename, fields=data_dict)
    size_target = np.prod(shell.shape)
    with h5py.File(f"{filename}.h5", "r") as fin:

        # mesh
        for var in ["x", "y", "z"]:
            assert fin[f"mesh/{var}"][()].size == size_target

        # other variables
        for var in ["dz", "random_external", "random_internal"]:
            assert fin[f"variables/{var}"][()].size == size_target
