import gmsh
import numpy as np
import openseespy.opensees as ops
from collections import defaultdict
from rich import print

OPS_GMSH_ELE_TYPE = [1, 2, 3, 4, 5, 9, 10, 11, 12, 16, 17]


class Gmsh2OPS:
    """Generate OpenSees code from GMSH."""

    def __init__(self, ndm: int = 3, ndf: int = 3):
        self.ndm = ndm
        self.ndf = ndf

        self.gmsh_entities = defaultdict(dict)
        self.gmsh_dim_entity_tags = None
        self.gmsh_nodes = defaultdict(dict)
        self.gmsh_eles = defaultdict(dict)
        self.out_file = None
        self.out_type = None

    def set_output_file(self, filename: str = "src.tcl"):
        """
        Parameters:
        ------------
        filename: str, default = "src.tcl".
            The output file-path-name, must end with ``.tcl`` or ``.py``.
        """
        self.out_file = filename
        if filename.endswith(".tcl"):
            self.out_type = "tcl"
        elif filename.endswith(".py"):
            self.out_type = "py"
        else:
            raise ValueError("output_filename must end with .tcl or .py!")
        with open(self.out_file, "w+") as outf:
            outf.write(f"# This file was created by {self.__class__.__name__}\n\n")
            if self.out_type == "py":
                outf.write("import openseespy.opensees as ops\n\n")
                outf.write("ops.wipe()\n")
                outf.write(
                    f"ops.model('basic', '-ndm', {self.ndm}, '-ndf', {self.ndf})\n\n"
                )
            else:
                outf.write("wipe\n")
                outf.write(f"model basic -ndm {self.ndm} -ndf {self.ndf}\n\n")

    def read_gmsh_file(self, file_path: str):
        """
        Read an ``.msh`` file generated by ``GMSH``.

        Parameters
        -----------
        file_path: str, the file path.
        """
        with open(file_path) as inf:
            lines = [ln.strip() for ln in inf.readlines()]
        # Remove comments
        lines = [ln for ln in lines if not ln.startswith("**")]
        node_idx, ele_idx, entities_idx = [], [], []
        node_idx.append(lines.index("$Nodes"))
        node_idx.append(lines.index("$EndNodes"))
        ele_idx.append(lines.index("$Elements"))
        ele_idx.append(lines.index("$EndElements"))
        entities_idx.append(lines.index("$Entities"))
        entities_idx.append(lines.index("$EndEntities"))

        self.gmsh_entities = _retrieve_entities(lines, entities_idx)
        self.gmsh_dim_entity_tags = list(self.gmsh_entities.keys())
        self.gmsh_nodes = _retrieve_nodes(lines, node_idx)
        self.gmsh_eles = _retrieve_eles(lines, ele_idx)

    def read_gmsh_data(self, print_info: bool = True):
        """
        Read data from ``GMSH`` at runtime.

        Parameters
        -----------
        print_info: bool, default=True
            Print info.
        """
        all_nodal_tags, all_ele_tags = [], []
        self.gmsh_dim_entity_tags = gmsh.model.getEntities()
        for dim, etag in self.gmsh_dim_entity_tags:
            bound_dimtags = gmsh.model.getBoundary(
                dimTags=[(dim, etag)], oriented=False
            )
            self.gmsh_entities[(dim, etag)]["BoundTags"] = [
                data[1] for data in bound_dimtags
            ]
            # //
            nodeTags, nodeCoords, nodeParams = gmsh.model.mesh.getNodes(dim, etag)
            nodeCoords = np.reshape(nodeCoords, (-1, 3))
            for tag, coord in zip(nodeTags, nodeCoords):
                all_nodal_tags.append(int(tag))
                self.gmsh_nodes[(dim, etag)][int(tag)] = list(coord)
            # //
            elemTypes, elemTags, elemNodeTags = gmsh.model.mesh.getElements(dim, etag)
            ele_tags = [item for row in elemTags for item in row]
            for tag in ele_tags:
                ele_type, node_tags, dim, etag = gmsh.model.mesh.getElement(tag)
                node_tags = [int(i) for i in node_tags]
                node_tags = _reshape_ele_node_order(ele_type, node_tags)
                node_tags.append(ele_type)
                self.gmsh_eles[(dim, etag)][int(tag)] = node_tags
                all_ele_tags.append(int(tag))
        # print info
        if print_info:
            num_points, num_curves, num_surf, num_vol, n = 0, 0, 0, 0, 0
            for dim, etag in self.gmsh_dim_entity_tags:
                if dim == 0:
                    num_points += 1
                elif dim == 1:
                    num_curves += 1
                elif dim == 2:
                    num_surf += 1
                elif dim == 3:
                    num_vol += 1
                n += 1
            print(
                f"Info:: {n} Entities: {num_points} Point; "
                f"{num_curves} Curves; "
                f"{num_surf} Surfaces; "
                f"{num_vol} Volumes."
            )
            print(
                f"Info:: {len(all_nodal_tags)} Nodes; MaxNodeTag {max(all_nodal_tags)}; "
                f"MinNodeTag {min(all_nodal_tags)}."
            )
            print(
                f"Info:: {len(all_ele_tags)} Elements; MaxEleTag {max(all_ele_tags)}; "
                f"MinEleTag {min(all_ele_tags)}."
            )

    def create_node_cmds(self) -> list:
        """
        Create ``OpenSeesPy`` nodes at runtime.
        Returns
        ---------
        node_tags: list, a list containing openseespy node tags.
        """
        node_tags = []
        for key in self.gmsh_nodes.keys():
            for tag, coords in self.gmsh_nodes[key].items():
                ops.node(tag, *coords)
                node_tags.append(tag)
        return node_tags

    def write_node_file(self):
        """
        Write node commands file, Tcl or Python.
        """
        with open(self.out_file, "a+") as outf:
            outf.write("\n# Create node commands\n\n")
            for key in self.gmsh_nodes.keys():
                for tag, coords in self.gmsh_nodes[key].items():
                    if self.out_type == "tcl":
                        coords = " ".join(map(str, coords[: self.ndm]))
                        outf.write(f"node {tag} {coords}\n")
                    else:
                        outf.write(f"ops.node({tag}, *{coords[:self.ndm]})\n")

    def create_element_cmds(
        self,
        dim_entity_tags: list,
        ops_ele_type: str,
        ops_ele_args: list,
    ) -> list:
        """
        Create ``OpenSeesPy`` elements at runtime.

        Parameters
        -----------
        dim_entity_tags: list, the GMSH [(dim, entity tag), ...].
            A list of dimension and entity tag containing element information that
            will be converted to OpenSeesPy elements.
            If None, all entities will be converted.
        ops_ele_type: str, the OpenSeesPy element type to generate.
        ops_ele_args: list,
            Parameters except OpenSeesPy element tag and coordinate points.

        Returns
        ---------
        ele_tags: list, a list containing ops element tags.
        """
        ele_tags = []
        if dim_entity_tags is None:
            for _, ele_nodes in self.gmsh_eles.items():
                if ele_nodes:
                    for tag, ntags in ele_nodes.items():
                        ops.element(ops_ele_type, tag, *ntags[:-1], *ops_ele_args)
                        ele_tags.append(tag)
        else:
            entity_tags = np.atleast_2d(dim_entity_tags)
            for etag in entity_tags:
                etag = (int(etag[0]), int(etag[1]))
                for tag, ntags in self.gmsh_eles[etag].items():
                    ops.element(ops_ele_type, tag, *ntags[:-1], *ops_ele_args)
                    ele_tags.append(tag)
        return ele_tags

    def write_element_file(
        self,
        dim_entity_tags: list,
        ops_ele_type: str,
        ops_ele_args: list,
    ):
        """
        Write elements commands file, Tcl or Python.

        Parameters
        -----------
        dim_entity_tags: list, the GMSH [(dim, entity tag), ...].
            A list of dimension and entity tag containing element information that
            will be converted to OpenSeesPy elements.
            If None, all entities will be converted.
        ops_ele_type: str, the OpenSeesPy element type to generate.
        ops_ele_args: list,
            Parameters except OpenSeesPy element tag and coordinate points.
        Returns
        ---------
        ele_tags: list, a list containing ops element tags.
        """
        if dim_entity_tags is None:
            entity_tags = self.gmsh_dim_entity_tags
        else:
            entity_tags = np.atleast_2d(dim_entity_tags)
        with open(self.out_file, "a+") as outf:
            outf.write(f"\n# Create element commands, type={ops_ele_type}\n\n")
            for etag in entity_tags:
                etag = (int(etag[0]), int(etag[1]))
                for tag, ntags in self.gmsh_eles[etag].items():
                    if self.out_type == "tcl":
                        nodetags = " ".join(map(str, ntags[:-1]))
                        ele_args = " ".join(map(str, ops_ele_args))
                        outf.write(
                            f"element {ops_ele_type} {tag} {nodetags} {ele_args}\n"
                        )
                    else:
                        outf.write(
                            f"ops.element({ops_ele_type}, {tag}, *{ntags[:-1]}, *{ops_ele_args})\n"
                        )

    def create_fix_cmds(
        self,
        dim_entity_tags: list,
        dofs: list,
    ) -> list:
        """
        Create fix constraints for OpenSeesPy at runtime.

        Parameters
        -----------
        dim_entity_tags: list, the GMSH [(dim, entity tag), ...].
            A list of dimension and entity tag containing node information that
            will be converted to OpenSeesPy fix.
        dofs: list, degrees of freedom to be constrained.

        Returns
        ---------
        node_tags: list, a list containing openseespy fixed node tags.
        """
        entity_tags = np.atleast_2d(dim_entity_tags)
        fixed_tags = []
        for etag in entity_tags:
            etag = (int(etag[0]), int(etag[1]))
            for tag in self.gmsh_nodes[etag].keys():
                if tag not in fixed_tags:
                    ops.fix(tag, *dofs)
                    fixed_tags.append(tag)
        return fixed_tags

    def write_fix_file(
        self,
        dim_entity_tags: list,
        dofs: list,
    ):
        """
        Write node fix commands file, Tcl or Python.

        Parameters
        -----------
        dim_entity_tags: list, the GMSH [(dim, entity tag), ...].
            A list of dimension and entity tag containing node information that
            will be converted to OpenSeesPy fix.
        dofs: list, degrees of freedom to be constrained.

        Returns
        ---------
        node_tags: list, a list containing openseespy fixed node tags.
        """
        entity_tags = np.atleast_2d(dim_entity_tags)
        fixed_tags = []
        with open(self.out_file, "a+") as outf:
            outf.write("\n# Create fix commands\n\n")
            for etag in entity_tags:
                etag = (int(etag[0]), int(etag[1]))
                for tag in self.gmsh_nodes[etag].keys():
                    if tag not in fixed_tags:
                        if self.out_type == "tcl":
                            dofs_ = " ".join(map(str, dofs))
                            outf.write(f"fix {tag} {dofs_}\n")
                        else:
                            outf.write(f"ops.fix({tag}, *{dofs})\n")
                        fixed_tags.append(tag)

    def get_dim_entity_tags(self, dim: int = None) -> list:
        """
        Get dim_entity_tags from GMSH.
        Parameters
        ----------
        dim: int, optional, default None
            The dimension tag. If None, all entities will be returned.
        Returns
        -------
        dim_entity_tags: list, the GMSH [(dim, entity tag), ...].
        """
        if dim is None:
            return self.gmsh_dim_entity_tags
        else:
            dim_entity_tags = []
            for dimi, etag in self.gmsh_dim_entity_tags:
                if dimi == dim:
                    dim_entity_tags.append((dimi, etag))
            return dim_entity_tags

    def get_boundary_dim_tags(
        self, dim_entity_tags: list, include_self: bool = False
    ) -> list:
        """
        Get all boundaries of the GMSH entities, including corner points.

        Parameters
        ----------
        dim_entity_tags: list, the GMSH [(dim, entity tag), ...].
            A list of GMSH dimension and entity tags.
        include_self: bool, default False
            If True, the output contains itself, which is dim_entity_tags.

        Returns
        -------
        Boundary dimension and entity tags list.
        """
        dim_entity_tags = np.atleast_2d(dim_entity_tags)
        boundary_dimtags = []
        if include_self:
            for dim, etag in dim_entity_tags:
                boundary_dimtags.append((dim, etag))
        _get_boundary_dim_tags(boundary_dimtags, dim_entity_tags, self.gmsh_entities)
        return list(set(boundary_dimtags))


