"""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("//"):
            continue
        vals.append(float(line.split()[0]))
    return np.array(vals)


def load_complex_field(path: Path) -> np.ndarray:
    """Load vertex complex field (Ex/Ey/Ez). Accepts 'a b' or '(a,b)'."""
    vals = []
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("//"):
            continue
        s = line.replace("(", "").replace(")", "").replace(",", " ")
        parts = s.split()
        if len(parts) >= 2:
            vals.append(complex(float(parts[0]), float(parts[1])))
        else:
            vals.append(complex(float(parts[0]), 0.0))
    return np.array(vals, dtype=np.complex128)


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 _barycentric_weights(p: np.ndarray, tet_idx: np.ndarray, V: np.ndarray) -> np.ndarray:
    v = V[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)
    return np.array([w0, w123[0], w123[1], w123[2]], dtype=float)


def interp_scalar(p: np.ndarray, tet_idx: np.ndarray, V: np.ndarray, field: np.ndarray) -> float:
    w = _barycentric_weights(p, tet_idx, V)
    return float(np.dot(w, field[tet_idx]))


def interp_complex(p: np.ndarray, tet_idx: np.ndarray, V: np.ndarray, field: np.ndarray) -> complex:
    w = _barycentric_weights(p, tet_idx, V)
    return complex(np.dot(w, field[tet_idx]))


def build_slice_tris(V: np.ndarray, T: np.ndarray, field: np.ndarray, axis: str, coord: float):
    """Slice by interpolating a nodal scalar (e.g. normE). Prefer build_slice_tris_from_E."""
    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_slice_tris_from_E(
    V: np.ndarray,
    T: np.ndarray,
    Ex: np.ndarray,
    Ey: np.ndarray,
    Ez: np.ndarray,
    axis: str,
    coord: float,
):
    """Slice |E| by interpolating Ex,Ey,Ez then taking modulus.

    Interpolating nodal |E| through long tets is wrong (|E| is not linear) and
    invents fake midplane peaks (seen on BELE z=0). Always prefer this path.
    """
    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 = []
        for p in poly_arr:
            e = np.array(
                [
                    interp_complex(p, tet, V, Ex),
                    interp_complex(p, tet, V, Ey),
                    interp_complex(p, tet, V, Ez),
                ],
                dtype=np.complex128,
            )
            local_vals.append(float(np.sqrt(np.sum(np.abs(e) ** 2))))
        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, flat: bool = False):
    ai, bi = plane_axes
    names = ["x", "y", "z"]
    fig, ax = plt.subplots(figsize=(7.5, 6.5), dpi=150)
    shading = "flat" if flat else "gouraud"
    lw = 0.35 if flat else 0.15
    tpc = ax.tripcolor(tri, zval, shading=shading, cmap="jet", edgecolors="k", linewidth=lw)
    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"{axis} = {coord:g} 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("--norme", type=Path, default=None, help="normE file (default: outdir/normE)")
    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("--flat", action="store_true", help="flat shading (show mesh patches)")
    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_path = args.norme if args.norme is not None else args.outdir / "normE"

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

    if args.mode == "surface":
        normE = load_scalar_field(norme_path)
        if normE.size != V.shape[0]:
            raise SystemExit(f"normE size {normE.size} != vertex count {V.shape[0]}")
        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:
        # Prefer Ex/Ey/Ez → |E| (correct). Fallback to nodal normE.
        ex_p, ey_p, ez_p = args.outdir / "Ex", args.outdir / "Ey", args.outdir / "Ez"
        if ex_p.is_file() and ey_p.is_file() and ez_p.is_file():
            Ex = load_complex_field(ex_p)
            Ey = load_complex_field(ey_p)
            Ez = load_complex_field(ez_p)
            if Ex.size != V.shape[0]:
                raise SystemExit(f"Ex size {Ex.size} != vertex count {V.shape[0]}")
            tri, zval, plane_axes = build_slice_tris_from_E(V, T, Ex, Ey, Ez, args.plane, args.coord)
            title = f"|E| on {args.plane}={args.coord} m (from ExEyEz)"
        else:
            normE = load_scalar_field(norme_path)
            if normE.size != V.shape[0]:
                raise SystemExit(f"normE size {normE.size} != vertex count {V.shape[0]}")
            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, flat=args.flat)


if __name__ == "__main__":
    main()
