import os

import pytest
import numpy as np

from kokiy.cartshell import CartShell
from kokiy.shell_utils import interpolate_solutions_on_shell
from kokiy.shell_utils import interpolate_solution_on_shell
from kokiy.shell_utils import infer_shell_from_xyz
from kokiy.shell_utils import shell_repr
from kokiy.shell_utils import shell_geom_type
from kokiy.shell_utils import _from_np_struct_to_dict


ATOL = 1e-6


def test_interpolate_solutions_on_shell_2d(datadir):

    # create fake shell to extract mesh
    n_trans, n_longi = 12, 24
    zero = np.array([0.00, 0.12, -0.2])
    u_max = np.array([0.00, 0.12, 0.2])
    v_max = np.array([0.20, 0.12, -0.2])
    init_shell = CartShell(n_trans, n_longi, zero, u_max, v_max)
    init_mesh = init_shell.xyz.reshape(-1, 3)

    # random data (a value for each node in the initial shell)
    np.random.seed(1)
    n_time, n_vars = 3, 2
    names = [f"random{i}" for i in range(n_vars)]
    dtype = dict(names=names, formats=[np.float32] * n_vars)
    shape = (n_time, init_mesh.shape[0])
    data = np.empty(shape, dtype=dtype)
    for name in names:
        data[name] = np.random.random(shape)

    # create new shell in which data will be interpolated
    n_trans, n_longi = 6, 12  # coarser than original
    zero = np.array([0.00, 0.12, -0.2])
    u_max = np.array([0.00, 0.12, 0.2])
    v_max = np.array([0.20, 0.12, -0.2])
    new_shell = CartShell(n_trans, n_longi, zero, u_max, v_max)

    # interpolate results (test array-like input)
    results = interpolate_solutions_on_shell(new_shell, init_mesh, data)

    # interpolate results (test dict-like input)
    data_dict = _from_np_struct_to_dict(data)
    results_dict = interpolate_solutions_on_shell(new_shell, init_mesh, data_dict)

    # compare with target
    filename = os.path.join(datadir, "interp_data.npz")
    target_results = np.load(filename, allow_pickle=True)
    for name in target_results.dtype.names:
        np.testing.assert_allclose(results[name], target_results[name], atol=ATOL)
        np.testing.assert_allclose(results_dict[name], target_results[name], atol=ATOL)

    # interpolate one solution interpolation
    var_name = "random0"
    result = interpolate_solution_on_shell(new_shell, init_mesh, data_dict[var_name])
    np.testing.assert_allclose(result, results_dict[var_name])

    # interpolate one solution at one time step only
    var_name = "random0"
    result = interpolate_solution_on_shell(new_shell, init_mesh, data_dict[var_name][0])
    np.testing.assert_allclose(result, results_dict[var_name][0])


def test_interpolate_solutions_on_shell_3d(thickshell):
    # just tests if it is able to interpolate itself
    # test on 2d ensures interpolation is not changing

    shell_type, init_shell = thickshell

    for name in ["frustum", "torch"]:  # TODO: make it work
        if name in shell_type:
            return

    new_shell = init_shell.replicate(init_shell.shape)

    init_mesh = init_shell.xyz.reshape(-1, 3)
    np.random.seed(1)
    data = np.random.random(
        (
            2,
            init_mesh.shape[0],
        )
    )

    # interpolate
    new_data = interpolate_solution_on_shell(new_shell, init_mesh, data)

    np.testing.assert_allclose(new_data.reshape(2, -1), data, atol=ATOL)


def test_cartshell_create_structured_base():
    """
    Test of generation of cartshell
    """

    n_y = 3
    n_z = 3

    mesh = _create_mesh()

    shell = infer_shell_from_xyz(mesh, "cart", (n_y, n_z))
    xyz_tgt = np.array(
        [
            [
                [0.0, 0.0, 0.0],
                [0.0, 0.5, 0.0],
                [0.0, 1.0, 0.0],
            ],
            [
                [0.0, 0.0, 0.5],
                [0.0, 0.5, 0.5],
                [0.0, 1.0, 0.5],
            ],
            [
                [0.0, 0.0, 1.0],
                [0.0, 0.5, 1.0],
                [0.0, 1.0, 1.0],
            ],
        ]
    )
    np.testing.assert_allclose(xyz_tgt, shell.xyz)


def test_axishell_create_structured_based():
    # TODO: need to test bnd_uv

    mesh = _create_mesh()

    n_r = 3
    n_theta = 3
    shell = infer_shell_from_xyz(mesh, "axicyl", (n_r, n_theta))
    xyz_tgt = np.array(
        [
            [
                [0.0, 7.08085476e-01, 0.0],
                [0.0, 8.54042738e-01, 0.0],
                [0.0, 1.0, 0.0],
            ],
            [
                [0.0, 5.00692042e-01, 5.00692042e-01],
                [0.0, 6.03899412e-01, 6.03899412e-01],
                [0.0, 7.07106781e-01, 7.07106781e-01],
            ],
            [
                [0.0, 0.0, 7.08085476e-01],
                [0.0, 0.0, 8.54042738e-01],
                [0.0, 0.0, 1.0],
            ],
        ]
    )
    np.testing.assert_allclose(xyz_tgt, shell.xyz, atol=ATOL)


@pytest.mark.parametrize(
    "load_shell_case",
    [
        ("cartshell", "cart"),
        ("axishell", "axicyl"),
        ("torch_shell", "torch"),
        ("frustum_shell", "frustum"),
    ],
    indirect=True,
)
def test_shell_repr(load_shell_case):
    # TODO: will be deprecated -> move test to within each of the objects
    (_, shell), geom_type = load_shell_case

    target_value = f"Shell\n-----\n\nshape : (8, 10)\ntype : {geom_type}"
    assert shell_repr(shell) == target_value
    assert repr(shell) == target_value


@pytest.mark.parametrize(
    "load_shell_case",
    [
        ("cartshell", "cart"),
        ("axishell", "axicyl"),
        ("torch_shell", "torch"),
        ("frustum_shell", "frustum"),
    ],
    indirect=True,
)
def test_shell_geom_type(load_shell_case):
    # TODO: will be deprecate (delete)
    (_, shell), geom_type = load_shell_case
    assert shell_geom_type(shell) == geom_type


def _create_mesh():
    yi = np.linspace(0, 1, 20)
    zi = np.linspace(1, 0, 20)
    xi = np.zeros_like(yi)

    return np.stack((xi, yi, zi), axis=-1)
