#!/usr/bin/env python3
"""Export MeshData2x.mat (Fem4) or MeshData2.mat (Fem2) to OpticsFEM Portmesh_*.dat."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np

try:
    import scipy.io as sio
except ImportError:
    sio = None


def _load_mat_mesh(mat_path: Path) -> object:
    """Load Mesh struct from .mat (v7 or v7.3)."""
    if sio is not None:
        try:
            data = sio.loadmat(str(mat_path), squeeze_me=True, struct_as_record=False)
            return data["Mesh"]
        except NotImplementedError:
            pass
    import h5py

    class MeshView:
        pass

    mesh = MeshView()
    with h5py.File(str(mat_path), "r") as f:
        m = f["Mesh"]
        mesh.Nodes = np.asarray(m["Nodes"]).T
        mesh.Elements = np.asarray(m["Elements"]).T
        mesh.Domains = np.asarray(m["Domains"]).reshape(-1)
        mesh.Faces = np.asarray(m["Faces"]).T
        mesh.FacesIndex = np.asarray(m["FacesIndex"]).reshape(-1)
    return mesh


def build_mesh_from_mat(mat_path: Path) -> dict:
    mesh = _load_mat_mesh(mat_path)
    vertex = np.asarray(mesh.Nodes, dtype=float)
    tet_raw = np.asarray(mesh.Elements, dtype=int)
    domain_of_tet = np.asarray(mesh.Domains, dtype=int).reshape(-1, 1)
    tri_raw = np.asarray(mesh.Faces, dtype=int)
    domain_of_tri = np.asarray(mesh.FacesIndex, dtype=int).reshape(-1, 1)

    nbr_vertex = vertex.shape[0]
    nbr_tet = tet_raw.shape[0]
    # Keep MATLAB Elements vertex order — sorting breaks Nedelec local edge BF mapping.
    tet = np.asarray(tet_raw, dtype=np.int32)

    # Match MATLAB PhysicMatrixAssembly: unique(...,'rows') sorts edge pairs
    # lexicographically; EdgesOfElements maps to that sorted order (not first-seen).
    el2no = tet.T
    n1 = el2no[[0, 0, 0, 1, 1, 2], :]
    n2 = el2no[[1, 2, 3, 2, 3, 3], :]
    el_ed2no = np.column_stack([n1.reshape(-1, order="F"), n2.reshape(-1, order="F")])
    edge, ic = np.unique(el_ed2no, axis=0, return_inverse=True)
    edge = edge.astype(np.int32)
    edge_of_tet = (ic.reshape(6, -1, order="F").T + 1).astype(np.int32)
    tri = np.asarray(tri_raw, dtype=np.int32)

    el2no = tet.T
    rows_f1 = [0, 0, 0, 1]
    rows_f2 = [1, 1, 2, 2]
    rows_f3 = [2, 3, 3, 3]
    f1 = el2no[rows_f1, :].reshape(-1, order="F")
    f2 = el2no[rows_f2, :].reshape(-1, order="F")
    f3 = el2no[rows_f3, :].reshape(-1, order="F")
    el_face_array = np.column_stack([f1, f2, f3])

    conn_of_tri = np.zeros((tri.shape[0], 2), dtype=np.int32)
    for i in range(tri.shape[0]):
        tri_sorted = np.sort(tri[i])
        matches = np.where(np.all(np.sort(el_face_array, axis=1) == tri_sorted, axis=1))[0]
        if matches.size == 0:
            raise RuntimeError(f"boundary tri {i + 1} not found in tet faces")
        index = int(matches[0]) + 1
        conn_of_tri[i, 0] = (index - 1) // 4 + 1
        conn_of_tri[i, 1] = index - (conn_of_tri[i, 0] - 1) * 4

    norm_of_face = np.zeros((24, 3), dtype=float)
    norm_of_face[0] = [-1, 0, 0]
    norm_of_face[1] = [0, -1, 0]
    norm_of_face[2] = [0, 0, -1]
    norm_of_face[3] = [0, 0, 1]
    norm_of_face[4] = [0, 1, 0]
    norm_of_face[23] = [1, 0, 0]

    return {
        "NbrVertex": np.int32(nbr_vertex),
        "Vertex": vertex,
        "NbrTet": np.int32(nbr_tet),
        "DomainOfTet": domain_of_tet,
        "Tet": tet.astype(np.int32),
        "NbrEdge": np.int32(edge.shape[0]),
        "Edge": edge,
        "EdgeOfTet": edge_of_tet,
        "NbrTri": np.int32(tri.shape[0]),
        "Tri": tri.astype(np.int32),
        "DomainOfTri": domain_of_tri,
        "ConnOfTri": conn_of_tri,
        "NormOfFace": norm_of_face,
    }


def write_dat(mesh: dict, dat_path: Path) -> None:
    lines: list[str] = []
    lines.append("NbrVertex")
    lines.append(str(int(mesh["NbrVertex"])))
    lines.append("Vertex")
    for row in mesh["Vertex"]:
        lines.append(f"{row[0]:.16g} {row[1]:.16g} {row[2]:.16g}")

    lines.append("NbrTet")
    lines.append(str(int(mesh["NbrTet"])))
    lines.append("Tet")
    for row in mesh["Tet"]:
        lines.append(f"{int(row[0])} {int(row[1])} {int(row[2])} {int(row[3])}")

    lines.append("DomainOfTet")
    for v in mesh["DomainOfTet"].reshape(-1):
        lines.append(str(int(v)))

    lines.append("NbrEdge")
    lines.append(str(int(mesh["NbrEdge"])))
    lines.append("Edge")
    for row in mesh["Edge"]:
        lines.append(f"{int(row[0])} {int(row[1])}")

    lines.append("EdgeOfTet")
    for row in mesh["EdgeOfTet"]:
        lines.append(f"{int(row[0])} {int(row[1])} {int(row[2])} {int(row[3])} {int(row[4])} {int(row[5])}")

    lines.append("NbrTri")
    lines.append(str(int(mesh["NbrTri"])))
    lines.append("Tri")
    for row in mesh["Tri"]:
        lines.append(f"{int(row[0])} {int(row[1])} {int(row[2])}")

    lines.append("DomainOfTri")
    for v in mesh["DomainOfTri"].reshape(-1):
        lines.append(str(int(v)))

    lines.append("ConnOfTri")
    for row in mesh["ConnOfTri"]:
        lines.append(f"{int(row[0])} {int(row[1])}")

    lines.append("NormOfFace")
    active_norms = [
        (idx + 1, row)
        for idx, row in enumerate(mesh["NormOfFace"])
        if np.linalg.norm(row) > 0.0
    ]
    lines.append(str(len(active_norms)))
    for domain_id, row in active_norms:
        lines.append(f"{int(domain_id)} {row[0]:.16g} {row[1]:.16g} {row[2]:.16g}")

    dat_path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--mat",
        type=Path,
        default=Path(__file__).resolve().parents[2]
        / "三维matlab代码/2023-2-端口激励问题(四面体网格)/MeshData2x.mat",
    )
    parser.add_argument(
        "--out",
        type=Path,
        default=Path(__file__).resolve().parents[1] / "port/Release/Portmesh_fem4.dat",
    )
    args = parser.parse_args()
    # If out still defaults to fem4 but mat is MeshData2, suggest fem2 name
    if args.out.name == "Portmesh_fem4.dat" and "MeshData2.mat" in args.mat.name and "MeshData2x" not in args.mat.name:
        args.out = args.out.with_name("Portmesh_fem2.dat")
    mesh = build_mesh_from_mat(args.mat)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    write_dat(mesh, args.out)
    print(f"Wrote {args.out}  (V={mesh['NbrVertex']}, T={mesh['NbrTet']}, E={mesh['NbrEdge']})")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
