#!/usr/bin/env python3
"""Export singlePBC_mesh.mat / doublePBC_mesh.mat to OpticsFEM .dat (all NormOfFace rows)."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import scipy.io as sio


def _as_rows(arr, ncol: int) -> list[list[int]]:
    if arr.ndim == 1:
        return [list(arr)]
    if arr.shape[0] == ncol and arr.shape[1] != ncol:
        arr = arr.T
    return [list(row) for row in arr]


def export_mat_to_dat(mat_path: Path, dat_path: Path) -> None:
    data = sio.loadmat(str(mat_path), squeeze_me=True, struct_as_record=False)
    mesh = data["mesh"]

    lines: list[str] = []
    lines.append("NbrVertex")
    lines.append(str(int(mesh.NbrVertex)))
    lines.append("Vertex")
    for row in _as_rows(mesh.Vertex, 3):
        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 _as_rows(mesh.Tet, 4):
        lines.append(f"{int(row[0])} {int(row[1])} {int(row[2])} {int(row[3])}")

    lines.append("DomainOfTet")
    for v in mesh.DomainOfTet.flatten():
        lines.append(str(int(v)))

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

    lines.append("EdgeOfTet")
    eot_rows = _as_rows(mesh.EdgeOfTet, 6)
    zero_based = any(int(v) == 0 for row in eot_rows for v in row)
    for row in eot_rows:
        vals = [int(x) + (1 if zero_based else 0) for x in row]
        lines.append(" ".join(str(x) for x in vals))

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

    lines.append("DomainOfTri")
    for v in mesh.DomainOfTri.flatten():
        lines.append(str(int(v)))

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

    norm_face = np.asarray(mesh.NormOfFace, dtype=float)
    n_norm = norm_face.shape[0]
    lines.append("NormOfFace")
    lines.append(str(n_norm))
    for d in range(1, n_norm + 1):
        n = norm_face[d - 1]
        nx, ny, nz = float(n[0]), float(n[1]), float(n[2])
        if abs(nx) + abs(ny) + abs(nz) < 1e-30:
            nx, ny, nz = 0.0, 0.0, 1.0
        lines.append(f"{d} {nx:.16g} {ny:.16g} {nz:.16g}")

    dat_path.parent.mkdir(parents=True, exist_ok=True)
    dat_path.write_text("\n".join(lines) + "\n", encoding="ascii")
    print(f"Wrote {dat_path} ({len(lines)} lines)")
    print(
        f"  NbrVertex={mesh.NbrVertex}, NbrEdge={mesh.NbrEdge}, "
        f"NbrTri={mesh.NbrTri}, NormOfFace={n_norm}"
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("mat_file", type=Path)
    parser.add_argument("-o", "--output", type=Path, default=None)
    args = parser.parse_args()
    out = args.output or args.mat_file.with_suffix(".dat")
    export_mat_to_dat(args.mat_file, out)


if __name__ == "__main__":
    main()