def _get_boundary_dim_tags(boundary_dimtags, dim_entity_tags, entites):
    for dim_etag in dim_entity_tags:
        dim, etag = int(dim_etag[0]), int(dim_etag[1])
        if dim > 0:
            bound_etags = entites[(dim, etag)]["BoundTags"]
            bound_dimtags = [(dim - 1, abs(data)) for data in bound_etags]
            boundary_dimtags.extend(bound_dimtags)
            _get_boundary_dim_tags(boundary_dimtags, bound_dimtags, entites)


# def _get_boundary_dim_tags(boundary_dimtags, dim_entity_tags):
#     for etag in dim_entity_tags:
#         etag = [(int(etag[0]), int(etag[1]))]
#         bound_dimtags = gmsh.model.getBoundary(dimTags=etag, oriented=False)
#         boundary_dimtags.extend(bound_dimtags)
#         for eetag in bound_dimtags:
#             if eetag[0] > 0:
#                 _get_boundary_dim_tags(boundary_dimtags, [eetag])


def _retrieve_entities(lines, entities_idx):
    entities = defaultdict(dict)
    idx = entities_idx[0] + 1
    contents = [int(data) for data in lines[idx].split(" ")]
    num_points, num_curves, num_surf, num_vol = contents
    n = num_points + num_curves + num_surf + num_vol
    print(
        f"Info:: {n} Entities: {num_points} Point; "
        f"{num_curves} Curves; "
        f"{num_surf} Surfaces; "
        f"{num_vol} Volumes."
    )
    idx = entities_idx[0] + 2
    while idx < entities_idx[1]:
        for i in range(num_points):
            contents = lines[idx].split(" ")
            tag = int(contents[0])
            entities[(0, tag)]["Coord"] = [float(data) for data in contents[1:4]]
            entities[(0, tag)]["numPhysicalTags"] = int(contents[4])
            entities[(0, tag)]["physicalTags"] = [
                int(contents[5 + iii]) for iii in range(int(contents[4]))
            ]
            entities[(0, tag)]["numBound"] = 0
            entities[(0, tag)]["BoundTags"] = []
            idx += 1
        for i in range(num_curves):
            contents = lines[idx].split(" ")
            tag = int(contents[0])
            entities[(1, tag)]["CoordBoundary"] = [
                float(data) for data in contents[1:7]
            ]
            entities[(1, tag)]["numPhysicalTags"] = int(contents[7])
            entities[(1, tag)]["physicalTags"] = [
                int(contents[8 + iii]) for iii in range(int(contents[7]))
            ]
            numBounds = int(contents[8 + int(contents[7])])
            entities[(1, tag)]["numBound"] = numBounds
            idxxxx = 9 + int(contents[7])
            entities[(1, tag)]["BoundTags"] = [
                int(contents[idxxxx + iii]) for iii in range(numBounds)
            ]
            idx += 1
        for i in range(num_surf):
            contents = lines[idx].split(" ")
            tag = int(contents[0])
            entities[(2, tag)]["CoordBoundary"] = [
                float(data) for data in contents[1:7]
            ]
            entities[(2, tag)]["numPhysicalTags"] = int(contents[7])
            entities[(2, tag)]["physicalTags"] = [
                int(contents[8 + iii]) for iii in range(int(contents[7]))
            ]
            numBounds = int(contents[8 + int(contents[7])])
            entities[(2, tag)]["numBound"] = numBounds
            idxxxx = 9 + int(contents[7])
            entities[(2, tag)]["BoundTags"] = [
                int(contents[idxxxx + iii]) for iii in range(numBounds)
            ]
            idx += 1
        for i in range(num_vol):
            contents = lines[idx].split(" ")
            tag = int(contents[0])
            entities[(3, tag)]["CoordBoundary"] = [
                float(data) for data in contents[1:7]
            ]
            entities[(3, tag)]["numPhysicalTags"] = int(contents[7])
            entities[(3, tag)]["physicalTags"] = [
                int(contents[8 + iii]) for iii in range(int(contents[7]))
            ]
            numBounds = int(contents[8 + int(contents[7])])
            entities[(3, tag)]["numBound"] = numBounds
            idxxxx = 9 + int(contents[7])
            entities[(3, tag)]["BoundTags"] = [
                int(contents[idxxxx + iii]) for iii in range(numBounds)
            ]
            idx += 1
    return entities


