import numpy as np
import matplotlib.pyplot as plt
import pandas as pd


def te10_impedance(f_hz, a_m, b_m, er=1.0, mur=1.0):
    """
    TE10 wave impedance of a rectangular waveguide.
    """
    # constants
    c0  = 299_792_458.0
    mu0 = 4e-7 * np.pi
    eps0 = 1.0 / (mu0 * c0**2)

    f = np.atleast_1d(f_hz).astype(float)
    omega = 2*np.pi*f

    mu  = mu0 * mur
    eps = eps0 * er
    eta = np.sqrt(mu/eps)
    k   = omega * np.sqrt(mu*eps)           # free-space wavenumber in the medium

    # TE10 cutoff
    fc = c0 / (2 * a_m * np.sqrt(er * mur))

    kc = np.pi / a_m
    kz = np.sqrt(k**2 - kc**2 + 0j)         # +0j handles evanescent region

    Z_te = omega*mu / kz                    # TE impedance

    extras = {'fc': fc, 'eta': eta, 'kz': kz}
    return Z_te, extras


if __name__ == "__main__":
    # Waveguide dims
    a = 200e-3   # m
    b = 10e-3    # m
    freqs = np.linspace(1e9, 5e9, 501)  # Hz

    Zte, info = te10_impedance(freqs, a, b, er=1.0, mur=1.0)

    print(f"TE10 cutoff fc = {info['fc']/1e9:.3f} GHz")

    # plotting
    fig, ax = plt.subplots(2, 1, figsize=(7,6), sharex=True)

    ax[0].plot(freqs/1e9, np.abs(Zte))
    ax[0].axvline(info['fc']/1e9, color='r', ls='--', label="Cutoff")
    ax[0].set_ylabel("|Z_TE| (Ohms)")
    ax[0].set_title("Rectangular WG TE10 Mode Impedance")
    ax[0].grid(True)
    ax[0].legend()

    ax[1].plot(freqs/1e9, np.angle(Zte, deg=True))
    ax[1].axvline(info['fc']/1e9, color='r', ls='--')
    ax[1].set_ylabel("∠Z_TE (deg)")
    ax[1].set_xlabel("Frequency (GHz)")
    ax[1].grid(True)

    plt.tight_layout()
    plt.savefig("IMPEDANCE.png")
    
    

    # Build DataFrame with impedance data
    df = pd.DataFrame({
        "Frequency (Hz)": freqs,
        "Frequency (GHz)": freqs/1e9,
        "|Z_TE| (Ohm)": np.abs(Zte),
        "Phase(Z_TE) (deg)": np.angle(Zte, deg=True),
        "Real(Z_TE)": np.real(Zte),
        "Imag(Z_TE)": np.imag(Zte)
    })

    # Create cutoff frequency metadata row
    cutoff_info = pd.DataFrame({
        "Frequency (Hz)": [""],
        "Frequency (GHz)": [""],
        "|Z_TE| (Ohm)": [""],
        "Phase(Z_TE) (deg)": [""],
        "Real(Z_TE)": [""],
        "Imag(Z_TE)": [""]
    })

    # Insert cutoff frequency note at the top
    cutoff_note = pd.DataFrame({"Frequency (Hz)": [f"Cutoff fc = {info['fc']:.3e} Hz ({info['fc']/1e9:.3f} GHz)"]})
    
    # Save to Excel with two sheets
    out_path = "TE10_Impedance.xlsx"
    with pd.ExcelWriter(out_path, engine="openpyxl") as writer:
        cutoff_note.to_excel(writer, sheet_name="Summary", index=False)
        df.to_excel(writer, sheet_name="Impedance_Data", index=False)

    print(f"Excel file saved to {out_path}")