import os
import json

import pytest
import numpy as np

from kokiy._mesh_utils import get_conns
from kokiy._mesh_utils import get_bound_nodes


def read_conns_data(datadir):

    filename = os.path.join(datadir, "conns.json")
    with open(filename, "r") as file:
        data = json.load(file)

    return data


@pytest.mark.parametrize("reverse", [False, True])
@pytest.mark.parametrize("elem_type", ["quad", "triangle"])
def test_conns_2d(datadir, elem_type, reverse):
    data = read_conns_data(datadir)

    c, r = 3, 4

    elem_type = "quad"
    conns = get_conns(c, r, elem_type=elem_type)

    expected_conns = data["conns_2d"][elem_type]

    np.testing.assert_array_equal(conns, expected_conns)


@pytest.mark.parametrize("reverse", [True, False])
@pytest.mark.parametrize("elem_type", ["tetra", "hexahedron"])
def test_conns_3d(datadir, elem_type, reverse):
    data = read_conns_data(datadir)

    c, r, n = 3, 4, 3
    conns = get_conns(c, r, n, elem_type=elem_type, reverse=reverse)

    expected_conns = data["conns_3d"][elem_type][str(reverse).lower()]

    np.testing.assert_array_equal(conns, expected_conns)


def test_bound_nodes_2d():
    c, r = 3, 4
    bnd_nodes = get_bound_nodes(c, r)

    # assumes dict as ordered
    exp_nodes = [[0, 1, 2, 3], [3, 7, 11], [8, 9, 10, 11], [0, 4, 8]]

    for bnd_node_group, exp_node_group in zip(bnd_nodes.values(), exp_nodes):
        np.testing.assert_array_equal(bnd_node_group, exp_node_group)


def test_bound_nodes_3d():
    c, r, n = 4, 3, 3

    # normal order
    bnd_nodes = get_bound_nodes(c, r, n, reverse=False)

    # assumes dict as ordered
    exp_nodes = [
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
        [24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35],
        [0, 1, 2, 12, 13, 14, 24, 25, 26],
        [9, 10, 11, 21, 22, 23, 33, 34, 35],
        [0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33],
        [2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35],
    ]

    for bnd_node_group, exp_node_group in zip(bnd_nodes.values(), exp_nodes):
        np.testing.assert_array_equal(bnd_node_group, exp_node_group)

    # reverse order (bottom swaps with top)
    bnd_nodes = get_bound_nodes(c, r, n, reverse=True)

    exp_nodes[0], exp_nodes[1] = exp_nodes[1], exp_nodes[0]
    for bnd_node_group, exp_node_group in zip(bnd_nodes.values(), exp_nodes):
        np.testing.assert_array_equal(bnd_node_group, exp_node_group)
