#
# This is the interface between the meshio package and our Discretization description
# May not work perfectly
#
from typing import Any, Dict, List, Optional, Tuple

import meshio
import numpy as np
from tqdm import tqdm

from .discretization import Discretization, Node
from .element.line2 import Line2
from .element.line3 import Line3
from .element.parse_element import create_element
from .element.quad4 import Quad4
from .element.tri3 import Tri3
from .element.tri6 import Tri6
from .nodeset import (
    LineNodeset,
    LineNodesetBuilder,
    PointNodeset,
    PointNodesetBuilder,
    SurfaceNodeset,
    SurfaceNodesetBuilder,
    VolumeNodeset,
    VolumeNodesetBuilder,
)

cell_nodes = {
    "vertex": 1,
    "line": 2,
    "line3": 3,
    "triangle": 3,
    "triangle6": 6,
    "tetra": 4,
    "tetra10": 10,
    "pyramid": 5,
    "hexahedron": 8,
    "hexahedron20": 20,
    "hexahedron27": 27,
    "wedge": 6,
    "quad": 4,
    "quad9": 9,
    "vertex": 1,
}

cell_disc_eles = {
    "vertex": "UNKNOWN",
    "line": "UNKNOWN",
    "line3": "UNKNOWN",
    "triangle": "UNKNOWN",
    "triangle6": "UNKNOWN",
    "tetra": "SOLID",
    "tetra10": "SOLID",
    "pyramid": "SOLID",
    "hexahedron": "SOLID",
    "hexahedron20": "SOLID",
    "hexahedron27": "SOLID",
    "wedge": "UNKNOWN",
    "quad": "UNKNOWN",
    "quad9": "UNKNOWN",
}

cell_disc_shape = {
    "vertex": "VERTEX1",
    "line": "LINE2",
    "line3": "LINE3",
    "triangle": "TRI3",
    "triangle6": "TRI6",
    "tetra": "TET4",
    "tetra10": "TET10",
    "pyramid": "PYRAMID5",
    "hexahedron": "HEX8",
    "hexahedron20": "HEX20",
    "hexahedron27": "HEX27",
    "wedge": "WEDGE6",
    "quad": "QUAD4",
    "quad9": "QUAD9",
}

_cell_data_id_names = ["medit:ref", "gmsh:geometrical"]

# invert cell_disc_shape
disc_shape_cell = {v: k for k, v in cell_disc_shape.items()}

cell_to_dim = {
    "vertex": 0,
    "line": 1,
    "line3": 1,
    "triangle": 2,
    "triangle6": 2,
    "tetra": 3,
    "tetra10": 3,
    "pyramid": 3,
    "hexahedron": 3,
    "hexahedron20": 3,
    "hexahedron27": 3,
    "wedge": 3,
    "quad": 2,
    "quad9": 2,
}

# meshio used vtk node ordering
# with this array, we reorder the nodes
ele_node_order_vtk2baci = {
    "VERTEX1": list(range(0, 1)),
    "LINE2": list(range(0, 2)),
    "LINE3": list(range(0, 2)),
    "TRI3": list(range(0, 3)),
    "TRI6": list(range(0, 6)),
    "TET4": list(range(0, 4)),
    "TET10": list(range(0, 10)),
    "PYRAMID5": list(range(0, 5)),
    "HEX8": list(range(0, 8)),
    "HEX20": [
        0,
        1,
        2,
        3,
        4,
        5,
        6,
        7,
        8,
        9,
        10,
        11,
        16,
        17,
        18,
        19,
        12,
        13,
        14,
        15,
    ],  # verify ?
    "HEX27": [
        0,
        1,
        2,
        3,
        4,
        5,
        6,
        7,
        8,
        9,
        10,
        11,
        12,
        13,
        14,
        15,
        16,
        17,
        18,
        19,
        21,
        25,
        24,
        26,
        23,
        22,
        20,
    ],
    "WEDGE6": list(range(0, 6)),
    "QUAD4": list(range(0, 4)),
    "QUAD9": list(range(0, 9)),
}


