"""Analytic image-theory arbiter for the Yang & Rahmat-Samii (2003) Fig. 3 case.

Computes the input impedance of a horizontal 0.40-lambda0 dipole at height
h = 0.06 lambda0 over an infinite PEC or PMC plane via the induced-EMF
method with an exact sinusoidal-current near field (Balanis ch. 8.5-8.6).

  Z_total(referred to current max) = Z_self(a) -/+ Z_mutual(d=2h)
      PEC image: anti-phase  -> minus
      PMC image: in-phase    -> plus
  Z_in = Z_total / sin^2(kH)   (transfer from current max to center feed)

The self term uses the wire radius a as the field-point offset; the mutual
term uses the image separation d = 2h (parallel, side-by-side, no stagger).
This is exact for an INFINITE plane and a sinusoidal current filament, so it
is the independent ground truth for RayRF's infinite PEC/PMC boundary runs.

RayRF models the arm as a flat strip of width w; equivalent round-wire
radius a_eq = w/4, so the paper's a = 0.005*lambda0 = 0.1249 mm maps to
w = 4a = 0.4997 mm.

No scipy needed: the EMF integral is done with complex trapezoid quadrature.
"""

import csv
import os

import numpy as np

C0 = 299_792_458.0
ETA0 = 376.730313668
F0 = 12e9
LAM0 = C0 / F0                  # 24.9827 mm

L_DIP = 0.40 * LAM0             # 9.9931e-3 m, fixed physical length
A_WIRE = 0.005 * LAM0           # 1.2491e-4 m, round-wire radius
H_GND = 0.06 * LAM0             # 1.4990e-3 m, height above plane
Z0 = 50.0

NZ = 4001                       # quadrature points along the dipole


def z_emf(freq, rho):
    """Induced-EMF impedance (referred to current maximum) between two
    identical parallel side-by-side sinusoidal-current dipoles of length
    L_DIP separated laterally by rho. rho=A_WIRE -> self impedance."""
    k = 2.0 * np.pi * freq / C0
    H = L_DIP / 2.0
    z = np.linspace(-H, H, NZ)
    r1 = np.sqrt(rho**2 + (z - H) ** 2)
    r2 = np.sqrt(rho**2 + (z + H) ** 2)
    r0 = np.sqrt(rho**2 + z**2)
    kern = (np.exp(-1j * k * r1) / r1
            + np.exp(-1j * k * r2) / r2
            - 2.0 * np.cos(k * H) * np.exp(-1j * k * r0) / r0)
    cur = np.sin(k * (H - np.abs(z)))
    return 1j * ETA0 / (4.0 * np.pi) * np.trapezoid(kern * cur, z)


def zin_from_zmax(freq, z_max):
    """Transfer impedance referred to current max to the center feed."""
    k = 2.0 * np.pi * freq / C0
    s = np.sin(k * L_DIP / 2.0)
    return z_max / (s * s)


def s11_db(zin, z0=Z0):
    gam = (zin - z0) / (zin + z0)
    return 20.0 * np.log10(np.maximum(np.abs(gam), 1e-12))


def sweep(freqs):
    rows = []
    for f in freqs:
        zs = z_emf(f, A_WIRE)
        zm = z_emf(f, 2.0 * H_GND)
        zin_fs = zin_from_zmax(f, zs)
        zin_pec = zin_from_zmax(f, zs - zm)
        zin_pmc = zin_from_zmax(f, zs + zm)
        rows.append({
            "freq_hz": f,
            "zs_re": zs.real, "zs_im": zs.imag,
            "zm_re": zm.real, "zm_im": zm.imag,
            "zin_fs_re": zin_fs.real, "zin_fs_im": zin_fs.imag,
            "zin_pec_re": zin_pec.real, "zin_pec_im": zin_pec.imag,
            "zin_pmc_re": zin_pmc.real, "zin_pmc_im": zin_pmc.imag,
            "s11_fs_db": s11_db(zin_fs),
            "s11_pec_db": s11_db(zin_pec),
            "s11_pmc_db": s11_db(zin_pmc),
        })
    return rows


