import os

import numpy as np
import h5py
import pytest

from kokiy.cartshell import CartShell


ATOL = 1e-8


@pytest.mark.parametrize(
    "load_shell_case",
    [
        ("cartshell", "cart"),
        ("axishell", "axicyl"),
        ("torch_shell", "torch"),
    ],
    indirect=True,
)
def test_shell_attrs(datadir, load_shell_case):
    # frustum shell is tested differently
    (shell_type, shell), geom_type = load_shell_case

    assert shell.geom_type == geom_type

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

    filename = os.path.join(datadir, f"{shell_type}.h5")

    target_vars = [
        "theta",
        "n_x",
        "n_y",
        "n_z",
        "du",
        "u",
        "v",
        "dv",
        "dwu",
        "dwv",
        "surf",
    ]
    if shell_type not in ["frustumshell", "torchshell"]:
        target_vars.extend(["n_r"])

    with h5py.File(filename) as file:

        # mesh
        for i, var in enumerate(["x", "y", "z"]):
            np.testing.assert_allclose(
                np.take(shell.xyz, i, -1).ravel(), file[f"mesh/{var}"][()], atol=1.0e-15
            )

        # radius (different name)
        np.testing.assert_allclose(shell.rad, file["variables/r"][()], atol=1.0e-15)

        # other vars
        for var in target_vars:
            print("var", var)
            np.testing.assert_allclose(
                getattr(shell, var), file[f"variables/{var}"][()], atol=1.0e-15
            )


def test_dump(tmpdir, shell_2d):
    shell_type, shell = shell_2d

    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

    # dump shell
    # TODO: delete after deprecation
    shell.dump_shell(f"{filename}_bc", fields=data_dict)
    shell.dump(filename, fields=data_dict)
    size_target = shell.shape[0] * shell.shape[1]

    target_vars = [
        "r",
        "theta",
        "n_x",
        "n_y",
        "n_z",
        "du",
        "u",
        "v",
        "dv",
        "dwu",
        "dwv",
        "surf",
        "random_external",
        "random_internal",
    ]
    if shell_type not in ["frustumshell", "torchshell"]:
        target_vars.extend(["n_r"])

    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 target_vars:
            assert fin[f"variables/{var}"][()].size == size_target


def test_average_over_dir(axishell):
    # TODO: is this test meaningful enough?
    _, axishell = axishell

    np.random.seed(1)
    target_field = np.random.random((3, axishell.shape[0], axishell.shape[1]))

    # with scale
    average = axishell.average_on_shell_over_dirs(target_field, ["v", "time"])
    target_val = np.array(
        [
            4.66739841e-05,
            1.27688882e-04,
            1.14211735e-04,
            1.08330533e-04,
            9.72293499e-05,
            1.22829723e-04,
            1.06270736e-04,
            1.32217197e-04,
            1.19892968e-04,
            6.73650501e-05,
        ]
    )
    np.testing.assert_allclose(average, target_val, atol=ATOL)

    # without scale
    average = axishell.average_on_shell_over_dirs(
        target_field, ["v", "time"], scale=False
    )
    target_val = np.array(
        [
            0.41297985,
            0.55279941,
            0.48937754,
            0.4651728,
            0.40128855,
            0.51424186,
            0.43755768,
            0.56629453,
            0.53040696,
            0.59220561,
        ]
    )
    np.testing.assert_allclose(average, target_val, atol=ATOL)

    # in case of bad definition of directions
    with pytest.raises(ValueError):
        axishell.average_on_shell_over_dirs(target_field, ["v", "k"])