def mesh2Discretization(mesh: meshio.Mesh) -> Discretization:
    # create empty discretization
    disc = Discretization()

    # mesh.points -> disc.nodes
    # mesh.cells -> disc.elements[Element.FieldTypeStructure]
    # node_sets -> are stored in disc.nodes

    # create Nodes from points
    for i, coord in tqdm(enumerate(mesh.points), desc="Create nodes"):
        n = Node(coord)

        for key, v in mesh.point_data.items():
            n.data[key] = v[i]

        disc.nodes.append(n)

    # get the maximum element dimension, which is the dimension of the mesh
    maxdim = max([cell_to_dim[cell_block.type] for cell_block in mesh.cells])

    # create Elements from cells
    disc.elements.structure = []
    pointnsbuilder = PointNodesetBuilder()
    linensbuilder = LineNodesetBuilder()
    surfnsbuilder = SurfaceNodesetBuilder()
    volumensbuilder = VolumeNodesetBuilder()

    for cellgroupid, cellblock in enumerate(mesh.cells):
        if cellblock.type not in cell_nodes:
            raise Exception(
                "The cell type {0} is currently not implemented".format(cellblock.type)
            )

        for cellid, cell in enumerate(cellblock.data):
            # read nodes that belong to the cells
            nodes = []

            # in case of a vertex, cell is just an integer
            cell = np.array(cell).reshape((-1))

            for nodeids in cell:
                nodes.append(disc.nodes[nodeids])

            # cells with the same dimension as the mesh dimensions are normal cells
            # cells with a lower dimension are treated as surface, line or node
            # definitions
            element_dim = cell_to_dim[cellblock.type]
            if maxdim == element_dim:
                # this is a normal element
                ele = create_element(
                    cell_disc_eles[cellblock.type],
                    cell_disc_shape[cellblock.type],
                    # also need to reorder nodes in a few element types
                    [
                        nodes[i]
                        for i in ele_node_order_vtk2baci[
                            cell_disc_shape[cellblock.type]
                        ]
                    ],
                )
                disc.elements.structure.append(ele)

                # extract material info from cell data
                matid: Optional[int] = _get_material_from_cell_data(
                    mesh.cell_data, cellgroupid, cellid
                )
                if matid == None:
                    matid = 1
                ele.options["MAT"] = matid

                # extract cell data
                for key, data in mesh.cell_data.items():
                    ele.data[key] = data[cellgroupid][cellid]

                # explicitly append the cell group id (block id) to the cell data
                ele.data["GROUP_ID"] = cellgroupid

                # create surface-nodesets
                if maxdim == 2 and len(mesh.cell_data) > 0:
                    nsid: Optional[int] = _get_nodesetid_from_cell_data(
                        mesh.cell_data, cellgroupid, cellid
                    )

                    if nsid is not None:
                        for node in nodes:
                            surfnsbuilder.add(node, nsid)

            else:
                # this is a lower-dimensional element
                # treat as surface-, line- or point-nodeset definition

                # extract nodeset info from cell data
                nsid: Optional[int] = _get_nodesetid_from_cell_data(
                    mesh.cell_data, cellgroupid, cellid
                )

                if nsid is not None:
                    for node in nodes:
                        if element_dim == 0:
                            pointnsbuilder.add(node, nsid)
                        elif element_dim == 1:
                            linensbuilder.add(node, nsid)
                        elif element_dim == 2:
                            surfnsbuilder.add(node, nsid)

    # go through nodesets
    for name, nodes in mesh.point_sets.items():
        if "volume" in name:
            nsid = volumensbuilder.get_unused_id()
            for n in nodes:
                volumensbuilder.add(disc.nodes[n], nsid)
        if "surface" in name:
            nsid = surfnsbuilder.get_unused_id()
            for n in nodes:
                surfnsbuilder.add(disc.nodes[n], nsid)
        elif "line" in name:
            nsid = linensbuilder.get_unused_id()
            for n in nodes:
                linensbuilder.add(disc.nodes[n], nsid)
        elif "point" in name:
            nsid = pointnsbuilder.get_unused_id()
            for n in nodes:
                pointnsbuilder.add(disc.nodes[n], nsid)

    disc.volumenodesets = volumensbuilder.finalize()
    disc.pointnodesets = pointnsbuilder.finalize()
    disc.linenodesets = linensbuilder.finalize()
    disc.surfacenodesets = surfnsbuilder.finalize()

    disc.finalize()
    return disc


def discretization2mesh(dis: Discretization) -> meshio.Mesh:
    dis.compute_ids(zero_based=True)

    points = dis.get_node_coords()

    point_data = {}

    for i, n in enumerate(dis.nodes):
        for k, v in n.data.items():
            if k not in point_data:
                point_data[k] = np.zeros(
                    tuple([len(dis.nodes)] + list(np.array(v).shape))
                )

            point_data[k][i] = v

    cells = []
    cell_data = {}

    id_to_index = {}

    # store cells
    for eletype in dis.elements.values():
        for ele in eletype:
            celltype = disc_shape_cell[ele.shape]

            newgroup = False

            if len(cells) == 0 or cells[-1][0] != celltype:
                cell_nodes = []
                cells.append((celltype, cell_nodes))
                newgroup = True
            else:
                cell_nodes = cells[-1][1]

            cell_nodes.append([n.id for n in ele.nodes])
            id_to_index[ele.id] = (len(cells) - 1, len(cell_nodes) - 1)

    # go again through the elements to store element data
    for eletype in dis.elements.values():
        for ele in eletype:
            block_id, cell_id = id_to_index[ele.id]

            if "MAT" in ele.options:
                if "material" not in cell_data:
                    cell_data["material"] = _create_cell_variable(
                        cells, tuple([1]), int
                    )

                cell_data["material"][block_id][cell_id] = int(
                    ele.options["MAT"][0]
                    if _isiter(ele.options["MAT"])
                    else ele.options["MAT"]
                )

            for variable_name, value in ele.data.items():
                value_reshaped = np.array(value).reshape((-1))

                if variable_name not in cell_data:
                    cell_data[variable_name] = _create_cell_variable(
                        cells, value_reshaped.shape, value_reshaped.dtype
                    )
                cell_data[variable_name][block_id][cell_id] = value_reshaped

    mesh: meshio.Mesh = meshio.Mesh(
        points, cells, cell_data=cell_data, point_data=point_data
    )

    dis.reset()

    return mesh


def _create_cell_variable(
    cells: List[Tuple[str, List[int]]], shape: Tuple[int], dtype
) -> List[np.ndarray]:
    data = []

    for block in cells:
        new_shape = tuple([len(block[1])] + list(shape))
        if len(shape) == 1 and shape[0] == 1:
            new_shape = tuple([len(block[1])])
        data.append(np.zeros(new_shape, dtype=dtype))

    return data


def _get_material_from_cell_data(
    celldata: Dict[str, List[List[Any]]], cellgroupid: int, cellid: int
) -> Optional[int]:
    for name in _cell_data_id_names:
        if name in celldata:
            return int(celldata[name][cellgroupid][cellid])

    return None


def _get_nodesetid_from_cell_data(
    celldata: Dict[str, List[List[Any]]], cellgroupid: int, cellid: int
) -> Optional[int]:
    for name in _cell_data_id_names:
        if name in celldata:
            return int(celldata[name][cellgroupid][cellid])

    return None


def _isiter(list: Any) -> bool:
    try:
        iter(list)
    except TypeError:
        return False

    return True