def _retrieve_nodes(lines, node_idx):
    nodes = defaultdict(dict)
    idx = node_idx[0] + 1
    contents = [int(data) for data in lines[idx].split(" ")]
    _, num_nodes, min_node_tag, max_node_tag = contents
    print(
        f"Info:: {num_nodes} Nodes; MaxNodeTag {max_node_tag}; MinNodeTag {min_node_tag}."
    )
    idx = node_idx[0] + 2
    while idx < node_idx[1]:
        contents = [int(data) for data in lines[idx].split(" ")]
        dim, etag, parametric, num_nodes_inblock = contents
        for i in range(num_nodes_inblock):
            tag = int(lines[idx + i + 1])
            coords = [
                float(data)
                for data in lines[idx + num_nodes_inblock + i + 1].split(" ")
            ]
            nodes[(dim, etag)][tag] = coords
        idx += 2 * num_nodes_inblock + 1
    return nodes


def _retrieve_eles(lines, ele_idx):
    eles = defaultdict(dict)
    idx = ele_idx[0] + 1
    contents = [int(data) for data in lines[idx].split(" ")]
    _, num_eles, min_ele_tag, max_ele_tag = contents
    print(
        f"Info:: {num_eles} Elements; MaxEleTag {max_ele_tag}; MinEleTag {min_ele_tag}."
    )
    idx = ele_idx[0] + 2
    while idx < ele_idx[1]:
        contents = [int(data) for data in lines[idx].split(" ")]
        dim, etag, ele_type, num_eles_inblock = contents
        if ele_type in OPS_GMSH_ELE_TYPE:
            for i in range(num_eles_inblock):
                info = [int(data) for data in lines[idx + i + 1].split(" ")]
                tag = info[0]
                node_tags = _reshape_ele_node_order(ele_type, info[1:])
                node_tags += [ele_type]
                eles[(dim, etag)][tag] = node_tags
        idx += num_eles_inblock + 1
    return eles


