"""Build SBCmesh.mat / SBCmesh.dat from COMSOL SBC.mph (same logic as getMesh.m)."""

from __future__ import annotations

import argparse
from pathlib import Path

import mph
import numpy as np
import scipy.io as sio


def _build_mesh_struct(mph_path: Path) -> dict:
    client = mph.start()
    model = client.load(str(mph_path))
    m = model.java.component("comp1").mesh("mesh1")

    vertex = np.array(m.getVertex(), dtype=float).T
    nbr_vertex = vertex.shape[0]

    tet_raw = np.array(m.getElem("tet"), dtype=int).T
    nbr_tet = tet_raw.shape[0]
    domain_of_tet = np.array(m.getElemEntity("tet"), dtype=int).reshape(-1, 1)
    tet = np.sort(tet_raw, axis=1) + 1

    el2no = tet.T
    rows_n1 = [0, 0, 0, 1, 1, 2]
    rows_n2 = [1, 2, 3, 2, 3, 3]
    n1 = el2no[rows_n1, :].reshape(-1, order="F")
    n2 = el2no[rows_n2, :].reshape(-1, order="F")
    el_ed2no_array = np.column_stack([n1, n2])
    edge, edge_of_tet_flat = np.unique(el_ed2no_array, axis=0, return_inverse=True)
    nbr_edge = edge.shape[0]
    edge_of_tet = edge_of_tet_flat.reshape(6, nbr_tet, order="F").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])

    tri_raw = np.array(m.getElem("tri"), dtype=int).T
    tri = np.sort(tri_raw, axis=1) + 1
    domain_of_tri = np.array(m.getElemEntity("tri"), dtype=int).reshape(-1, 1)
    nbr_tri = tri.shape[0]

    conn_of_tri = np.zeros((nbr_tri, 2), dtype=int)
    for i in range(nbr_tri):
        matches = np.where(np.all(el_face_array == tri[i], 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((14, 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[13] = [1, 0, 0]

    return {
        "NbrVertex": np.int32(nbr_vertex),
        "Vertex": vertex,
        "NbrTet": np.int32(nbr_tet),
        "DomainOfTet": domain_of_tet,
        "Tet": tet,
        "NbrEdge": np.int32(nbr_edge),
        "Edge": edge.astype(np.int32),
        "EdgeOfTet": edge_of_tet.astype(np.int32),
        "NbrTri": np.int32(nbr_tri),
        "Tri": tri.astype(np.int32),
        "DomainOfTri": domain_of_tri,
        "ConnOfTri": conn_of_tri.astype(np.int32),
        "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(" ".join(str(int(x) + 1) for x in row))

    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])}")

    norm_rows = [1, 2, 3, 4, 5, 14]
    lines.append("NormOfFace")
    lines.append(str(len(norm_rows)))
    for d in norm_rows:
        n = mesh["NormOfFace"][d - 1]
        lines.append(f"{d} {n[0]:.16g} {n[1]:.16g} {n[2]:.16g}")

    dat_path.parent.mkdir(parents=True, exist_ok=True)
    dat_path.write_text("\n".join(lines) + "\n", encoding="ascii")


def build_sbcmesh(mph_path: Path, out_dir: Path, deploy: Path | None = None) -> dict:
    mesh = _build_mesh_struct(mph_path)
    out_dir.mkdir(parents=True, exist_ok=True)

    mat_path = out_dir / "SBCmesh.mat"
    dat_path = out_dir / "SBCmesh.dat"
    sio.savemat(str(mat_path), {"mesh": mesh}, do_compression=True)
    _write_dat(mesh, dat_path)

    if deploy is not None:
        deploy.mkdir(parents=True, exist_ok=True)
        import shutil

        shutil.copy2(mat_path, deploy / "SBCmesh.mat")
        shutil.copy2(dat_path, deploy / "SBCmesh.dat")

    print(f"Wrote {mat_path}")
    print(f"Wrote {dat_path}")
    print(
        f"  NbrVertex={mesh['NbrVertex']}, NbrTet={mesh['NbrTet']}, "
        f"NbrEdge={mesh['NbrEdge']}, NbrTri={mesh['NbrTri']}"
    )
    print(f"  A COO estimate (NbrTet*36) = {int(mesh['NbrTet']) * 36}")
    if deploy is not None:
        print(f"Deployed to {deploy}")
    return mesh


def main() -> None:
    parser = argparse.ArgumentParser(description="Build SBCmesh from COMSOL mph")
    parser.add_argument("mph_file", type=Path)
    parser.add_argument("-o", "--output-dir", type=Path, default=None)
    parser.add_argument("--deploy", type=Path, default=None)
    args = parser.parse_args()
    out_dir = args.output_dir or args.mph_file.parent
    build_sbcmesh(args.mph_file, out_dir, deploy=args.deploy)


if __name__ == "__main__":
    main()