#!/usr/bin/env python3
import re
import argparse
import numpy as np
import matplotlib.pyplot as plt

LINE_RE = re.compile(
    r"""r=\(\s*([+-]?\d*\.?\d+e[+-]?\d+)\s*,\s*([+-]?\d*\.?\d+e[+-]?\d+)\s*,\s*([+-]?\d*\.?\d+e[+-]?\d+)\s*\)\s*\|\s*
        Et_pre=\(\s*([+-]?\d*\.?\d+e[+-]?\d+)\s*,\s*([+-]?\d*\.?\d+e[+-]?\d+)\s*,\s*([+-]?\d*\.?\d+e[+-]?\d+)\s*\)""",
    re.IGNORECASE | re.VERBOSE
)

def parse_file(path):
    xs, ys, zs, ex, ey, ez = [], [], [], [], [], []
    with open(path, "r") as f:
        for line in f:
            m = LINE_RE.search(line)
            if not m:
                continue
            rx, ry, rz, etx, ety, etz = map(float, m.groups())
            xs.append(rx); ys.append(ry); zs.append(rz)
            ex.append(etx); ey.append(ety); ez.append(etz)
    return np.array(xs), np.array(ys), np.array(zs), np.array(ex), np.array(ey), np.array(ez)

def main():
    ap = argparse.ArgumentParser(description="Plot Et quiver from CUDA log lines.")
    ap.add_argument("--txt", required=True, help="Input text file with lines containing r=(...) | Et_pre=(...) ...")
    ap.add_argument("--out", default="et_quiver.png", help="Output PNG path")
    ap.add_argument("--normalize", action="store_true", help="Normalize arrows to unit length (color still |Et|)")
    ap.add_argument("--scale", type=float, default=1.0, help="Multiply arrow length by this factor")
    ap.add_argument("--width", type=float, default=0.008, help="Arrow thickness")
    ap.add_argument("--headlength", type=float, default=6, help="Arrow head length (points)")
    ap.add_argument("--headwidth", type=float, default=4, help="Arrow head width (points)")
    ap.add_argument("--cmap", default="viridis", help="Matplotlib colormap")
    args = ap.parse_args()

    x, y, z, Et_x, Et_y, Et_z = parse_file(args.txt)
    if x.size == 0:
        raise SystemExit("No valid lines parsed. Check the file format.")

    mag = np.hypot(Et_x, Et_y)
    u, v = Et_x.copy(), Et_y.copy()
    if args.normalize:
        nz = mag > 0
        u[nz] /= mag[nz]
        v[nz] /= mag[nz]
    u *= args.scale
    v *= args.scale

    fig, ax = plt.subplots(figsize=(7, 7))
    q = ax.quiver(
        x, y, u, v, mag, cmap=args.cmap,
        angles="xy", scale_units="xy", scale=1.0,
        width=args.width, headlength=args.headlength, headwidth=args.headwidth
    )

    ax.set_aspect("equal", adjustable="box")
    ax.set_xlabel("r_x")
    ax.set_ylabel("r_y")
    ax.grid(True, ls="--", alpha=0.4)
    ax.set_title("|Et| quiver at quadrature points")
    # No colorbar by default; uncomment to show:
    # fig.colorbar(q, ax=ax, label="|Et|")

    fig.savefig(args.out, dpi=220, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved {args.out}   (parsed {x.size} points)")

if __name__ == "__main__":
    main()