from __future__ import annotations

import os
from typing import TYPE_CHECKING, TypedDict, cast

import gmsh
import numpy as np


if TYPE_CHECKING:
    from collections.abc import Sequence


OUTPUT_TEMPLATE = """# GENERATED by gmsh_interop/contrib/gmsh-node-tuples.py
# GMSH_VERSION: %s
# DO NOT EDIT

from __future__ import annotations

from typing import TYPE_CHECKING, TypedDict


if TYPE_CHECKING:
    from collections.abc import Sequence


class ElementInfo(TypedDict):
    element_name: str
    '''Name of the element in GMSH.'''
    element_type: int
    '''An integer denoting the type of the element.'''
    node_tuples: Sequence[tuple[int, ...]]
    '''Index tuples that describe the ordering of the nodes.'''


triangle_data: dict[int, ElementInfo] = %s

tetrahedron_data: dict[int, ElementInfo] = %s

quadrangle_data: dict[int, ElementInfo] = %s

hexahedron_data: dict[int, ElementInfo] = %s
"""


TRIANGLE_ELEMENTS = {
        "MSH_TRI_3": 2,
        "MSH_TRI_6": 9,
        "MSH_TRI_10": 21,
        "MSH_TRI_15": 23,
        "MSH_TRI_21": 25,
        "MSH_TRI_28": 42,
        "MSH_TRI_36": 43,
        "MSH_TRI_45": 44,
        "MSH_TRI_55": 45,
        "MSH_TRI_66": 46,
        }


TETRAHEDRON_ELEMENTS = {
        "MSH_TET_4": 4,
        "MSH_TET_10": 11,
        "MSH_TET_20": 29,
        "MSH_TET_35": 30,
        "MSH_TET_56": 31,
        "MSH_TET_84": 71,
        "MSH_TET_120": 72,
        "MSH_TET_165": 73,
        "MSH_TET_220": 74,
        "MSH_TET_286": 75,
        }


QUADRANGLE_ELEMENTS = {
        "MSH_QUA_4": 3,
        "MSH_QUA_9": 10,
        "MSH_QUA_16": 36,
        "MSH_QUA_25": 37,
        "MSH_QUA_36": 38,
        "MSH_QUA_49": 47,
        "MSH_QUA_64": 48,
        "MSH_QUA_81": 49,
        "MSH_QUA_100": 50,
        "MSH_QUA_121": 51,
        }


HEXAHEHEDRON_ELEMENTS = {
        "MSH_HEX_8": 5,
        "MSH_HEX_27": 12,
        "MSH_HEX_64": 92,
        "MSH_HEX_125": 93,
        "MSH_HEX_216": 94,
        "MSH_HEX_343": 95,
        "MSH_HEX_512": 96,
        "MSH_HEX_729": 97,
        "MSH_HEX_1000": 98,
        }


class ElementInfo(TypedDict):
    node_tuples: Sequence[tuple[int, ...]]
    element_type: int
    element_name: str


def generate_node_tuples_from_gmsh(
        eltype: int,
        eldim: int,
        elvertices: int,
        domain: str = "unit") -> Sequence[tuple[int, ...]]:
    # {{{ get element

    _name, dim, order, nnodes, nodes, nvertices = cast(
        "tuple[str, int, int, int, np.typing.NDArray[np.floating], int]",
        gmsh.model.mesh.getElementProperties(eltype))

    assert dim == eldim
    assert nvertices == elvertices

    if domain == "unit":
        pass
    elif domain == "biunit":
        nodes = (1.0 + nodes) / 2.0
    else:
        raise ValueError(f"unknown domain: '{domain}'")

    nodes = nodes.reshape(nnodes, dim) * order

    # }}}

    return [tuple(int(x) for x in node) for node in nodes.astype(np.int32)]


def generate_node_tuples(filename: str | None, *, overwrite: bool = False) -> int:
    if not overwrite and filename is not None and os.path.exists(filename):
        print(f"ERROR: File already exists (use --force): '{filename}'")
        return 1

    tri_data: dict[int, ElementInfo] = {}
    tet_data: dict[int, ElementInfo] = {}
    qua_data: dict[int, ElementInfo] = {}
    hex_data: dict[int, ElementInfo] = {}

    gmsh.initialize()
    gmsh.option.setNumber("General.Terminal", 1)

    for order, (name, eltype) in enumerate(TRIANGLE_ELEMENTS.items()):
        node_tuples = generate_node_tuples_from_gmsh(eltype, 2, 3)
        tri_data[order + 1] = {
                "node_tuples": node_tuples,
                "element_type": eltype,
                "element_name": name,
                }

    for order, (name, eltype) in enumerate(TETRAHEDRON_ELEMENTS.items()):
        node_tuples = generate_node_tuples_from_gmsh(eltype, 3, 4)
        tet_data[order + 1] = {
                "node_tuples": node_tuples,
                "element_type": eltype,
                "element_name": name,
                }

    for order, (name, eltype) in enumerate(QUADRANGLE_ELEMENTS.items()):
        node_tuples = generate_node_tuples_from_gmsh(eltype, 2, 4, domain="biunit")
        qua_data[order + 1] = {
                "node_tuples": node_tuples,
                "element_type": eltype,
                "element_name": name,
                }

    for order, (name, eltype) in enumerate(HEXAHEHEDRON_ELEMENTS.items()):
        node_tuples = generate_node_tuples_from_gmsh(eltype, 3, 8, domain="biunit")
        hex_data[order + 1] = {
                "node_tuples": node_tuples,
                "element_type": eltype,
                "element_name": name,
                }

    gmsh.finalize()

    from pprint import pformat

    txt = (OUTPUT_TEMPLATE % (
        gmsh.GMSH_API_VERSION,
        pformat(tri_data, width=80),
        pformat(tet_data, width=80),
        pformat(qua_data, width=80),
        pformat(hex_data, width=80),
        )).replace('"', "")

    if filename is None:
        print(txt)
    else:
        with open(filename, "w") as fd:
            fd.write(txt)

    return 0


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("filename", nargs="?", default=None)
    parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        help="Overwrite existing files",
    )
    args = parser.parse_args()

    raise SystemExit(generate_node_tuples(args.filename, overwrite=args.force))
