"""Plot normE on a plane slice (x / y / z = const), COMSOL-style."""
from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.tri import Triangulation


def load_mesh(path: Path):
    lines = path.read_text().splitlines()
    i = 0

    def need(tag: str):
        nonlocal i
        if lines[i] != tag:
            raise ValueError(f"expected {tag!r}, got {lines[i]!r}")
        i += 1

    need("NbrVertex")
    nv = int(lines[i])
    i += 1
    need("Vertex")
    V = np.array([list(map(float, lines[i + k].split())) for k in range(nv)])
    i += nv
    need("NbrTet")
    nt = int(lines[i])
    i += 1
    need("Tet")
    T = np.array([list(map(int, lines[i + k].split())) for k in range(nt)], int) - 1
    i += nt
    need("DomainOfTet")
    i += nt
    need("NbrEdge")
    ne = int(lines[i])
    i += 1
    need("Edge")
    i += ne
    need("EdgeOfTet")
    i += nt
    need("NbrTri")
    ntri = int(lines[i])
    i += 1
    need("Tri")
    Tri = np.array([list(map(int, lines[i + k].split())) for k in range(ntri)], int) - 1
    i += ntri
    need("DomainOfTri")
    dom_tri = np.array([int(lines[i + k]) for k in range(ntri)])
    i += ntri
    if i < len(lines) and lines[i] == "ConnOfTri":
        i += 1 + ntri
    return V, T, Tri, dom_tri


def load_scalar_field(path: Path) -> np.ndarray:
    vals = []
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line:
            continue
        if line.startswith("//"):
            line = line[2:].strip()
        vals.append(float(line.split()[0]))
    return np.array(vals)


def axis_index(axis: str) -> int:
    return {"x": 0, "y": 1, "z": 2}[axis.lower()]


def other_axes(axis: str) -> tuple[int, int]:
    if axis == "x":
        return 1, 2
    if axis == "y":
        return 0, 2
    return 0, 1


def triangulate_polygon(pts: list[np.ndarray]) -> list[tuple[int, int, int]]:
    if len(pts) < 3:
        return []
    if len(pts) == 3:
        return [(0, 1, 2)]
    # fan triangulation for convex quad
    return [(0, 1, 2), (0, 2, 3)]


def tet_plane_poly(verts: np.ndarray, axis: int, coord: float, tol: float = 1e-9):
    """Return polygon vertices (3 or 4) where tet intersects plane, or None."""
    c = verts[:, axis]
    cmin, cmax = c.min(), c.max()
    if coord < cmin - tol or coord > cmax + tol:
        return None

    if np.all(np.abs(c - coord) < 1e-6):
        # face on plane: pick 3 verts on plane (rare duplicate handling)
        idx = np.where(np.abs(c - coord) < 1e-6)[0]
        if len(idx) >= 3:
            return [verts[idx[0]], verts[idx[1]], verts[idx[2]]]
        return None

    edges = [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]
    pts = []
    for a, b in edges:
        ca, cb = c[a], c[b]
        if abs(ca - coord) < 1e-9:
            pts.append(verts[a])
        elif abs(cb - coord) < 1e-9:
            pts.append(verts[b])
        elif (ca - coord) * (cb - coord) < 0:
            t = (coord - ca) / (cb - ca)
            pts.append(verts[a] + t * (verts[b] - verts[a]))

    if len(pts) < 3:
        return None

    # remove near-duplicates
    uniq = []
    for p in pts:
        if not any(np.linalg.norm(p - q) < 1e-8 for q in uniq):
            uniq.append(p)
    if len(uniq) < 3:
        return None
    if len(uniq) > 4:
        uniq = uniq[:4]
    return uniq


def interp_scalar(p: np.ndarray, tet_idx: np.ndarray, V: np.ndarray, field: np.ndarray) -> float:
    v = V[tet_idx]
    f = field[tet_idx]
    A = np.column_stack([v[1] - v[0], v[2] - v[0], v[3] - v[0]])
    w123 = np.linalg.solve(A, p - v[0])
    w0 = 1.0 - np.sum(w123)
    w = np.array([w0, w123[0], w123[1], w123[2]])
    return float(np.dot(w, f))