def main():
    out_dir = os.path.dirname(os.path.abspath(__file__))
    freqs = np.linspace(10e9, 20e9, 201)
    rows = sweep(freqs)

    csv_path = os.path.join(out_dir, "image_theory.csv")
    with open(csv_path, "w", newline="") as fh:
        wr = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
        wr.writeheader()
        wr.writerows(rows)

    # Spot values at 12 GHz for the report
    i12 = int(np.argmin(np.abs(freqs - 12e9)))
    r = rows[i12]
    print(f"lambda0 = {LAM0*1e3:.4f} mm  L = {L_DIP*1e3:.4f} mm  "
          f"a = {A_WIRE*1e3:.4f} mm (strip w = {4*A_WIRE*1e3:.4f} mm)  "
          f"h = {H_GND*1e3:.4f} mm")
    print(f"12 GHz: Zself = {r['zs_re']:.1f}{r['zs_im']:+.1f}j  "
          f"Zmut = {r['zm_re']:.1f}{r['zm_im']:+.1f}j")
    print(f"12 GHz: Zin free-space = {r['zin_fs_re']:.1f}{r['zin_fs_im']:+.1f}j"
          f"  S11 = {r['s11_fs_db']:.2f} dB")
    print(f"12 GHz: Zin PEC  = {r['zin_pec_re']:.1f}{r['zin_pec_im']:+.1f}j"
          f"  S11 = {r['s11_pec_db']:.2f} dB")
    print(f"12 GHz: Zin PMC  = {r['zin_pmc_re']:.1f}{r['zin_pmc_im']:+.1f}j"
          f"  S11 = {r['s11_pmc_db']:.2f} dB")

    fghz = freqs / 1e9
    s_fs = [x["s11_fs_db"] for x in rows]
    s_pec = [x["s11_pec_db"] for x in rows]
    s_pmc = [x["s11_pmc_db"] for x in rows]
    pmc_min = int(np.argmin(s_pmc))
    fs_min = int(np.argmin(s_fs))
    print(f"PMC best match: {s_pmc[pmc_min]:.2f} dB at {fghz[pmc_min]:.2f} GHz")
    print(f"Free-space resonance dip: {s_fs[fs_min]:.2f} dB at "
          f"{fghz[fs_min]:.2f} GHz")
    print(f"PEC band range: [{min(s_pec):.2f}, {max(s_pec):.2f}] dB")

    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4.5))
    ax1.plot(fghz, s_pec, "b-", label="PEC (image anti-phase)")
    ax1.plot(fghz, s_pmc, "r-", label="PMC (image in-phase)")
    ax1.plot(fghz, s_fs, "k--", lw=1, label="free space (no image)")
    ax1.set_xlabel("Frequency (GHz)")
    ax1.set_ylabel("S11 (dB)")
    ax1.set_title("Image-theory return loss, infinite plane, h = 0.06$\\lambda_0$")
    ax1.set_ylim(-30, 0)
    ax1.grid(True, alpha=0.3)
    ax1.legend()

    ax2.plot(fghz, [x["zin_pec_re"] for x in rows], "b-", label="R PEC")
    ax2.plot(fghz, [x["zin_pec_im"] for x in rows], "b--", label="X PEC")
    ax2.plot(fghz, [x["zin_pmc_re"] for x in rows], "r-", label="R PMC")
    ax2.plot(fghz, [x["zin_pmc_im"] for x in rows], "r--", label="X PMC")
    ax2.axhline(50, color="gray", lw=0.8)
    ax2.axhline(0, color="gray", lw=0.8)
    ax2.set_xlabel("Frequency (GHz)")
    ax2.set_ylabel("Zin (ohm)")
    ax2.set_title("Image-theory input impedance")
    ax2.grid(True, alpha=0.3)
    ax2.legend(ncol=2)

    fig.tight_layout()
    png_path = os.path.join(out_dir, "image_theory.png")
    fig.savefig(png_path, dpi=140)
    print(f"wrote {csv_path}")
    print(f"wrote {png_path}")


if __name__ == "__main__":
    main()
