#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Compute port power from time-modulated modal fields sampled at triangle quadrature points.

Inputs (in working dir):
  - computeDGTD.log    : contains PERFORMING INFORMATION + PORT line (see regex below)
  - port_inc_field.csv : columns:
        exc_face_id, local_face_idx, quad_idx,
        x1,y1,z1, x2,y2,z2, x3,y3,z3,   # triangle vertices (meters)
        nx,ny,nz,                       # QUADRATURE-POINT coordinates (meters)
        Et_x,Et_y,Et_z,                 # modal tangential E at quad (shape only)
        Ht_x,Ht_y,Ht_z                  # modal tangential H at quad (shape only)

Outputs:
  - port_power_spectrum.csv/.png        : per-frequency incident power (W) vs f
  - port_power_time_trace.csv           : instantaneous P(t) (W) and time-axis
  - console print of <P(t)> and sum P_k consistency check

Notes:
  - If you want TEM-consistent amplitudes, set Hmag = Emag / Zc instead of Hmag = Emag.
  - Ensure quad_idx ordering matches your 6/9-pt rule tables (it does in this script).
"""

from pathlib import Path
import re
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import fnmatch
import os
import os
os.environ["MPLBACKEND"] = "Agg"   # must be set before importing pyplot

import matplotlib
matplotlib.use("Agg", force=True)  # belt-and-suspenders
import matplotlib.pyplot as plt

import imageio.v2 as imageio
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 (enables 3D)
from matplotlib import cm



# =========================
# Config / paths
# =========================
LOG_PATH = Path("computeDGTD.log")
CSV_PATH = Path("port_inc_field.csv")
PROBE_FOLDER = Path("./PROBES")
SAVE_PREFIX = "port"

# =========================
# Constants
# =========================
Pi = np.pi
c0 = 299792458.0  # m/s

# =========================
# Parse computeDGTD.log
# =========================
def extract_selected_log_values(filepath: Path):
    targets = {
        "Final Time(sec)": None,
        "dt_sample": None,
        "tsPerSampling": None
    }
    with filepath.open("r") as f:
        lines = f.readlines()

    in_block = False
    for line in lines:
        if "==========================================" in line and "PERFORMING INFORMATION" in line:
            in_block = True
            continue
        if in_block and "==========================================" in line:
            break

        for key in targets:
            if key in line:
                m = re.search(r'=\s*([\d.eE+-]+)', line)
                if m:
                    val = float(m.group(1))
                    if key == "tsPerSampling":
                        targets[key] = int(val)
                    else:
                        targets[key] = val

    return {
        "FinalTime": targets["Final Time(sec)"],
        "dt_sample": targets["dt_sample"],
        "tsPerSampling": targets["tsPerSampling"],
    }

PORT_RE = re.compile(
    r"PORT:\s*ro=\[(?P<ro>[^]]+)\],\s*r1=\[(?P<r1>[^]]+)\],\s*Z=(?P<Z>[+\-0-9.eE]+),\s*\|E\|=(?P<Emag>[+\-0-9.eE]+),\s*Flag=(?P<Flag>\d+),\s*Tdist=(?P<Tdist>\d+),\s*f=(?P<f>[+\-0-9.eE]+),\s*t0=(?P<t0>[+\-0-9.eE]+),\s*tau=(?P<tau>[+\-0-9.eE]+)"
)

def _parse_bracket_vec(s: str):
    parts = [float(p) for p in s.split(",")]
    if len(parts) != 3:
        raise ValueError(f"Expected 3 components in vector, got: {s}")
    return np.array(parts, dtype=float)

def parse_port_from_log(log_path: Path):
    with log_path.open("r") as f:
        for line in f:
            if "PORT:" in line:
                m = PORT_RE.search(line)
                if not m:
                    continue
                ro  = _parse_bracket_vec(m.group("ro"))
                r1  = _parse_bracket_vec(m.group("r1"))
                Z   = float(m.group("Z"))
                Emag= float(m.group("Emag"))
                Flag= int(m.group("Flag"))
                Tdist=int(m.group("Tdist"))
                fMHz = float(m.group("f"))     # MHz
                t0   = float(m.group("t0"))
                tau  = float(m.group("tau"))

                khat = r1 - ro
                nrm  = np.linalg.norm(khat)
                khat = khat / nrm if nrm > 0 else np.array([0.0, 0.0, 1.0], float)

                return {
                    "ro": ro, "r1": r1, "Z": Z, "Emag": Emag,
                    "Flag": Flag, "Tdist": Tdist, "fMHz": fMHz,
                    "t0": t0, "tau": tau, "khat": khat
                }
    raise RuntimeError("No parsable PORT line found in log.")

# =========================
# Quadrature (6-pt / 9-pt) weights in same ordering as your C macros
# =========================
G6_W = np.array([0.109951743655322]*3 + [0.223381589678011]*3)   # sum = 1/2 on ref triangle
G9_W = np.array([0.205950504760887]*3 + [0.063691414286223]*6)   # sum = 1/2 on ref triangle

# =========================
# Geometry helpers
# =========================
def tri_area(v0, v1, v2):
    return 0.5 * np.linalg.norm(np.cross(v1 - v0, v2 - v0))

def tri_normal(v0, v1, v2):
    n = np.cross(v1 - v0, v2 - v0)
    nn = np.linalg.norm(n)
    if nn == 0:
        raise ValueError("Degenerate triangle")
    return n / nn

def build_port_quadrature_from_csv(df):
    """
    df columns required:
      exc_face_id, local_face_idx, quad_idx,
      x1,y1,z1, x2,y2,z2, x3,y3,z3,
      nx,ny,nz,   # QUADRATURE-POINT coordinates on the face
      Et_x,Et_y,Et_z, Ht_x,Ht_y,Ht_z
    Returns:
      rQ     (Q,3)    quad-point positions (nx,ny,nz)
      EtQ    (Q,3)    modal tangential E at quad points
      HtQ    (Q,3)    modal tangential H at quad points
      w_phys (Q,)     physical weights = (2*A_face) * W_ref[quad_idx]
      n_face (Q,3)    per-quad face unit normal (same per face),
                      flipped to point 'into device' if provided
    """
    needed = ["exc_face_id","local_face_idx","quad_idx",
              "x1","y1","z1","x2","y2","z2","x3","y3","z3",
              "nx","ny","nz","Et_x","Et_y","Et_z","Ht_x","Ht_y","Ht_z"]
    for col in needed:
        if col not in df.columns:
            raise ValueError(f"Missing column in CSV: {col}")

    rQ  = df[["nx","ny","nz"]].to_numpy(float)
    EtQ = df[["Et_x","Et_y","Et_z"]].to_numpy(float)
    HtQ = df[["Ht_x","Ht_y","Ht_z"]].to_numpy(float)

    w_phys = np.zeros(len(df), float)
    n_face = np.zeros((len(df), 3), float)

    # normalize reference direction (into device) if provided
    u_enf = np.array([0,0,1])

    # group by face
    groups = {}
    for i, row in df.iterrows():
        key = (int(row["exc_face_id"]), int(row["local_face_idx"]))
        groups.setdefault(key, []).append(i)

    for key, idxs in groups.items():
        qidxs = df.loc[idxs, "quad_idx"].to_numpy(int)
        Nq = int(qidxs.max() + 1)
        if Nq == 6:
            Wref = G6_W
        elif Nq == 9:
            Wref = G9_W
        else:
            raise ValueError(f"Face {key} has {Nq} quad points; expected 6 or 9.")

        row0 = df.loc[idxs[0]]
        v0 = np.array([row0["x1"], row0["y1"], row0["z1"]], float)
        v1 = np.array([row0["x2"], row0["y2"], row0["z2"]], float)
        v2 = np.array([row0["x3"], row0["y3"], row0["z3"]], float)

        A = tri_area(v0, v1, v2)
        n = tri_normal(v0, v1, v2)
        if u_enf is not None and np.dot(n, u_enf) < 0:
            n = -n

        w_face = (2.0 * A) * Wref[qidxs]   # physical quad weights
        w_phys[idxs] = w_face
        n_face[idxs, :] = n

        # Optional quick check: per-face sum weights ≈ area
        # import numpy as _np
        # _np.testing.assert_allclose(w_face.sum(), A, rtol=1e-10, atol=1e-12)

    return rQ, EtQ, HtQ, w_phys, n_face




Pi = np.pi
MEGA = 1e6
Vo = c0

def time_modulation_inc(flag, t, t0, khat, r, r0, freq_m, Emag, Hmag, tau):
    omega = 2.0 * Pi * freq_m * MEGA
    t = np.asarray(t, float)

    if flag == 0:
        phase = - omega * t[:, None]
        IncE = Emag * np.sin(phase)
        IncH = Hmag * np.sin(phase)
    elif flag == 1:
        Exponent = t[:, None] - t0 
        CosMod  = np.cos(omega * (t[:, None] - t0))
        env = np.exp(-(Exponent**2) / (tau*tau))
        IncE = Emag * CosMod * env
        IncH = Hmag * CosMod * env
    elif flag == 2:
        Exponent = t[:, None] - t0
        Neuman = (2.0 * Exponent) / (tau*tau)
        env = np.exp(-(Exponent**2) / (tau*tau))
        IncE = Emag * Neuman * env
        IncH = Hmag * Neuman * env
    else:
        raise ValueError("Unknown TimeDistributionFlag")
    return IncE, IncH    
    


# =========================
# Main
# =========================

# ---- parse log ----
sim = extract_selected_log_values(LOG_PATH)
port = parse_port_from_log(LOG_PATH)

dt_sample = sim["dt_sample"]
steps     = sim["tsPerSampling"]

# Sampling frequency f_s = 1 / dt
f_s = 1.0 / dt_sample
# Nyquist frequency
f_nyq = f_s / 2.0

# Count .curJ files
nfiles = len(fnmatch.filter(os.listdir(PROBE_FOLDER), '*.csv'))
nfiles = int(nfiles)
print('Number of time domain files:', nfiles)


print("\n---- Parsed sim/log ----")
print(f"FinalTime   = {sim['FinalTime']}")
print(f"dt_sample   = {dt_sample}")
print(f"steps       = {steps}")
print("\n---- Parsed PORT ----")
for k, v in port.items():
    print(f"{k:>6} : {v}")

# Number of FFT samples
NFFT = nfiles * 5

# Time axis from sampling
t_idx = np.arange(nfiles, dtype=int)       # 0..steps-1
t_sec = t_idx * dt_sample                 # seconds


# ---- load CSV ----
df = pd.read_csv(CSV_PATH)

# ---- build quad geometry/weights/normals ----
rQ, EtQ, HtQ, w_phys, n_face = build_port_quadrature_from_csv(df)

# ---- apply time modulation at quad points ----
Emag = port["Emag"]
# TEM-consistent option: Hmag = Emag / port["Z"]
Hmag = Emag
IncE, IncH = time_modulation_inc(
    port["Tdist"], t_sec, port["t0"], port["khat"],
    rQ, port["ro"], port["fMHz"], Emag, Hmag, port["tau"]
)  # (T,Q)

# build vector fields (T,Q,3)
E_inc = IncE[:, :, None] * EtQ[None, :, :]
H_inc = IncH[:, :, None] * HtQ[None, :, :]

# ---- time-domain instantaneous & average power (surface Poynting) ----
S_t   = np.cross(E_inc, H_inc)                         # (T,Q,3)
S_n_t = np.einsum('tqj,qj->tq', S_t, n_face)           # project onto normal
P_t   = np.sum(S_n_t * w_phys[None, :], axis=1)        # (T,)
P_timeavg_W = float(np.mean(P_t))
print(f"\n[Port] Time-average incident power: {P_timeavg_W:.6e} W")

# ---- frequency-domain per-bin power (single-sided) ----
freqs = np.fft.rfftfreq(NFFT, d=dt_sample)             # Hz
freqs_ghz = freqs * 1e-9

E_f = np.fft.rfft(E_inc, n=NFFT, axis=0) / NFFT        # (F,Q,3)
H_f = np.fft.rfft(H_inc, n=NFFT, axis=0) / NFFT        # (F,Q,3)
Cross_F = np.cross(E_f, np.conj(H_f))                  # (F,Q,3) note H*
S_n_F   = np.einsum('fqj,qj->fq', Cross_F, n_face)     # (F,Q)
Sum_F   = np.sum(S_n_F * w_phys[None, :], axis=1)      # (F,) complex

scale = np.ones_like(Sum_F, float)
if scale.size > 2:
    scale[1:-1] = 2.0                                  # single-sided doubling
P_per_freq_W = scale * np.real(Sum_F)                  # (F,)
P_total_bins_W = float(np.sum(P_per_freq_W))

print(f"[Port] Sum of per-frequency power:  {P_total_bins_W:.6e} W")
if P_timeavg_W != 0:
    rel = abs(P_total_bins_W - P_timeavg_W)/abs(P_timeavg_W)
    print(f"[check] rel diff sum(P_k) vs <P(t)> = {rel:.3e}")

# ---- save outputs ----
# spectrum CSV
pd.DataFrame({
    "f_Hz": freqs,
    "f_GHz": freqs_ghz,
    "P_W": P_per_freq_W
}).to_csv(f"{SAVE_PREFIX}_power_spectrum.csv", index=False)

# spectrum PNG
plt.figure(figsize=(9,5))
plt.plot(freqs_ghz, P_per_freq_W, lw=1.4)
plt.grid(True)
plt.xlabel("Frequency (GHz)")
plt.ylabel("Port power per bin (W)")
plt.title("Incident Port Power Spectrum (surface Poynting, single-sided)")
plt.tight_layout()
plt.savefig(f"{SAVE_PREFIX}_power_spectrum.png", dpi=180)
plt.close()

# time trace CSV
pd.DataFrame({"t_s": t_sec, "P_inst_W": P_t}).to_csv(f"{SAVE_PREFIX}_power_time_trace.csv", index=False)

print("\nSaved:")
print(f" - {SAVE_PREFIX}_power_spectrum.csv / .png")
print(f" - {SAVE_PREFIX}_power_time_trace.csv")




# ======= Area-weighted average port fields vs time (add after E_inc/H_inc) =======
A_total = np.sum(w_phys)                               # total port area

# Area-weighted averages of the *vector* fields at each time sample
# Shapes: (T,3)
E_avg_t = np.sum(E_inc * w_phys[None, :, None], axis=1) / A_total
H_avg_t = np.sum(H_inc * w_phys[None, :, None], axis=1) / A_total

# Magnitudes vs time
E_avg_mag = np.linalg.norm(E_avg_t, axis=1)            # (T,)
H_avg_mag = np.linalg.norm(H_avg_t, axis=1)            # (T,)

# Save to CSV
df_avg = pd.DataFrame({
    "t_s": t_sec,
    "Eavg_x": E_avg_t[:, 0], "Eavg_y": E_avg_t[:, 1], "Eavg_z": E_avg_t[:, 2],
    "Eavg_mag": E_avg_mag,
    "Havg_x": H_avg_t[:, 0], "Havg_y": H_avg_t[:, 1], "Havg_z": H_avg_t[:, 2],
    "Havg_mag": H_avg_mag,
})
df_avg.to_csv(f"{SAVE_PREFIX}_avg_field_time.csv", index=False)

# Plot magnitudes vs time
# Plot E and H magnitudes vs time with separate axes
fig, ax1 = plt.subplots(figsize=(9,5))

color_e = "tab:blue"
ax1.set_xlabel("Time (s)")
ax1.set_ylabel("|E_avg(t)|", color=color_e)
ax1.plot(t_sec, E_avg_mag, color=color_e, lw=1.4)
ax1.tick_params(axis="y", labelcolor=color_e)

# Twin y-axis for H
ax2 = ax1.twinx()
color_h = "tab:red"
ax2.set_ylabel("|H_avg(t)|", color=color_h)
ax2.plot(t_sec, H_avg_mag, color=color_h, lw=1.4, linestyle="--")
ax2.tick_params(axis="y", labelcolor=color_h)

plt.title("Area-weighted average tangential fields vs time")
fig.tight_layout()
plt.savefig(f"{SAVE_PREFIX}_avg_field_time.png", dpi=180)
plt.close(fig)


print(f" - {SAVE_PREFIX}_avg_field_time.csv / .png")






# =========================
# TOTAL FIELD processing @ port (power, averages, S11)
# =========================

import glob
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

# ---- Config for total-field files ----
FILE_BASENAME = "patch_model"   # matches "patch_model.probe" and "Probes_patch_model_XXXX.csv"
SAVE_PREFIX_TOT = "total"

# ---- Helpers: probe IO & series reader (fixing names/case) ----
def read_probe_coords(path="patch_model.probe") -> pd.DataFrame:
    dfp = pd.read_csv(path, sep=None, engine="python")
    dfp.columns = [c.strip().upper() for c in dfp.columns]
    assert set(["X","Y","Z"]).issubset(dfp.columns), "Probe file must have X,Y,Z columns"
    dfp = dfp[["X","Y","Z"]].copy()
    dfp.to_csv("probe_locations_clean.csv", index=False)
    return dfp

_TF_COLS = ["Ex","Ey","Ez","Hx","Hy","Hz"]

def _extract_tidx(fname: str) -> int:
    m = re.search(r"_([0-9]+)\.csv$", fname)
    if not m:
        raise ValueError(f"Cannot parse time index from filename: {fname}")
    return int(m.group(1))

def _read_one_probe_csv(path: str) -> pd.DataFrame:
    d = pd.read_csv(path, sep=None, engine="python", header=0)
    d.columns = [c.strip() for c in d.columns]
    keep = [c for c in _TF_COLS if c in d.columns]
    if len(keep) < 6:
        raise ValueError(f"Missing expected field columns in {path}")
    d = d[keep].copy()
    d["probe_idx"] = np.arange(len(d), dtype=int)
    d["t_idx"] = _extract_tidx(Path(path).name)
    return d

def read_total_field_series(probe_folder: Path, file_basename: str, workers: int = 12):
    pat1 = str(probe_folder / f"Probes_{file_basename}_*.csv")
    pat2 = str(probe_folder / f"Currents_{file_basename}_*.csv")
    files = sorted(glob.glob(pat1)) or sorted(glob.glob(pat2))
    if not files:
        raise FileNotFoundError(f"No probe CSVs found in {probe_folder} matching {pat1} or {pat2}")

    results = []
    with ThreadPoolExecutor(max_workers=workers) as ex:
        futures = {ex.submit(_read_one_probe_csv, f): f for f in files}
        for fut in tqdm(as_completed(futures), total=len(futures),
                        desc="Reading probe CSVs", unit="file"):
            results.append(fut.result())

    total_df = pd.concat(results, ignore_index=True)
    total_df.sort_values(["t_idx","probe_idx"], inplace=True, kind="mergesort")
    total_df.reset_index(drop=True, inplace=True)

    t_vals = np.sort(total_df["t_idx"].unique())
    p_vals = np.sort(total_df["probe_idx"].unique())
    T, P = len(t_vals), len(p_vals)
    t_map = {t:i for i,t in enumerate(t_vals)}

    E = np.empty((T, P, 3), dtype=float)
    H = np.empty((T, P, 3), dtype=float)

    # fast fill: rely on sort (t_idx first, then probe_idx)
    arr = total_df[["t_idx","probe_idx"] + _TF_COLS].to_numpy()
    # block-write is fastest, but just loop is fine:
    for row in arr:
        ti = t_map[int(row[0])]; pi = int(row[1])
        E[ti, pi, :] = row[2:5]
        H[ti, pi, :] = row[5:8]

    return total_df, E, H

# ---- Load probe coordinates + total fields ----
probes_df = read_probe_coords(f"{FILE_BASENAME}.probe")
total_df, E_tot, H_tot = read_total_field_series(PROBE_FOLDER, FILE_BASENAME, workers=12)

print("total_df shape:", total_df.shape)
print("E_tot shape:", E_tot.shape, "H_tot shape:", H_tot.shape)

# ---- Align probes to quadrature points (so weights/normals match) ----
# rQ: (Q,3) from earlier; probes_df: (P,3). We expect Q == P for a face probe export.
probes_xyz = probes_df[["X","Y","Z"]].to_numpy(float)
Q = rQ.shape[0]
P = probes_xyz.shape[0]

if P != Q:
    raise ValueError(f"Probe count ({P}) != quadrature count ({Q}). Ensure .probe uses port quad points.")

# map each quadrature point to its nearest probe index
dists = np.linalg.norm(rQ[:, None, :] - probes_xyz[None, :, :], axis=2)  # (Q,P)
order_idx = np.argmin(dists, axis=1)                                     # (Q,)
max_mm = 1e3 * float(np.max(dists[np.arange(Q), order_idx]))
if max_mm > 0.05:
    print(f"[warn] Max rQ↔probe mismatch ≈ {max_mm:.3f} mm")

# reorder total fields along probe axis to match (rQ, w_phys, n_face) order
E_tot = E_tot[:, order_idx, :]
H_tot = H_tot[:, order_idx, :]

# ---- Time axis for total fields (use dt_sample from sim) ----
T_tot = E_tot.shape[0]
t_sec_tot = np.arange(T_tot, dtype=int) * dt_sample

# ---- Instantaneous NET power from total fields: P_net(t) = ∬ (E×H)·n dA ----
S_t_tot   = np.cross(E_tot, H_tot)                             # (T,Q,3)
S_n_t_tot = np.einsum("tqj,qj->tq", S_t_tot, n_face)           # (T,Q)
P_net_t   = np.sum(S_n_t_tot * w_phys[None, :], axis=1)        # (T,)
P_net_avg = float(np.mean(P_net_t))
print(f"[Port] Time-average NET power (total fields): {P_net_avg:.6e} W")

# ---- Frequency-domain NET power (single-sided) from total fields ----
E_f_tot = np.fft.rfft(E_tot, n=NFFT, axis=0) / NFFT            # (F,Q,3) complex
H_f_tot = np.fft.rfft(H_tot, n=NFFT, axis=0) / NFFT            # (F,Q,3) complex
Cross_F_tot = np.cross(E_f_tot, np.conj(H_f_tot))              # (F,Q,3)
S_n_F_tot   = np.einsum("fqj,qj->fq", Cross_F_tot, n_face)     # (F,Q)
Sum_F_tot   = np.sum(S_n_F_tot * w_phys[None, :], axis=1)      # (F,) complex

scale = np.ones_like(Sum_F_tot.real)
if scale.size > 2:
    scale[1:-1] = 2.0
P_net_per_freq_W = scale * np.real(Sum_F_tot)                  # (F,)

# ---- Save total-field outputs ----
pd.DataFrame({"t_s": t_sec_tot, "P_net_W": P_net_t}).to_csv(f"{SAVE_PREFIX_TOT}_power_time_trace.csv", index=False)

plt.figure(figsize=(9,5))
plt.plot(freqs_ghz, P_net_per_freq_W, lw=1.4)
plt.grid(True)
plt.xlabel("Frequency (GHz)")
plt.ylabel("Net power per bin (W)")
plt.title("NET Port Power Spectrum from Total Fields (single-sided)")
plt.tight_layout()
plt.savefig(f"{SAVE_PREFIX_TOT}_power_spectrum.png", dpi=180)
plt.close()

pd.DataFrame({
    "f_Hz": freqs,
    "f_GHz": freqs_ghz,
    "P_net_W": P_net_per_freq_W
}).to_csv(f"{SAVE_PREFIX_TOT}_power_spectrum.csv", index=False)

print("\nSaved:")
print(f" - {SAVE_PREFIX_TOT}_power_spectrum.csv / .png")
print(f" - {SAVE_PREFIX_TOT}_power_time_trace.csv")

# ---- Area-weighted average TOTAL fields vs time ----
A_total = np.sum(w_phys)
E_avg_tot_t = np.sum(E_tot * w_phys[None, :, None], axis=1) / A_total   # (T,3)
H_avg_tot_t = np.sum(H_tot * w_phys[None, :, None], axis=1) / A_total   # (T,3)
E_avg_tot_mag = np.linalg.norm(E_avg_tot_t, axis=1)
H_avg_tot_mag = np.linalg.norm(H_avg_tot_t, axis=1)

df_avg_tot = pd.DataFrame({
    "t_s": t_sec_tot,
    "Eavg_x": E_avg_tot_t[:,0], "Eavg_y": E_avg_tot_t[:,1], "Eavg_z": E_avg_tot_t[:,2],
    "Eavg_mag": E_avg_tot_mag,
    "Havg_x": H_avg_tot_t[:,0], "Havg_y": H_avg_tot_t[:,1], "Havg_z": H_avg_tot_t[:,2],
    "Havg_mag": H_avg_tot_mag,
})
df_avg_tot.to_csv(f"{SAVE_PREFIX_TOT}_avg_field_time.csv", index=False)

fig, ax1 = plt.subplots(figsize=(9,5))
ax1.set_xlabel("Time (s)")
ax1.set_ylabel("|E_avg,tot(t)|")
ax1.plot(t_sec_tot, E_avg_tot_mag, lw=1.4)
ax2 = ax1.twinx()
ax2.set_ylabel("|H_avg,tot(t)|")
ax2.plot(t_sec_tot, H_avg_tot_mag, lw=1.4, linestyle="--")
plt.title("Area-weighted TOTAL tangential fields vs time")
fig.tight_layout()
plt.savefig(f"{SAVE_PREFIX_TOT}_avg_field_time.png", dpi=180)
plt.close(fig)

print(f" - {SAVE_PREFIX_TOT}_avg_field_time.csv / .png")







# =========================
# Per-quadrature comparison (no area weighting)
# =========================

# Align lengths
T_common = min(E_inc.shape[0], E_tot.shape[0])
t_common = t_sec[:T_common]
E_inc_c  = E_inc[:T_common]               # (T,Q,3)
E_tot_c  = E_tot[:T_common]               # (T,Q,3)
E_scat_c = E_tot_c - E_inc_c              # (T,Q,3)

# Magnitudes per point over time: shapes (T,Q)
Einc_mag_tq  = np.linalg.norm(E_inc_c,  axis=2)
Etot_mag_tq  = np.linalg.norm(E_tot_c,  axis=2)
Escat_mag_tq = np.linalg.norm(E_scat_c, axis=2)

# Ratios (guard against zero)
with np.errstate(divide="ignore", invalid="ignore"):
    ratio_tot_inc_tq  = Etot_mag_tq / np.maximum(Einc_mag_tq, 1e-30)
    ratio_scat_inc_tq = Escat_mag_tq / np.maximum(Einc_mag_tq, 1e-30)

# ---------- Per-quadrature stats ----------
Q = Einc_mag_tq.shape[1]
rms = lambda X: np.sqrt(np.mean(X**2, axis=0))        # (Q,)
stats_df = pd.DataFrame({
    "q_index":         np.arange(Q, dtype=int),
    "rms_|E_inc|":     rms(Einc_mag_tq),
    "rms_|E_tot|":     rms(Etot_mag_tq),
    "rms_|E_scat|":    rms(Escat_mag_tq),
    "median_ratio_tot_inc":  np.median(ratio_tot_inc_tq,  axis=0),
    "max_ratio_tot_inc":     np.max(ratio_tot_inc_tq,     axis=0),
    "median_ratio_scat_inc": np.median(ratio_scat_inc_tq, axis=0),
    "max_ratio_scat_inc":    np.max(ratio_scat_inc_tq,    axis=0),
})
stats_df.to_csv("E_per_quadrature_stats.csv", index=False)
print("Saved: E_per_quadrature_stats.csv")

# ---------- Save raw arrays (compact) ----------
np.savez_compressed(
    "E_per_quadrature_time_series.npz",
    t_s=t_common,
    Einc_mag_tq=Einc_mag_tq,
    Etot_mag_tq=Etot_mag_tq,
    Escat_mag_tq=Escat_mag_tq,
    ratio_tot_inc_tq=ratio_tot_inc_tq,
    ratio_scat_inc_tq=ratio_scat_inc_tq,
)
print("Saved: E_per_quadrature_time_series.npz (magnitudes + ratios)")

# ---------- Heatmaps (time × quadrature index) ----------
def plot_heatmap(M, title, fname, vmin=None, vmax=None):
    plt.figure(figsize=(10, 4.5))
    # imshow expects (rows, cols) => (time, q)
    im = plt.imshow(
        M, aspect="auto", origin="lower",
        extent=[0, Q-1, float(t_common[0]), float(t_common[-1])],
        vmin=vmin, vmax=vmax
    )
    plt.xlabel("Quadrature index q")
    plt.ylabel("Time (s)")
    plt.title(title)
    cbar = plt.colorbar(im)
    cbar.ax.set_ylabel("Magnitude")
    plt.tight_layout()
    plt.savefig(fname, dpi=180)
    plt.close()
    print(f"Saved: {fname}")

# Choose reasonable caps for visibility
plot_heatmap(Einc_mag_tq,  "Incident |E|(t,q)", "E_heatmap_incident.png")
plot_heatmap(Etot_mag_tq,  "Total |E|(t,q)",    "E_heatmap_total.png")
plot_heatmap(Escat_mag_tq, "Scattered |E|(t,q)","E_heatmap_scattered.png")

# Ratio heatmap (clip to avoid a few outliers dominating color scale)
ratio_cap = 4.0
plot_heatmap(
    np.clip(ratio_tot_inc_tq, 0, ratio_cap),
    f"|E_tot|/|E_inc| (clipped to ≤{ratio_cap})",
    "E_heatmap_ratio_tot_over_inc.png",
    vmin=0.0, vmax=ratio_cap
)
plot_heatmap(
    np.clip(ratio_scat_inc_tq, 0, ratio_cap),
    f"|E_scat|/|E_inc| (clipped to ≤{ratio_cap})",
    "E_heatmap_ratio_scat_over_inc.png",
    vmin=0.0, vmax=ratio_cap
)

# ---------- (Optional) export a few raw traces for sanity checks ----------
sample_qs = [0, Q//4, Q//2, 3*Q//4, Q-1] if Q >= 5 else list(range(Q))
for q in sample_qs:
    df_q = pd.DataFrame({
        "t_s": t_common,
        "Einc_mag":  Einc_mag_tq[:, q],
        "Etot_mag":  Etot_mag_tq[:, q],
        "Escat_mag": Escat_mag_tq[:, q],
        "ratio_tot_inc":  ratio_tot_inc_tq[:, q],
        "ratio_scat_inc": ratio_scat_inc_tq[:, q],
    })
    df_q.to_



# =========================
# Time-domain E-field plots: incident, total, and scattered
# =========================

# Align time lengths just in case
T_common = min(E_inc.shape[0], E_tot.shape[0])
t_common = t_sec[:T_common]  # use incident's time axis (same dt)

E_inc_c  = E_inc[:T_common]               # (T,Q,3)
E_tot_c  = E_tot[:T_common]               # (T,Q,3)
E_scat_c = E_tot_c - E_inc_c              # (T,Q,3)

# -------- Area-weighted |E|(t) magnitudes --------
A_total = float(np.sum(w_phys))

def area_weighted_vec(E_TQ3):
    # (T,Q,3) -> (T,3) area-weighted average vector
    return np.sum(E_TQ3 * w_phys[None, :, None], axis=1) / A_total

Eavg_inc_t  = area_weighted_vec(E_inc_c)      # (T,3)
Eavg_tot_t  = area_weighted_vec(E_tot_c)      # (T,3)
Eavg_scat_t = area_weighted_vec(E_scat_c)     # (T,3)

Eavg_inc_mag  = np.linalg.norm(Eavg_inc_t,  axis=1)
Eavg_tot_mag  = np.linalg.norm(Eavg_tot_t,  axis=1)
Eavg_scat_mag = np.linalg.norm(Eavg_scat_t, axis=1)

# Save CSV of area-weighted magnitudes
pd.DataFrame({
    "t_s": t_common,
    "|E_inc_avg|(t)":  Eavg_inc_mag,
    "|E_tot_avg|(t)":  Eavg_tot_mag,
    "|E_scat_avg|(t)": Eavg_scat_mag,
    "E_inc_avg_x":  Eavg_inc_t[:,0], "E_inc_avg_y":  Eavg_inc_t[:,1], "E_inc_avg_z":  Eavg_inc_t[:,2],
    "E_tot_avg_x":  Eavg_tot_t[:,0], "E_tot_avg_y":  Eavg_tot_t[:,1], "E_tot_avg_z":  Eavg_tot_t[:,2],
    "E_scat_avg_x": Eavg_scat_t[:,0], "E_scat_avg_y": Eavg_scat_t[:,1], "E_scat_avg_z": Eavg_scat_t[:,2],
}).to_csv("E_time_area_weighted.csv", index=False)

# Plot area-weighted magnitudes
plt.figure(figsize=(9,5))
plt.plot(t_common, Eavg_inc_mag,  lw=1.6, label="|E_inc,avg|(t)")
plt.plot(t_common, Eavg_tot_mag,  lw=1.6, label="|E_tot,avg|(t)")
plt.plot(t_common, Eavg_scat_mag, lw=1.6, label="|E_scat,avg|(t)")
plt.grid(True, alpha=0.4)
plt.xlabel("Time (s)")
plt.ylabel("Area-weighted |E|(t)")
plt.title("Area-weighted |E|(t): incident, total, scattered")
plt.legend(loc="best")
plt.tight_layout()
plt.savefig("E_time_area_weighted.png", dpi=180)
plt.close()

# -------- Sample quadrature point: components vs time --------
# choose the quadrature point closest to the face centroid for a representative sample
face_centroid = np.average(rQ, axis=0, weights=None)
q_sample = int(np.argmin(np.linalg.norm(rQ - face_centroid[None, :], axis=1)))

Einc_q  = E_inc_c[:, q_sample, :]     # (T,3)
Etot_q  = E_tot_c[:, q_sample, :]     # (T,3)
Escat_q = E_scat_c[:, q_sample, :]    # (T,3)

# Save CSV for the sample point
pd.DataFrame({
    "t_s": t_common,
    "Einc_x":  Einc_q[:,0],  "Einc_y":  Einc_q[:,1],  "Einc_z":  Einc_q[:,2],
    "Etot_x":  Etot_q[:,0],  "Etot_y":  Etot_q[:,1],  "Etot_z":  Etot_q[:,2],
    "Escat_x": Escat_q[:,0], "Escat_y": Escat_q[:,1], "Escat_z": Escat_q[:,2],
}).to_csv("E_time_sample_point.csv", index=False)

# Plot components at the sample point (three stacked panels)
fig, axs = plt.subplots(3, 1, figsize=(10, 7), sharex=True)
labels = ["x", "y", "z"]
for i, ax in enumerate(axs):
    ax.plot(t_common, Einc_q[:, i],  lw=1.3, label=f"E_inc_{labels[i]}")
    ax.plot(t_common, Etot_q[:, i],  lw=1.3, label=f"E_tot_{labels[i]}")
    ax.plot(t_common, Escat_q[:, i], lw=1.3, label=f"E_scat_{labels[i]}")
    ax.grid(True, alpha=0.4)
    ax.set_ylabel(f"{labels[i]}-comp")
    ax.legend(loc="best", ncol=3, fontsize=9)
axs[-1].set_xlabel("Time (s)")
fig.suptitle(f"E(t) at sample quadrature point q={q_sample}")
plt.tight_layout()
plt.savefig("E_time_sample_components.png", dpi=180)
plt.close()

print("Saved: E_time_area_weighted.csv/.png, E_time_sample_point.csv, E_time_sample_components.png")




print("\n============= Calculate S Parameters (E-only inner product) ======================\n")



# One-sided frequency axis
dt   = dt_sample
f_s  = 1.0 / dt
f_nyq = f_s / 2.0
freqs = np.fft.rfftfreq(NFFT, d=dt)      # Hz
freqs_ghz = freqs * 1e-9

# FFT of incident field (NO WINDOW)
# Shape: (F, Q, 3), complex
E_f = np.fft.rfft(E_inc, n=NFFT, axis=0) / NFFT
H_f = np.fft.rfft(H_inc, n=NFFT, axis=0) / NFFT

# Mean spectrum energy across probes & components
mean_spec_energy = np.mean(np.abs(E_f)**2, axis=(1, 2))   # (F,)

# Threshold at 10% of peak
peakE = float(mean_spec_energy.max()) if mean_spec_energy.size else 0.0
thresh = 0.10 * (peakE if peakE > 0 else 1.0)

mask = mean_spec_energy >= thresh

# Turn mask into contiguous frequency bands
def mask_to_bands(mask_bool, xvals):
    idx = np.flatnonzero(mask_bool)
    if idx.size == 0:
        return []
    bands = []
    start = prev = idx[0]
    for k in idx[1:]:
        if k == prev + 1:
            prev = k
        else:
            bands.append((start, prev))
            start = prev = k
    bands.append((start, prev))
    # return with indices and frequency endpoints
    return [(i0, i1, xvals[i0], xvals[i1]) for (i0, i1) in bands]

bands = mask_to_bands(mask, freqs)

print("=== Bands where mean |E_inc(f)|^2 ≥ 10% of peak ===")
if not bands:
    print("None found. (Check dt/NFFT/time record.)")
else:
    for (i0, i1, f0, f1) in bands:
        print(f"Idx {i0}–{i1}  |  {f0*1e-9:.6f}–{f1*1e-9:.6f} GHz  "
              f"(points: {i1 - i0})")

# If you want a single sweep range (union):
if bands:
    idx_min, fmin_hz = bands[0][0], bands[0][2]
    idx_max, fmax_hz = bands[-1][1], bands[-1][3]
    print(f"\nOverall sweep range: idx {idx_min}–{idx_max}  |  "
          f"{fmin_hz*1e-9:.6f}–{fmax_hz*1e-9:.6f} GHz  "
          f"(points: {idx_max - idx_min})")
else:
    idx_min = idx_max = 0
    fmin_hz = fmax_hz = 0.0

# Plot normalized mean spectrum energy, highlight all ≥10% bands
norm_energy = mean_spec_energy / (peakE if peakE > 0 else 1.0)

plt.figure(figsize=(10, 5))
plt.plot(freqs_ghz, norm_energy, label="Mean |E_inc(f)|² (normalized)")
for (i0, i1, f0, f1) in bands:
    plt.axvspan(f0*1e-9, f1*1e-9, alpha=0.25, label="≥10% band" if i0 == bands[0][0] else None)
plt.axhline(0.10, linestyle="--", linewidth=1)






# =========================
# S11 from E-only inner product, restricted to detected bands
# =========================
print("\n============= S11(E-only) over detected bands ======================\n")

# Reflected phasors
E_f_ref = E_f_tot - E_f     # (F, Q, 3), complex

def e_inner(A_fq3: np.ndarray, B_fq3: np.ndarray) -> np.ndarray:
    """
    Weighted inner product: ∬ A·B* dA using quadrature weights.
    A_fq3, B_fq3: (F, Q, 3) complex
    returns (F,) complex
    """
    prod_fq = np.sum(A_fq3 * np.conj(B_fq3), axis=2)      # (F, Q)
    return np.sum(prod_fq * w_phys[None, :], axis=1)      # (F,)

# Full-spectrum S11(E-only) once
num_full = 4 * e_inner(E_f_ref, E_f)                      # (F,)
den_full = e_inner(E_f, E_f)                              # (F,)
S11_full = num_full / np.maximum(den_full, 1e-30)         # (F,)
S11_mag_full = np.abs(S11_full)
S11_dB_full  = 20.0 * np.log10(np.clip(S11_mag_full, 1e-12, None))



# =========================
# Save & plot inner-product numerator/denominator for S11 (E-only)
# =========================
num_ip = num_full.copy()
den_ip = den_full.copy()

ip_df = pd.DataFrame({
    "f_Hz": freqs, "f_GHz": freqs_ghz,
    "num_real": np.real(num_ip), "num_imag": np.imag(num_ip),
    "num_abs": np.abs(num_ip), "num_phase_rad": np.angle(num_ip),
    "den_real": np.real(den_ip), "den_imag": np.imag(den_ip),
    "den_abs": np.abs(den_ip), "den_phase_rad": np.angle(den_ip),
})
ip_df.to_csv("S11_inner_products.csv", index=False)

plt.figure(figsize=(10,5))
plt.plot(freqs_ghz, np.abs(num_ip), lw=1.6, label="|⟨E_ref, E_inc⟩| (numerator)")
plt.plot(freqs_ghz, np.abs(den_ip), lw=1.6, linestyle="--", label="|⟨E_inc, E_inc⟩| (denominator)")
plt.xlabel("Frequency (GHz)"); plt.ylabel("Inner-product magnitude (SI units)")
plt.title("S11 (E-only) inner products — magnitudes")
plt.grid(True, alpha=0.4); plt.legend(); plt.tight_layout()
plt.savefig("S11_inner_products_mag.png", dpi=180); plt.close()

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10,7), sharex=True)
ax1.plot(freqs_ghz, np.real(num_ip), lw=1.4, label="Re{numerator}")
ax1.plot(freqs_ghz, np.real(den_ip), lw=1.4, linestyle="--", label="Re{denominator}")
ax1.set_ylabel("Real part"); ax1.grid(True, alpha=0.4); ax1.legend()

ax2.plot(freqs_ghz, np.imag(num_ip), lw=1.4, label="Im{numerator}")
ax2.plot(freqs_ghz, np.imag(den_ip), lw=1.4, linestyle="--", label="Im{denominator}")
ax2.set_xlabel("Frequency (GHz)"); ax2.set_ylabel("Imag part"); ax2.grid(True, alpha=0.4); ax2.legend()

fig.suptitle("S11 (E-only) inner products — real/imag parts")
plt.tight_layout(); plt.savefig("S11_inner_products_re_im.png", dpi=180); plt.close()

# Optional carrier print
f0_Hz = port["fMHz"] * 1e6
k0 = int(np.argmin(np.abs(freqs - f0_Hz)))
print(f"Carrier ~ {f0_Hz*1e-9:.6f} GHz → idx {k0}")
print(f"  num = {num_ip[k0].real:+.4e} {num_ip[k0].imag:+.4e}j  |num|={abs(num_ip[k0]):.4e}")
print(f"  den = {den_ip[k0].real:+.4e} {den_ip[k0].imag:+.4e}j  |den|={abs(den_ip[k0]):.4e}")





# =========================
# S11 from POWER spectra (single-sided bins)
#   P_inc  = P_per_freq_W
#   P_net  = P_net_per_freq_W
#   P_ref  = max(P_inc - P_net, 0)
#   |S11|  = sqrt(P_ref / P_inc)
# =========================
eps = 1e-30
P_inc = np.asarray(P_per_freq_W)          # (F,)
P_net = np.asarray(P_net_per_freq_W)      # (F,)
P_ref = np.maximum(abs(P_inc - P_net), 0.0)    # (F,)

S11_pow_mag_full = np.sqrt(np.clip(P_ref / np.maximum(P_inc, eps), 0.0, 1.0))
S11_pow_dB_full  = 20.0 * np.log10(np.clip(S11_pow_mag_full, 1e-12, None))

# Save a full-spectrum comparison CSV
pd.DataFrame({
    "f_Hz":      freqs,
    "f_GHz":     freqs_ghz,
    "S11_Eonly_real": np.real(S11_full),
    "S11_Eonly_imag": np.imag(S11_full),
    "S11_Eonly_mag":  S11_mag_full,
    "S11_Eonly_dB":   S11_dB_full,
    "S11_power_mag":  S11_pow_mag_full,
    "S11_power_dB":   S11_pow_dB_full,
    "P_inc_W":        P_inc,
    "P_net_W":        P_net,
    "P_ref_W":        P_ref,
}).to_csv("S11_compare_full.csv", index=False)


# =========================
# Plot/export by detected bands
# =========================
if not bands:
    print("No bands found — exported full-spectrum S11 comparison (S11_compare_full.csv).")

    # Full-spectrum overlay plot
    plt.figure(figsize=(10,5))
    plt.plot(freqs_ghz, S11_dB_full,     lw=1.6, label="S11 (E-only)")
    plt.plot(freqs_ghz, S11_pow_dB_full, lw=1.6, linestyle="--", label="S11 (power)")
    plt.xlabel("Frequency (GHz)"); plt.ylabel("S11 (dB)")
    plt.title("S11 comparison — full spectrum")
    plt.grid(True, alpha=0.4); plt.legend(loc="best"); plt.tight_layout()
    plt.savefig("S11_compare_full.png", dpi=180); plt.close()
else:
    # Per-band CSVs and multi-panel plot (E-only vs power)
    nB = len(bands)
    fig, axs = plt.subplots(nB, 1, figsize=(10, max(4, 2.6*nB)), sharey=True)
    if nB == 1:
        axs = [axs]

    concat_rows = []
    for bi, (i0, i1, f0_hz, f1_hz) in enumerate(bands, start=1):
        sl = slice(i0, i1+1)

        f_hz  = freqs[sl];   f_ghz = freqs_ghz[sl]
        S11E  = S11_dB_full[sl]
        S11P  = S11_pow_dB_full[sl]

        # Save CSV per band (both methods)
        df_band = pd.DataFrame({
            "f_Hz":  f_hz,
            "f_GHz": f_ghz,
            "S11_Eonly_mag": S11_mag_full[sl],
            "S11_Eonly_dB":  S11E,
            "S11_power_mag": S11_pow_mag_full[sl],
            "S11_power_dB":  S11P,
            "P_inc_W":       P_inc[sl],
            "P_net_W":       P_net[sl],
            "P_ref_W":       P_ref[sl],
        })
        df_band.to_csv(f"S11_compare_band_{bi}.csv", index=False)
        concat_rows.append(df_band)

        # Plot this band
        ax = axs[bi-1]
        ax.plot(f_ghz, S11E, lw=1.6, label="E-only")
        ax.plot(f_ghz, S11P, lw=1.6, linestyle="--", label="Power")
        ax.set_xlim(f0_hz*1e-9, f1_hz*1e-9)
        ax.set_xlabel("Frequency (GHz)")
        ax.set_ylabel("S11 (dB)")
        ax.grid(True, alpha=0.4)
        ax.legend(loc="best")
        ax.set_title(f"Band {bi}: {f0_hz*1e-9:.3f}–{f1_hz*1e-9:.3f} GHz")

    fig.suptitle("S11 comparison per detected band (E-only vs Power)")
    plt.tight_layout()
    plt.savefig("S11_compare_bands.png", dpi=180)
    plt.close(fig)

    # Concatenate bands and save
    df_concat = pd.concat(concat_rows, ignore_index=True)
    df_concat.to_csv("S11_compare_bands_concat.csv", index=False)

    # Union overlay across all bands
    idx_min, fmin_hz = bands[0][0], bands[0][2]
    idx_max, fmax_hz = bands[-1][1], bands[-1][3]
    sl_union = slice(idx_min, idx_max+1)

    plt.figure(figsize=(10,5))
    plt.plot(freqs_ghz[sl_union], S11_dB_full[sl_union],     lw=1.6, label="E-only")
    plt.plot(freqs_ghz[sl_union], S11_pow_dB_full[sl_union], lw=1.6, linestyle="--", label="Power")
    for (i0, i1, f0, f1) in bands:
        plt.axvspan(f0*1e-9, f1*1e-9, alpha=0.12)
    plt.xlabel("Frequency (GHz)")
    plt.ylabel("S11 (dB)")
    plt.title(f"S11 comparison (union): {fmin_hz*1e-9:.3f}–{fmax_hz*1e-9:.3f} GHz")
    plt.grid(True, alpha=0.4); plt.legend(loc="best"); plt.tight_layout()
    plt.savefig("S11_compare_union.png", dpi=180); plt.close()

    # Optional: report at carrier if inside any band
    f0_Hz = port["fMHz"] * 1e6
    if any((f0_Hz >= f0) and (f0_Hz <= f1) for (_, _, f0, f1) in bands):
        k0 = int(np.argmin(np.abs(freqs - f0_Hz)))
        print(f"S11(E-only) @ f0={f0_Hz*1e-9:.3f} GHz: {S11_dB_full[k0]:.2f} dB")
        print(f"S11(power) @ f0={f0_Hz*1e-9:.3f} GHz: {S11_pow_dB_full[k0]:.2f} dB")
        print(f"P_inc={P_inc[k0]:.4e} W  P_net={P_net[k0]:.4e} W  P_ref={P_ref[k0]:.4e} W")

print("Saved: inc_energy_bands.png, S11_compare_full.csv, "
      "S11_compare_bands.png (if bands), S11_compare_bands_concat.csv (if bands), "
      "S11_compare_union.png (if bands), and per-band CSVs (S11_compare_band_*.csv).")



















# =========================
# Animated 3D scatter GIFs of |E| and |H| over time (incident & total)
# =========================

# Field magnitudes at each quadrature point
E_inc_mag_q = np.linalg.norm(E_inc, axis=2)   # (T,Q)
H_inc_mag_q = np.linalg.norm(H_inc, axis=2)   # (T,Q)
E_tot_mag_q = np.linalg.norm(E_tot, axis=2)   # (T,Q)
H_tot_mag_q = np.linalg.norm(H_tot, axis=2)   # (T,Q)

# Consistent color limits across time
def finite_minmax(A):
    a = np.asarray(A).ravel()
    a = a[np.isfinite(a)]
    if a.size == 0:
        return 0.0, 1.0
    return float(np.min(a)), float(np.max(a))

lims = {
    "E_inc": finite_minmax(E_inc_mag_q),
    "H_inc": finite_minmax(H_inc_mag_q),
    "E_tot": finite_minmax(E_tot_mag_q),
    "H_tot": finite_minmax(H_tot_mag_q),
}

# 3D scatter movie helper
def make_3d_scatter_gif(vals_TQ, fname, title, vmin, vmax, step=1,
                        elev=90, azim=0, duration=0.15):
    
    print("Creating GIF for ", title)
    frames = []
    cmap = cm.get_cmap("viridis")
    # axes limits
    xmin,xmax = float(rQ[:,0].min()), float(rQ[:,0].max())
    ymin,ymax = float(rQ[:,1].min()), float(rQ[:,1].max())
    zmin,zmax = float(rQ[:,2].min()), float(rQ[:,2].max())

    for t in range(0, vals_TQ.shape[0], step):
        vals = vals_TQ[t]
        fig = plt.figure(figsize=(6,5), dpi=120)
        ax = fig.add_subplot(111, projection="3d")
        sc = ax.scatter(rQ[:,0], rQ[:,1], rQ[:,2],
                        c=vals, s=14, cmap=cmap, vmin=vmin, vmax=vmax)
        ax.set_xlim(xmin, xmax); ax.set_ylim(ymin, ymax); ax.set_zlim(zmin, zmax)
        ax.set_xlabel("x (m)"); ax.set_ylabel("y (m)"); ax.set_zlabel("z (m)")
        ax.set_title(f"{title}\n t = {t_sec[t]:.3e} s")
        cb = plt.colorbar(sc, ax=ax, fraction=0.046, pad=0.04)
        cb.set_label("|Field|")
        # Top view
        ax.view_init(elev=elev, azim=azim)
        fig.canvas.draw()
        w, h = fig.canvas.get_width_height()
        img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8).reshape(h, w, 3)
        frames.append(img)
        plt.close(fig)
    imageio.mimsave(fname, frames, duration=duration)  # slower playback
    print(f"Saved {fname}  ({len(frames)} frames, duration={duration}s per frame)")

# Keep GIFs small: cap to ~500 frames
Tgif = min(len(t_sec), E_tot_mag_q.shape[0], H_tot_mag_q.shape[0])
stride = max(1, Tgif // 500)

make_3d_scatter_gif(E_inc_mag_q[:Tgif], "gif_E_inc_mag.gif",
                    "|E_inc|(r,t)", *lims["E_inc"], step=stride,
                    elev=90, azim=0, duration=0.5)
make_3d_scatter_gif(E_tot_mag_q[:Tgif], "gif_E_tot_mag.gif",
                    "|E_tot|(r,t)", *lims["E_tot"], step=stride,
                    elev=90, azim=0, duration=0.5)
make_3d_scatter_gif(H_inc_mag_q[:Tgif], "gif_H_inc_mag.gif",
                    "|H_inc|(r,t)", *lims["H_inc"], step=stride,
                    elev=90, azim=0, duration=0.5)
make_3d_scatter_gif(H_tot_mag_q[:Tgif], "gif_H_tot_mag.gif",
                    "|H_tot|(r,t)", *lims["H_tot"], step=stride,
                    elev=90, azim=0, duration=0.5)