def build_slice_tris(V: np.ndarray, T: np.ndarray, field: np.ndarray, axis: str, coord: float):
    ai, bi = other_axes(axis)
    ax = axis_index(axis)

    polys = []
    values = []
    for tet in T:
        verts = V[tet]
        poly = tet_plane_poly(verts, ax, coord)
        if poly is None:
            continue
        poly_arr = np.array(poly)
        local_vals = [interp_scalar(p, tet, V, field) for p in poly_arr]
        tris = triangulate_polygon(poly)
        for t0, t1, t2 in tris:
            polys.append(poly_arr[[t0, t1, t2]])
            values.append([local_vals[t0], local_vals[t1], local_vals[t2]])

    if not polys:
        return None, None, None

    polys = np.array(polys)
    values = np.array(values)
    xy = polys.reshape(-1, 3)[:, [ai, bi]]
    zval = values.reshape(-1)
    tri = Triangulation(xy[:, 0], xy[:, 1], np.arange(len(zval)).reshape(-1, 3))
    return tri, zval, (ai, bi)


def build_surface_tris(V: np.ndarray, Tri: np.ndarray, field: np.ndarray, axis: str, coord: float, tol: float):
    ai, bi = other_axes(axis)
    ax = axis_index(axis)

    xs, ys, zs, tris = [], [], [], []
    node = 0
    for tri in Tri:
        verts = V[tri]
        if not np.all(np.abs(verts[:, ax] - coord) < tol):
            continue
        xs.extend(verts[:, ai])
        ys.extend(verts[:, bi])
        zs.extend(field[tri])
        tris.append([node, node + 1, node + 2])
        node += 3

    if not tris:
        return None, None, None

    xy = np.column_stack([xs, ys])
    tri = Triangulation(xy[:, 0], xy[:, 1], np.array(tris))
    return tri, np.array(zs), (ai, bi)


def plot_slice(tri, zval, plane_axes, out_path: Path, coord: float, axis: str, title: str):
    ai, bi = plane_axes
    names = ["x", "y", "z"]
    fig, ax = plt.subplots(figsize=(7.5, 6.5), dpi=150)
    tpc = ax.tripcolor(tri, zval, shading="gouraud", cmap="jet", edgecolors="k", linewidth=0.15)
    ax.set_aspect("equal")
    ax.set_xlabel(f"{names[ai]} (m)")
    ax.set_ylabel(f"{names[bi]} (m)")
    ax.set_title(title)
    cbar = fig.colorbar(tpc, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label("|E| (V/m)")
    fig.text(0.02, 0.98, f"lambda0 = 0.8 m\n{axis} = {coord} m", va="top", fontsize=9)
    fig.tight_layout()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, bbox_inches="tight")
    plt.close(fig)
    print(f"saved: {out_path}")
    print(f"  |E| min={zval.min():.4f} max={zval.max():.4f} mean={zval.mean():.4f}")


def main():
    root = Path(__file__).resolve().parents[1]
    parser = argparse.ArgumentParser(description="Plot normE on a plane slice")
    parser.add_argument("--mesh", type=Path, default=root / "build" / "Release" / "SBCmesh.dat")
    parser.add_argument("--outdir", type=Path, default=root / "build" / "Release" / "OutFile")
    parser.add_argument("--plane", choices=["x", "y", "z"], default="z")
    parser.add_argument("--coord", type=float, default=0.5, help="plane coordinate (m)")
    parser.add_argument("--mode", choices=["slice", "surface"], default="slice",
                        help="slice=cut volume; surface=boundary tris on plane")
    parser.add_argument("--tol", type=float, default=1e-3, help="tolerance for surface mode")
    parser.add_argument("--output", type=Path, default=None)
    args = parser.parse_args()

    V, T, Tri, _ = load_mesh(args.mesh)
    normE = load_scalar_field(args.outdir / "normE")
    if normE.size != V.shape[0]:
        raise SystemExit(f"normE size {normE.size} != vertex count {V.shape[0]}")

    if args.output is None:
        args.output = args.outdir / f"normE_{args.plane}{args.coord:.3f}.png"

    if args.mode == "surface":
        tri, zval, plane_axes = build_surface_tris(V, Tri, normE, args.plane, args.coord, args.tol)
        title = f"normE on {args.plane}={args.coord} m (surface)"
    else:
        tri, zval, plane_axes = build_slice_tris(V, T, normE, args.plane, args.coord)
        title = f"normE on {args.plane}={args.coord} m (slice)"

    if tri is None:
        raise SystemExit(f"no data on plane {args.plane}={args.coord}")

    plot_slice(tri, zval, plane_axes, args.output, args.coord, args.plane, title)


if __name__ == "__main__":
    main()