def _reshape_ele_node_order(ele_type, node_tags):
    if ele_type == 4:
        tags = _reshape_tet_n4(node_tags)
    elif ele_type == 11:
        tags = _reshape_tet_n10(node_tags)
    elif ele_type == 5:
        tags = _reshape_hex_n8(node_tags)
    elif ele_type == 17:
        tags = _reshape_hex_n20(node_tags)
    elif ele_type == 12:
        tags = _reshape_hex_n27(node_tags)
    else:
        tags = node_tags
    return tags


def _reshape_tet_n4(node_tags):
    tags = [node_tags[3], node_tags[1], node_tags[0], node_tags[2]]
    return tags


def _reshape_tet_n10(node_tags):
    tags = [
        node_tags[3],
        node_tags[1],
        node_tags[0],
        node_tags[2],
        node_tags[9],
        node_tags[4],
        node_tags[7],
        node_tags[8],
        node_tags[5],
        node_tags[6],
    ]
    return tags


def _reshape_hex_n8(node_tags):
    tags = [
        node_tags[4],
        node_tags[5],
        node_tags[1],
        node_tags[0],
        node_tags[7],
        node_tags[6],
        node_tags[2],
        node_tags[3],
    ]
    return tags


def _reshape_hex_n20(node_tags):
    tags = [
        node_tags[4],
        node_tags[5],
        node_tags[1],
        node_tags[0],
        node_tags[7],
        node_tags[6],
        node_tags[2],
        node_tags[3],
        node_tags[16],
        node_tags[22],
        node_tags[8],
        node_tags[10],
        node_tags[19],
        node_tags[14],
        node_tags[13],
        node_tags[15],
        node_tags[17],
        node_tags[18],
        node_tags[11],
        node_tags[9],
    ]
    return tags


def _reshape_hex_n27(node_tags):
    tags = [
        node_tags[4],
        node_tags[5],
        node_tags[1],
        node_tags[0],
        node_tags[7],
        node_tags[6],
        node_tags[2],
        node_tags[3],
        node_tags[16],
        node_tags[22],
        node_tags[8],
        node_tags[10],
        node_tags[19],
        node_tags[14],
        node_tags[13],
        node_tags[15],
        node_tags[17],
        node_tags[18],
        node_tags[11],
        node_tags[9],
        node_tags[23],
        node_tags[20],
        node_tags[24],
        node_tags[22],
        node_tags[25],
        node_tags[21],
    ]
    return tags