def test_integrate_over_dir():
    # create a very simple shell
    n_trans, n_longi = 5, 6
    zero = np.zeros(3)
    umax = np.array([1.0, 0.0, 0.0])
    vmax = np.array([0.0, 1.0, 0.0])
    shell = CartShell(n_trans, n_longi, zero, umax, vmax)

    # create a field
    def field(x, t):
        return x[:, 0] * t  # just a linear field

    ts = [1.0, 2.0, 3.0]
    n_t = len(ts)
    fields = np.array(
        [field(shell.xyz.reshape(-1, 3), t).reshape(shell.shape) for t in ts]
    )

    # check average u
    average_u = shell.operate_on_shell_over_dirs(
        fields, directions=["u"], operator=np.mean
    )
    assert average_u.shape == (n_t, n_trans)
    for k in range(n_t):
        np.testing.assert_allclose(average_u[k], fields[k, :, 0])

    # check average v
    average_v = shell.operate_on_shell_over_dirs(
        fields, directions=["v"], operator=np.mean
    )
    assert average_v.shape == (n_t, n_longi)
    for k, t in enumerate(ts):
        assert (average_v[k][0] - t / 2) < ATOL

    # check average uv
    average_uv = shell.operate_on_shell_over_dirs(
        fields, directions=["v", "u"], operator=np.mean
    )
    assert average_uv.shape == (n_t,)
    for k, t in enumerate(ts):
        assert (average_uv[k] - t / 2) < ATOL

    # check average t (order do not matter)
    average_t = shell.operate_on_shell_over_dirs(
        fields, directions=["v", "u", "time"], operator=np.mean
    )
    assert (average_t - 1.0) < ATOL

    # check var u
    var_u = shell.operate_on_shell_over_dirs(fields, directions=["u"], operator=np.var)
    assert np.sum(np.abs(var_u)) == 0.0

    # check var v
    var_v = shell.operate_on_shell_over_dirs(fields, directions=["v"], operator=np.var)
    for k, var in enumerate([0.125, 0.5, 1.125]):
        assert (var_v[k][0] - var) < ATOL

    # check var v time
    var_vt = shell.operate_on_shell_over_dirs(
        fields, directions=["v", "time"], operator=np.var
    )
    assert var_vt.shape == (n_longi,)
    assert (np.mean(var_vt) - 0.17013889) < ATOL
    assert (var_vt[0] - 0.17013889) < ATOL

    # check var time v
    var_tv = shell.operate_on_shell_over_dirs(
        fields, directions=["time", "v"], operator=np.var
    )
    assert var_tv.shape == (n_longi,)
    assert (np.mean(var_tv) - 0.06041667) < ATOL
    assert (var_tv[0] - 0.06041667) < ATOL


def test_plot(tmpdir, shell_2d):
    _, shell = shell_2d
    shell.plot(os.path.join(tmpdir, "fig.png"))


def test_invert_normals(shell_2d):
    shell_type, shell = shell_2d

    normal_vars = ["n_x", "n_y", "n_z"]
    if shell_type not in ["frustumshell", "torchshell"]:
        normal_vars.append("n_z")

    target_val = -np.array([getattr(shell, var) for var in normal_vars])

    shell.invert_normals()
    new_normals = np.array([getattr(shell, var) for var in normal_vars])

    np.testing.assert_allclose(new_normals, target_val)


def test_baking(tmpdir, axishell):
    _, shell_2d = axishell

    # test add_curviwidth
    width_profile = (
        (0.0, -0.04),
        (0.2, -0.04),
        (0.4, -0.07),
        (0.6, -0.04),
        (0.8, -0.04),
    )
    shell_2d.add_curviwidth("width_profile", width_profile)
    target_val = np.array(
        [
            -0.04,
            -0.04,
            -0.04333333,
            -0.06,
            -0.06333333,
            -0.04666667,
            -0.04,
            -0.04,
            -0.04,
            -0.04,
        ]
    )
    np.testing.assert_allclose(
        shell_2d.width_matrix["width_profile"][0], target_val, atol=ATOL
    )

    # test addition of same name (stateful)
    shell_2d.add_curviwidth("width_profile", width_profile)
    np.testing.assert_allclose(
        shell_2d.width_matrix["width_profile"][0], 2 * target_val, atol=ATOL
    )

    # test bake_millefeuille
    cake = shell_2d.bake_millefeuille("width_profile", 4)

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

    # check dict like behavior
    cake["random_internal"] = np.random.random(cake.shape)

    filename = os.path.join(tmpdir, "shell")
    cake.dump(filename=filename, fields=data_dict)

    size_target = np.prod(cake.shape)
    with h5py.File(f"{filename}.h5", "r") as fin:

        # assert xyz do not exist
        with pytest.raises(KeyError):
            fin["variables/xyz"]

        # 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

    # test replicate
    new_shape = (4, 6, 3)
    new_cake = cake.replicate(shape=new_shape)
    np.testing.assert_array_equal(new_cake.shape, new_shape)


def test_rad_theta():
    # create a very simple shell
    n_trans, n_longi = 3, 3
    zero = np.array([1.0, 1.0, 0.0])
    umax = np.array([2.0, 1.0, 0.0])
    vmax = np.array([1.0, 2.0, 0.0])
    shell_xy = CartShell(n_trans, n_longi, zero, umax, vmax)

    vel_y = np.full(shell_xy.shape, 1.0)
    vel_z = np.full(shell_xy.shape, 0.0)
    vel_rad, vel_tht = shell_xy.rad_theta_components(vel_y, vel_z)
    assert np.allclose(vel_rad, 1.0)
    assert np.allclose(vel_tht, 0.0)

    vel_y = np.full(shell_xy.shape, 0.0)
    vel_z = np.full(shell_xy.shape, 1.0)
    vel_rad, vel_tht = shell_xy.rad_theta_components(vel_y, vel_z)
    assert np.allclose(vel_rad, 0.0)
    assert np.allclose(vel_tht, 1.0)
