"""
plotting.py
-----------

Plotting functions for notebooks or scripts after running fits.
"""
from __future__ import annotations
import numpy as np
import matplotlib.pyplot as plt
import skrf as rf


_PLOT_STYLE = {
    "figure.dpi":       160,
    "savefig.dpi":      300,
    "font.size":        8,
    "axes.labelsize":   8,
    "axes.titlesize":   8,
    "legend.fontsize":  7,
    "xtick.labelsize":  7,
    "ytick.labelsize":  7,
    "lines.linewidth":  1.4,
    "axes.linewidth":   0.8,
    "xtick.major.width": 0.8,
    "ytick.major.width": 0.8,
    "xtick.minor.width": 0.4,
    "ytick.minor.width": 0.4,
}


def plot_re_im(
    network: rf.Network,
    m: int = 0,
    n: int = 0,
    dpi: int = 300,
    title: str = None,
) -> None:
    """
    Plot Re(S[m,n]) and Im(S[m,n]) side by side.

    Parameters
    ----------
    network : rf.Network
    m, n : int
        S-parameter indices (0-based).
    dpi : int
        Save DPI (does not affect the displayed figure).
    title : str, optional
        Super-title for the figure.
    """
    plt.rcParams.update({**_PLOT_STYLE, "savefig.dpi": dpi})
    f = network.frequency.f / 1e9
    s = network.s[:, m, n]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.3))
    ax1.plot(f, np.real(s), color="red", label="Re(S{}{})".format(m + 1, n + 1))
    ax1.set_xlabel("Frequency (GHz)")
    ax1.set_ylabel("Real")
    ax1.legend()

    ax2.plot(f, np.imag(s), color="red", label="Im(S{}{})".format(m + 1, n + 1))
    ax2.set_xlabel("Frequency (GHz)")
    ax2.set_ylabel("Imaginary")
    ax2.legend()

    if title:
        fig.suptitle(title)
    plt.tight_layout()
    plt.show()


def plot_lom_vs_data_re_im(
    lom_network: rf.Network,
    data_network: rf.Network,
    m: int = 0,
    n: int = 0,
    lom_label: str = "LOM",
    data_label: str = "Data",
) -> None:
    """
    Overlay Re and Im of a fitted LOM network against a reference (data/CPW).

    Parameters
    ----------
    lom_network : rf.Network
        The fitted lumped-element model.
    data_network : rf.Network
        The reference (measured or CPW simulation).
    m, n : int
        S-parameter indices (0-based).
    lom_label, data_label : str
        Legend labels.
    """

    plt.rcParams.update(_PLOT_STYLE)
    f_lom  = lom_network.frequency.f / 1e9
    f_data = data_network.frequency.f / 1e9
    s_lom  = lom_network.s[:, m, n]
    s_data = data_network.s[:, m, n]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.3))

    ax1.plot(f_lom,  np.real(s_lom),  color="red",  label="{} Re(S{}{})".format(lom_label,  m+1, n+1))
    ax1.plot(f_data, np.real(s_data), color="blue", label="{} Re(S{}{})".format(data_label, m+1, n+1), linestyle="--")
    ax1.set_xlabel("Frequency (GHz)")
    ax1.set_ylabel("Real")
    ax1.legend()

    ax2.plot(f_lom,  np.imag(s_lom),  color="red",  label="{} Im(S{}{})".format(lom_label,  m+1, n+1))
    ax2.plot(f_data, np.imag(s_data), color="blue", label="{} Im(S{}{})".format(data_label, m+1, n+1), linestyle="--")
    ax2.set_xlabel("Frequency (GHz)")
    ax2.set_ylabel("Imaginary")
    ax2.legend()

    plt.tight_layout()
    plt.show()


def plot_all_models(
    networks: dict[str, rf.Network],
    m: int = 0,
    n: int = 0,
    quantity: str = "re",
    title: str = None,
) -> None:
    """
    Plot one S-parameter quantity for multiple networks on a single axes.

    Useful for comparing CPW, FosterFit, OptimizedFit, and AnalyticalFit
    in one call.

    Parameters
    ----------
    networks : dict[str, rf.Network]
        Mapping of label → network, e.g.
        {"CPW": cpw_net, "Foster": foster_net, "Optimized": opt_net}
    m, n : int
        S-parameter indices (0-based).
    quantity : {"re", "im", "db", "abs"}
        Which quantity to plot.
    title : str, optional
        Plot title.
    """
    plt.rcParams.update(_PLOT_STYLE)

    _extractors = {
        "re":  lambda s: np.real(s),
        "im":  lambda s: np.imag(s),
        "db":  lambda s: 20 * np.log10(np.maximum(np.abs(s), 1e-300)),
        "abs": lambda s: np.abs(s),
    }
    _ylabels = {"re": "Re", "im": "Im", "db": "|S| (dB)", "abs": "|S|"}

    if quantity not in _extractors:
        raise ValueError("quantity must be one of: {}".format(list(_extractors)))

    extract = _extractors[quantity]

    fig, ax = plt.subplots(figsize=(7, 4))
    for label, net in networks.items():
        f = net.frequency.f / 1e9
        s = net.s[:, m, n]
        ax.plot(f, extract(s), label=label)

    ax.set_xlabel("Frequency (GHz)")
    ax.set_ylabel("{} (S{}{})".format(_ylabels[quantity], m + 1, n + 1))
    ax.legend()
    if title:
        ax.set_title(title)
    plt.tight_layout()
    plt.show()



import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.ticker as mticker
from mpl_toolkits.axes_grid1 import make_axes_locatable
import skrf as rf
from scipy.signal import find_peaks

# ── Shared style ───────────────────────────────────────────────────────────────

_PAPER_RC = {
    "figure.dpi":        150,
    "savefig.dpi":       300,
    "font.family":       "sans-serif",
    "font.size":         9,
    "axes.labelsize":    9,
    "axes.titlesize":    9,
    "legend.fontsize":   8,
    "xtick.labelsize":   8,
    "ytick.labelsize":   8,
    "lines.linewidth":   1.4,
    "axes.linewidth":    0.8,
    "xtick.major.width": 0.8,
    "ytick.major.width": 0.8,
    "xtick.minor.width": 0.5,
    "ytick.minor.width": 0.5,
    "xtick.direction":   "in",
    "ytick.direction":   "in",
    "axes.spines.top":   True,
    "axes.spines.right": True,
}

_MODEL_COLORS = {
    "cpw":        "#2c7bb6",
    "optimized":  "#d7191c",
    "foster":     "#1a9641",
    "analytical": "#ff7f00",
}

_MODEL_LABELS = {
    "cpw":        "CPW",
    "optimized":  "Optimized",
    "foster":     "Foster",
    "analytical": "Analytical",
}


def _apply_paper_style():
    plt.rcParams.update(_PAPER_RC)


def _save_or_show(fig, save_path):
    plt.tight_layout()
    if save_path:
        fig.savefig(save_path, dpi=300, bbox_inches="tight")
        print(f"Saved: {save_path}")
    plt.show()


# ──────────────────────────────────────────────────────────────────────────────
# 1. plot_fit_residuals
# ──────────────────────────────────────────────────────────────────────────────

def plot_fit_residuals(
    data_ntw,
    lom_networks: dict,
    m: int = 0,
    n: int = 0,
    title: str = None,
    save_path: str = None,
):
    """
    Plot the complex residual S_lom - S_data for one or more LOM networks.

    Residuals are shown as Re and Im separately.  A perfect fit is a flat
    zero line.  Systematic errors (frequency offset, linewidth mismatch)
    appear as structured curves — much more diagnostic than an overlay.

    Parameters
    ----------
    data_ntw : rf.Network
        Reference (CPW or measured) network.
    lom_networks : dict[str, rf.Network]
        Mapping of label → fitted LOM network.
        Labels should be "optimized", "foster", or "analytical" to get
        automatic colours, or any string for a custom label.
    m, n : int
        S-parameter indices (0-based).
    title : str, optional
    save_path : str, optional
        If provided, saves figure to this path at 300 dpi.
    """
    _apply_paper_style()

    fig, axes = plt.subplots(2, 1, figsize=(6, 4.5), sharex=True)

    f_data = data_ntw.frequency.f / 1e9
    s_data = data_ntw.s[:, m, n]

    for label, lom_ntw in lom_networks.items():
        # Interpolate lom onto data frequency grid if needed
        f_lom = lom_ntw.frequency.f / 1e9
        s_lom = lom_ntw.s[:, m, n]
        if not np.allclose(f_lom, f_data):
            re_interp = np.interp(f_data, f_lom, np.real(s_lom))
            im_interp = np.interp(f_data, f_lom, np.imag(s_lom))
            s_lom = re_interp + 1j * im_interp

        residual = s_lom - s_data
        color    = _MODEL_COLORS.get(label, None)
        disp_label = _MODEL_LABELS.get(label, label)

        axes[0].plot(f_data, np.real(residual), color=color, label=disp_label)
        axes[1].plot(f_data, np.imag(residual), color=color, label=disp_label)

    for ax in axes:
        ax.axhline(0, color="black", linewidth=0.7, linestyle="--", alpha=0.5)
        ax.yaxis.set_minor_locator(mticker.AutoMinorLocator())
        ax.xaxis.set_minor_locator(mticker.AutoMinorLocator())

    axes[0].set_ylabel(r"$\mathrm{Re}(S_\mathrm{LOM} - S_\mathrm{ref})$")
    axes[1].set_ylabel(r"$\mathrm{Im}(S_\mathrm{LOM} - S_\mathrm{ref})$")
    axes[1].set_xlabel("Frequency (GHz)")
    axes[0].legend(framealpha=0.9, edgecolor="0.8")

    port_str = f"S{m+1}{n+1}"
    fig_title = title or f"Fit residuals — {port_str}"
    axes[0].set_title(fig_title)

    _save_or_show(fig, save_path)


# ──────────────────────────────────────────────────────────────────────────────
# 2. plot_fit_summary
# ──────────────────────────────────────────────────────────────────────────────

def plot_fit_summary(
    results: dict,
    port: str = "s11",
    title: str = None,
    save_path: str = None,
):
    """
    Three-panel summary of fitting quality for all three LC models vs CPW.

    Panels:
      Left   — Overlay of Re(S) for CPW + all three models
      Centre — f₀ and κ percent errors as a grouped bar chart
      Right  — Complex residual magnitude |S_lom - S_cpw| across the window

    Parameters
    ----------
    results : dict
        Output dict from analyze_system().
    port : {"s11", "s22"}
        Which reflection port to summarise.
    title : str, optional
    save_path : str, optional
    """
    _apply_paper_style()

    port = port.lower()
    m, n = (0, 0) if port == "s11" else (1, 1)

    cpw_net  = results["cpw_network"]
    networks = {
        "optimized":  results["optimized_network"],
        "foster":     results["foster_network"],
        "analytical": results["analytical_network"],
    }

    f = cpw_net.frequency.f / 1e9
    s_cpw = cpw_net.s[:, m, n]

    fig, axes = plt.subplots(1, 3, figsize=(12, 3.8))

    # ── Panel 1: Re(S) overlay ────────────────────────────────────────────────
    ax = axes[0]
    ax.plot(f, np.real(s_cpw), color=_MODEL_COLORS["cpw"],
            lw=2.0, label="CPW", zorder=5)
    for label, net in networks.items():
        f_lom = net.frequency.f / 1e9
        s_lom = net.s[:, m, n]
        ax.plot(f_lom, np.real(s_lom),
                color=_MODEL_COLORS[label], lw=1.2, linestyle="--",
                label=_MODEL_LABELS[label])
    ax.set_xlabel("Frequency (GHz)")
    ax.set_ylabel(r"$\mathrm{Re}(" + f"S_{{{m+1}{n+1}}}" + r")$")
    ax.set_title(f"Re(S{m+1}{n+1}) overlay")
    ax.legend(framealpha=0.9, edgecolor="0.8")
    ax.xaxis.set_minor_locator(mticker.AutoMinorLocator())
    ax.yaxis.set_minor_locator(mticker.AutoMinorLocator())

    # ── Panel 2: f₀ and κ error bar chart ────────────────────────────────────
    ax = axes[1]
    PORT = port.upper()
    model_keys = ["optimized", "foster", "analytical"]
    f0_errors = [
        results.get(f"{_MODEL_LABELS[k]} f0 error {PORT} (%)", np.nan)
        for k in model_keys
    ]
    k_errors = [
        results.get(f"{_MODEL_LABELS[k]} kappa error {PORT} (%)", np.nan)
        for k in model_keys
    ]

    x      = np.arange(len(model_keys))
    width  = 0.32
    bars_f = ax.bar(x - width/2, f0_errors, width,
                    color=[_MODEL_COLORS[k] for k in model_keys],
                    alpha=0.85, label="f₀ error")
    bars_k = ax.bar(x + width/2, k_errors, width,
                    color=[_MODEL_COLORS[k] for k in model_keys],
                    alpha=0.45, edgecolor=[_MODEL_COLORS[k] for k in model_keys],
                    linewidth=1.2, label="κ error")

    ax.axhline(0, color="black", linewidth=0.7)
    ax.set_xticks(x)
    ax.set_xticklabels([_MODEL_LABELS[k] for k in model_keys])
    ax.set_ylabel("Percent error (%)")
    ax.set_title(f"f₀ and κ errors ({PORT})")

    # Solid = f0, hatched = kappa legend
    from matplotlib.patches import Patch
    ax.legend(handles=[
        Patch(facecolor="grey", alpha=0.85, label="f₀ error"),
        Patch(facecolor="grey", alpha=0.45, edgecolor="grey", label="κ error"),
    ], framealpha=0.9, edgecolor="0.8")
    ax.yaxis.set_minor_locator(mticker.AutoMinorLocator())

    # ── Panel 3: residual magnitude ───────────────────────────────────────────
    ax = axes[2]
    for label, net in networks.items():
        f_lom = net.frequency.f / 1e9
        s_lom = net.s[:, m, n]
        if not np.allclose(f_lom, f):
            re_i = np.interp(f, f_lom, np.real(s_lom))
            im_i = np.interp(f, f_lom, np.imag(s_lom))
            s_lom = re_i + 1j * im_i
        ax.plot(f, np.abs(s_lom - s_cpw),
                color=_MODEL_COLORS[label], label=_MODEL_LABELS[label])

    ax.set_xlabel("Frequency (GHz)")
    ax.set_ylabel(r"$|S_\mathrm{LOM} - S_\mathrm{CPW}|$")
    ax.set_title(f"|Residual| magnitude ({PORT})")
    ax.legend(framealpha=0.9, edgecolor="0.8")
    ax.xaxis.set_minor_locator(mticker.AutoMinorLocator())
    ax.yaxis.set_minor_locator(mticker.AutoMinorLocator())

    fig.suptitle(title or "Fit summary", fontsize=10, y=1.01)
    _save_or_show(fig, save_path)


# ──────────────────────────────────────────────────────────────────────────────
# 3. plot_transmission_spectrum
# ──────────────────────────────────────────────────────────────────────────────

def plot_transmission_spectrum(
    networks: dict,
    m: int = 0,
    n: int = 1,
    annotate_resonances: bool = True,
    prominence: float = 3.0,
    title: str = None,
    save_path: str = None,
):
    """
    Plot |S_mn| in dB for one or more networks, resembling a VNA screenshot.

    Resonance dips/peaks are found automatically and annotated with their
    frequency and depth.  Designed to look like a measurement so that
    simulation and data can be compared at a glance.

    Parameters
    ----------
    networks : dict[str, rf.Network]
        Mapping of label → network.  Use "cpw", "optimized", etc. for
        automatic colours, or any string.
    m, n : int
        S-parameter indices (0-based).  Default S21 (0,1).
    annotate_resonances : bool
        If True, marks each resonance with a vertical line and frequency label.
    prominence : float
        Minimum prominence in dB for peak/dip detection.
    title : str, optional
    save_path : str, optional
    """
    _apply_paper_style()

    fig, ax = plt.subplots(figsize=(8, 4))

    all_annotations = []

    for label, net in networks.items():
        f_ghz  = net.frequency.f / 1e9
        s_mn   = net.s[:, m, n]
        mag_db = 20 * np.log10(np.maximum(np.abs(s_mn), 1e-300))
        color  = _MODEL_COLORS.get(label, None)
        disp   = _MODEL_LABELS.get(label, label)

        ax.plot(f_ghz, mag_db, color=color, label=disp, zorder=3)

        if annotate_resonances:
            # Detect dips (invert for find_peaks)
            dip_idx, props = find_peaks(-mag_db, prominence=prominence)
            for idx in dip_idx:
                all_annotations.append((f_ghz[idx], mag_db[idx], color))

    # Draw annotations (de-duplicated by proximity)
    if annotate_resonances and all_annotations:
        # Sort by frequency and deduplicate within 0.05 GHz
        all_annotations.sort(key=lambda x: x[0])
        kept = [all_annotations[0]]
        for ann in all_annotations[1:]:
            if ann[0] - kept[-1][0] > 0.05:
                kept.append(ann)

        for f_ann, db_ann, color in kept:
            ax.axvline(f_ann, color=color, linewidth=0.7,
                       linestyle=":", alpha=0.7, zorder=2)
            ax.annotate(
                f"{f_ann:.3f} GHz",
                xy=(f_ann, db_ann),
                xytext=(4, 8),
                textcoords="offset points",
                fontsize=6.5,
                color=color,
                rotation=90,
                va="bottom",
            )

    ax.set_xlabel("Frequency (GHz)")
    ax.set_ylabel(f"|S{m+1}{n+1}| (dB)")
    ax.set_title(title or f"Transmission spectrum — S{m+1}{n+1}")
    ax.legend(framealpha=0.9, edgecolor="0.8")
    ax.xaxis.set_minor_locator(mticker.AutoMinorLocator())
    ax.yaxis.set_minor_locator(mticker.AutoMinorLocator())

    # Light grid to mimic VNA display
    ax.grid(True, which="major", linestyle="-",  linewidth=0.4, alpha=0.5)
    ax.grid(True, which="minor", linestyle="--", linewidth=0.2, alpha=0.3)

    _save_or_show(fig, save_path)


# ──────────────────────────────────────────────────────────────────────────────
# 4. plot_error_heatmap
# ──────────────────────────────────────────────────────────────────────────────

# Valid error metric keys and their display strings
_METRIC_LABELS = {
    "f0_s11":       "f₀ error S11 (%)",
    "f0_s22":       "f₀ error S22 (%)",
    "kappa_s11":    "κ error S11 (%)",
    "kappa_s22":    "κ error S22 (%)",
    "shift_mode1":  "Mode 1 shift error (%)",
    "shift_mode2":  "Mode 2 shift error (%)",
    "shift_mode3":  "Mode 3 shift error (%)",
    "shift_max":    "Max shift error (%)",
}


def _extract_error(result: dict, model: str, metric: str) -> float:
    """
    Pull a signed percent error from an analyze_system() result dict.

    Signed error = (model_value - cpw_value) / |cpw_value| * 100

    Parameters
    ----------
    result : dict
        Single analyze_system() output.
    model : str
        One of "optimized", "foster", "analytical".
    metric : str
        One of the keys in _METRIC_LABELS.

    Returns
    -------
    float
        Signed percent error.  NaN if the key is missing.
    """
    label = _MODEL_LABELS.get(model, model)

    if metric in ("f0_s11", "f0_s22"):
        port    = metric.split("_")[1].upper()
        cpw_val = result.get(f"CPW f0 {port} (GHz)",       np.nan)
        mod_val = result.get(f"{label} f0 {port} (GHz)",   np.nan)

    elif metric in ("kappa_s11", "kappa_s22"):
        port    = metric.split("_")[1].upper()
        cpw_val = result.get(f"CPW kappa {port} (MHz)",     np.nan)
        mod_val = result.get(f"{label} kappa {port} (MHz)", np.nan)

    elif metric in ("shift_mode1", "shift_mode2", "shift_mode3"):
        idx = int(metric[-1]) - 1
        cpw_shifts = result.get("CPW all shifts (GHz)", np.array([np.nan]*3))
        mod_shifts = result.get(f"{label} all shifts (GHz)", np.array([np.nan]*3))
        cpw_val = cpw_shifts[idx] if len(cpw_shifts) > idx else np.nan
        mod_val = mod_shifts[idx] if len(mod_shifts) > idx else np.nan

    elif metric == "shift_max":
        cpw_shifts = result.get("CPW all shifts (GHz)", np.array([np.nan]*3))
        mod_shifts = result.get(f"{label} all shifts (GHz)", np.array([np.nan]*3))
        # Find the mode with the largest CPW shift (most constrained)
        if np.all(np.isnan(cpw_shifts)):
            return np.nan
        idx     = np.nanargmax(np.abs(cpw_shifts))
        cpw_val = cpw_shifts[idx]
        mod_val = mod_shifts[idx]

    else:
        return np.nan

    if np.isnan(cpw_val) or np.isnan(mod_val) or cpw_val == 0:
        return np.nan

    return (mod_val - cpw_val) / abs(cpw_val) * 100


def plot_error_heatmap(
    results_grid: list[list[dict]],
    param1_values,
    param2_values,
    param1_label: str,
    param2_label: str,
    model: str,
    metric: str,
    param1_scale: float = 1.0,
    param2_scale: float = 1.0,
    param1_unit: str = "",
    param2_unit: str = "",
    vmin: float = None,
    vmax: float = None,
    cmap: str = "RdBu_r",
    title: str = None,
    ax=None,
    save_path: str = None,
):
    """
    Heatmap of signed percent error from a 2D grid of analyze_system() runs.

    Parameters
    ----------
    results_grid : list[list[dict]]
        2D list of analyze_system() output dicts.
        Shape: [len(param1_values)][len(param2_values)].
        Axis 0 = param1 (y-axis), axis 1 = param2 (x-axis).
    param1_values, param2_values : array-like
        The swept parameter values.  param1 → y-axis, param2 → x-axis.
    param1_label, param2_label : str
        Human-readable axis labels (e.g. "Cc1", "Load frequency").
    model : str
        "optimized", "foster", or "analytical".
    metric : str
        One of: "f0_s11", "f0_s22", "kappa_s11", "kappa_s22",
                "shift_mode1", "shift_mode2", "shift_mode3", "shift_max".
    param1_scale, param2_scale : float
        Multiply parameter values before display (e.g. 1e15 to show fF).
    param1_unit, param2_unit : str
        Unit string appended to axis labels (e.g. "fF", "GHz").
    vmin, vmax : float, optional
        Colour scale limits.  If None, uses symmetric limits around zero
        based on the data range.
    cmap : str
        Matplotlib colormap name.  Default "RdBu_r" (red = positive error,
        blue = negative error, white = zero).
    title : str, optional
    ax : matplotlib.axes.Axes, optional
        If provided, draws into this axes (used by plot_error_heatmap_trio).
    save_path : str, optional
        Only used when ax is None (i.e. this is a standalone call).

    Returns
    -------
    matplotlib.image.AxesImage
        The image object (needed by plot_error_heatmap_trio to set shared limits).
    """
    _apply_paper_style()

    p1 = np.asarray(param1_values) * param1_scale
    p2 = np.asarray(param2_values) * param2_scale

    # Build the error matrix
    error_matrix = np.full((len(p1), len(p2)), np.nan)
    for i, row in enumerate(results_grid):
        for j, result in enumerate(row):
            error_matrix[i, j] = _extract_error(result, model, metric)

    # Symmetric colour limits centred on zero if not specified
    if vmin is None and vmax is None:
        abs_max = np.nanmax(np.abs(error_matrix))
        vmin, vmax = -abs_max, abs_max
    elif vmin is None:
        vmin = -vmax
    elif vmax is None:
        vmax = -vmin

    standalone = ax is None
    if standalone:
        fig, ax = plt.subplots(figsize=(5.5, 4.5))

    im = ax.pcolormesh(
        p2, p1, error_matrix,
        cmap=cmap, vmin=vmin, vmax=vmax,
        shading="auto",
    )

    # Axes labels and ticks
    p1_ax_label = f"{param1_label} ({param1_unit})" if param1_unit else param1_label
    p2_ax_label = f"{param2_label} ({param2_unit})" if param2_unit else param2_label
    ax.set_xlabel(p2_ax_label)
    ax.set_ylabel(p1_ax_label)

    metric_label = _METRIC_LABELS.get(metric, metric)
    model_label  = _MODEL_LABELS.get(model, model)
    ax.set_title(title or f"{model_label} — {metric_label}")

    ax.xaxis.set_minor_locator(mticker.AutoMinorLocator())
    ax.yaxis.set_minor_locator(mticker.AutoMinorLocator())

    if standalone:
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.08)
        cb  = plt.colorbar(im, cax=cax)
        cb.set_label(metric_label, fontsize=8)
        cb.ax.yaxis.set_minor_locator(mticker.AutoMinorLocator())
        plt.tight_layout()
        _save_or_show(plt.gcf(), save_path)

    return im


# ──────────────────────────────────────────────────────────────────────────────
# 5. plot_error_heatmap_trio
# ──────────────────────────────────────────────────────────────────────────────

def plot_error_heatmap_trio(
    results_grid: list[list[dict | None]],
    param1_values,
    param2_values,
    param1_label: str,
    param2_label: str,
    metric: str,
    param1_scale: float = 1.0,
    param2_scale: float = 1.0,
    param1_unit: str = "",
    param2_unit: str = "",
    cmap: str = "RdBu_r",
    title: str = None,
    save_path: str = None,
):
    _apply_paper_style()

    models = ["optimized", "foster", "analytical"]

    p1 = np.asarray(param1_values) * param1_scale
    p2 = np.asarray(param2_values) * param2_scale

    all_errors = {}
    for model in models:
        mat = np.full((len(p1), len(p2)), np.nan)
        for i, row in enumerate(results_grid):
            for j, result in enumerate(row):
                if result is None:
                    continue
                try:
                    mat[i, j] = _extract_error(result, model, metric)
                except (KeyError, TypeError, ValueError):
                    mat[i, j] = np.nan
        all_errors[model] = mat

    combined = np.concatenate([m.ravel() for m in all_errors.values()])

    # Handle case where everything is NaN
    if np.all(np.isnan(combined)):
        vmin, vmax = -1.0, 1.0
    else:
        abs_max = np.nanmax(np.abs(combined))
        vmin, vmax = -abs_max, abs_max

    fig, axes = plt.subplots(
        1, 3,
        figsize=(14, 4.5),
        gridspec_kw={"wspace": 0.35},
    )

    images = []
    for ax, model in zip(axes, models):
        im = ax.pcolormesh(
            p2, p1, all_errors[model],
            cmap=cmap, vmin=vmin, vmax=vmax,
            shading="auto",
        )
        images.append(im)

        p1_ax_label = f"{param1_label} ({param1_unit})" if param1_unit else param1_label
        p2_ax_label = f"{param2_label} ({param2_unit})" if param2_unit else param2_label
        ax.set_xlabel(p2_ax_label)
        ax.set_ylabel(p1_ax_label)
        ax.set_title(_MODEL_LABELS[model], fontsize=9)
        ax.xaxis.set_minor_locator(mticker.AutoMinorLocator())
        ax.yaxis.set_minor_locator(mticker.AutoMinorLocator())

    metric_label = _METRIC_LABELS.get(metric, metric)
    fig.subplots_adjust(right=0.88)
    cbar_ax = fig.add_axes([0.91, 0.12, 0.018, 0.75])
    cb = fig.colorbar(images[-1], cax=cbar_ax)
    cb.set_label(metric_label, fontsize=8)
    cb.ax.yaxis.set_minor_locator(mticker.AutoMinorLocator())

    sup_title = title or f"Signed percent error: {metric_label}"
    fig.suptitle(sup_title, fontsize=10, y=1.01)

    if save_path:
        fig.savefig(save_path, dpi=300, bbox_inches="tight")
        print(f"Saved: {save_path}")
    plt.show()

def plot_shift_errors(
    results: dict,
    title: str = None,
    save_path: str = None,
):
    raise NotImplementedError
    """
    Grouped bar chart of hybridized mode shift errors vs CPW reference.

    Parameters
    ----------
    results : dict
        Output from analyze_system().
    title : str, optional
    save_path : str, optional
    """
    _apply_paper_style()

    models = ["Foster", "Optimized", "Analytical"]
    errors = {m: results.get(f"{m} shift errors (%)", [float("nan")]*3) for m in models}

    x     = np.arange(3)
    width = 0.25
    fig, ax = plt.subplots(figsize=(6, 4))

    for i, model in enumerate(models):
        offset = (i - 1) * width
        ax.bar(
            x + offset, errors[model], width,
            color=_MODEL_COLORS[model.lower()],
            label=model, alpha=0.85,
        )

    ax.axhline(0, color="black", linewidth=0.7)
    ax.set_xticks(x)
    ax.set_xticklabels(["Mode 1", "Mode 2", "Mode 3"])
    ax.set_ylabel("Frequency shift error (%)")
    ax.set_title(title or "Hybridized mode shift errors vs CPW reference")
    ax.legend(framealpha=0.9, edgecolor="0.8")
    ax.yaxis.set_minor_locator(mticker.AutoMinorLocator())

    _save_or_show(fig, save_path)


def plot_pole_fit(network: rf.Network, pole: complex):
    freqs = network.f
    S21 = network.s[:, 1, 0]
    s = 1j * 2 * np.pi * freqs

    alpha = -pole.real
    omega_0 = pole.imag
    poles = np.array([complex(-alpha,  omega_0),
                      complex(-alpha, -omega_0)])

    # Reconstruct the fitted S21 from the final coefficients
    phi = np.column_stack([
        1 / (s - poles[0]),
        1 / (s - poles[1]),
        np.ones(len(s))
    ])
    A = np.vstack([phi.real, phi.imag])
    b = np.hstack([S21.real, S21.imag])
    coeffs, *_ = np.linalg.lstsq(A, b, rcond=None)
    S21_fit = phi @ coeffs

    f_GHz = freqs / 1e9

    fig, axes = plt.subplots(1, 3, figsize=(14, 4))
    fig.suptitle(f"Pole fit — f₀={omega_0/(2*np.pi)/1e9:.4f} GHz, "
                 f"Q={omega_0/(2*alpha):.1f}, "
                 f"κ/2π={alpha/np.pi/1e6:.1f} MHz")

    # 1. Magnitude
    ax = axes[0]
    ax.plot(f_GHz, 20*np.log10(np.abs(S21)),     label='Data', lw=2)
    ax.plot(f_GHz, 20*np.log10(np.abs(S21_fit)), label='Fit',  lw=1.5, ls='--')
    ax.axvline(omega_0/(2*np.pi)/1e9, color='r', ls=':', label='f₀')
    ax.set_xlabel('Frequency (GHz)')
    ax.set_ylabel('|S21| (dB)')
    ax.legend()
    ax.set_title('Magnitude')

    # 2. Phase
    ax = axes[1]
    ax.plot(f_GHz, np.angle(S21,     deg=True), label='Data', lw=2)
    ax.plot(f_GHz, np.angle(S21_fit, deg=True), label='Fit',  lw=1.5, ls='--')
    ax.axvline(omega_0/(2*np.pi)/1e9, color='r', ls=':', label='f₀')
    ax.set_xlabel('Frequency (GHz)')
    ax.set_ylabel('Phase (deg)')
    ax.legend()
    ax.set_title('Phase')

    # 3. IQ circle
    ax = axes[2]
    ax.plot(S21.real,     S21.imag,     label='Data', lw=2)
    ax.plot(S21_fit.real, S21_fit.imag, label='Fit',  lw=1.5, ls='--')

    # Mark the resonance point on the circle
    pk = np.argmin(np.abs(freqs - omega_0/(2*np.pi)))
    ax.plot(S21_fit[pk].real, S21_fit[pk].imag, 'ro', ms=8, label='f₀')
    ax.set_xlabel('Re(S21)')
    ax.set_ylabel('Im(S21)')
    ax.set_aspect('equal')
    ax.legend()
    ax.set_title('IQ plane')

    plt.tight_layout()
    plt.show()
    print(f"f0 = {S21_fit[pk].imag}")

def plot_complex_network_comparison(
    network1: rf.Network,
    network2: rf.Network,
    label1: str = "Network 1",
    label2: str = "Network 2",
    color1: str = "blue",
    color2: str = "red",
):
    freqs1 = network1.f
    freqs2 = network2.f
    S21_1 = network1.s[:, 1, 0]
    S21_2 = network2.s[:, 1, 0]
    f_GHz1 = freqs1 / 1e9
    f_GHz2 = freqs2 / 1e9

    fig, axes = plt.subplots(1, 3, figsize=(14, 4))

    # 1. Magnitude
    ax = axes[0]
    ax.plot(f_GHz1, 20*np.log10(np.abs(S21_1)), label=label1, lw=2,   color=color1)
    ax.plot(f_GHz2, 20*np.log10(np.abs(S21_2)), label=label2, lw=1.5, color=color2, ls='--')
    ax.set_xlabel('Frequency (GHz)')
    ax.set_ylabel('|S21| (dB)')
    ax.legend()
    ax.set_title('Magnitude')
    ax.grid(True, alpha=0.3)

    # 2. Phase
    ax = axes[1]
    ax.plot(f_GHz1, np.angle(S21_1, deg=True), label=label1, lw=2,   color=color1)
    ax.plot(f_GHz2, np.angle(S21_2, deg=True), label=label2, lw=1.5, color=color2, ls='--')
    ax.set_xlabel('Frequency (GHz)')
    ax.set_ylabel('Phase (deg)')
    ax.legend()
    ax.set_title('Phase')
    ax.grid(True, alpha=0.3)

    # 3. IQ circle
    ax = axes[2]
    ax.plot(S21_1.real, S21_1.imag, label=label1, lw=2,   color=color1)
    ax.plot(S21_2.real, S21_2.imag, label=label2, lw=1.5, color=color2, ls='--')
    ax.set_xlabel('Re(S21)')
    ax.set_ylabel('Im(S21)')
    ax.set_aspect('equal')
    ax.legend()
    ax.set_title('IQ plane')
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.show()







def upgraded_plot_lom_vs_data_re_im(
    lom_network: rf.Network,
    data_network: rf.Network,
    m: int = 0,
    n: int = 0,
    lom_label: str = "LOM",
    data_label: str = "Data",
) -> None:
    """
    Overlay Re and Im of a fitted LOM network against a reference (data/CPW).

    Parameters
    ----------
    lom_network : rf.Network
        The fitted lumped-element model.
    data_network : rf.Network
        The reference (measured or CPW simulation).
    m, n : int
        S-parameter indices (0-based).
    lom_label, data_label : str
        Legend labels.
    """
    
    plt.rcParams.update(_PLOT_STYLE)
    f_lom  = lom_network.frequency.f / 1e9
    f_data = data_network.frequency.f / 1e9
    s_lom  = lom_network.s[:, m, n]
    s_data = data_network.s[:, m, n]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.3))

    ax1.plot(f_lom,  np.real(s_lom),  color="red",  label="{} Re(S{}{})".format(lom_label,  m+1, n+1))
    ax1.plot(f_data, np.real(s_data), color="blue", label="{} Re(S{}{})".format(data_label, m+1, n+1), linestyle="--")
    ax1.set_xlabel("Frequency (GHz)")
    ax1.set_ylabel("Real")
    ax1.legend()

    ax2.plot(f_lom,  np.imag(s_lom),  color="red",  label="{} Im(S{}{})".format(lom_label,  m+1, n+1))
    ax2.plot(f_data, np.imag(s_data), color="blue", label="{} Im(S{}{})".format(data_label, m+1, n+1), linestyle="--")
    ax2.set_xlabel("Frequency (GHz)")
    ax2.set_ylabel("Imaginary")
    ax2.legend()

    plt.tight_layout()
    plt.show()


def fancy_plot(
    lom_network: rf.Network,
    data_network: rf.Network,
    m: int = 0,
    n: int = 0,
    lom_label: str = "LOM",
    data_label: str = "Data",
) -> None:

    # USC color scheme: cardinal red, gold, dark navy text
    USC_RED    = "#AF1E75"
    USC_GOLD   = "#2765AB"
    USC_DARK   = "#000000"
    GRID_COLOR = "#e0e0e0"

    f_lom  = lom_network.frequency.f / 1e9
    f_data = data_network.frequency.f / 1e9
    s_lom  = lom_network.s[:, m, n]
    s_data = data_network.s[:, m, n]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.3), dpi=300)

    for ax, y_lom, y_data, ylabel in [
        (ax1, np.real(s_lom), np.real(s_data), f"Re(S{m+1}{n+1})"),
        (ax2, np.imag(s_lom), np.imag(s_data), f"Im(S{m+1}{n+1})"),
    ]:
        ax.plot(f_data, y_data, color=USC_GOLD,  linewidth=2.0,
                label=f"{data_label}", linestyle="-")
        ax.plot(f_lom,  y_lom,  color=USC_RED,   linewidth=2.0,
                label=f"{lom_label}",  linestyle="--")

        # Faint grid
        ax.grid(True, color=GRID_COLOR, linewidth=0.8, zorder=0)
        ax.set_axisbelow(True)

        # Large titles and labels
        ylabel_re = rf"$\mathrm{{Re}}(S_{{{m+1}{n+1}}})$"
        ylabel_im = rf"$\mathrm{{Im}}(S_{{{m+1}{n+1}}})$"

        for ax, y_lom, y_data, ylabel in [
            (ax1, np.real(s_lom), np.real(s_data), ylabel_re),
            (ax2, np.imag(s_lom), np.imag(s_data), ylabel_im),
        ]:
            ax.set_title(ylabel, fontsize=16, fontweight="bold", color=USC_DARK, pad=8)
            ax.set_xlabel("Frequency (GHz)", fontsize=13, color=USC_DARK)
            ax.set_ylabel(ylabel, fontsize=13, color=USC_DARK)
            ax.tick_params(labelsize=11, colors=USC_DARK)
            ax.legend(fontsize=11, framealpha=0.9)

        # Spine colors
        for spine in ax.spines.values():
            spine.set_edgecolor(GRID_COLOR)

    # Gold accent line under the figure
    fig.patch.set_facecolor("white")
    plt.tight_layout()
    plt.show()

    # Shared style constants
_COLORS = {
    'data'  : '#1f77b4',   # strong blue
    'fit'   : '#f86d9b',   # strong red
    'f0'    : "#d62728",   # green accent for f0 marker
    'grid'  : '#e0e0e0',
}

def _apply_style(ax, title, xlabel, ylabel):
    ax.set_title(title,   fontsize=24, pad=10)  # bigger, not bold
    ax.set_xlabel(xlabel, fontsize=15)
    ax.set_ylabel(ylabel, fontsize=15)
    ax.tick_params(labelsize=13)
    ax.grid(True, color=_COLORS['grid'], linewidth=0.8, zorder=0)
    ax.set_axisbelow(True)
    ax.legend(fontsize=13, framealpha=0.9)
    for spine in ax.spines.values():
        spine.set_edgecolor(_COLORS['grid'])


def plot_pole_fit(network: rf.Network, pole: complex):
    freqs  = network.f
    S21    = network.s[:, 1, 0]
    s      = 1j * 2 * np.pi * freqs
    alpha  = -pole.real
    omega_0 = pole.imag

    poles = np.array([complex(-alpha, omega_0),
                      complex(-alpha, -omega_0)])

    phi = np.column_stack([
        1 / (s - poles[0]),
        1 / (s - poles[1]),
        np.ones(len(s))
    ])
    A = np.vstack([phi.real, phi.imag])
    b = np.hstack([S21.real, S21.imag])
    coeffs, *_ = np.linalg.lstsq(A, b, rcond=None)
    S21_fit = phi @ coeffs

    f_GHz = freqs / 1e9
    f0_GHz = omega_0 / (2 * np.pi) / 1e9

    fig, axes = plt.subplots(1, 3, figsize=(16, 5), dpi=150)
    fig.suptitle(
        rf"Pole fit — $f_0={f0_GHz:.4f}$ GHz,  "
        rf"$Q={omega_0/(2*alpha):.1f}$,  "
        rf"$\kappa/2\pi={alpha/np.pi/1e6:.1f}$ MHz",
        fontsize=16, fontweight='bold'
    )

    # 1. Magnitude
    ax = axes[0]
    ax.plot(f_GHz, 20*np.log10(np.abs(S21)),
            color=_COLORS['data'], lw=2.5, label='Data')
    ax.plot(f_GHz, 20*np.log10(np.abs(S21_fit)),
            color=_COLORS['fit'],  lw=2.0, ls='--', label='Fit')
    ax.axvline(f0_GHz, color=_COLORS['f0'], ls=':', lw=1.5, label=r'$f_0$')
    _apply_style(ax,
                 title=r'$|S_{21}|$',
                 xlabel=r'Frequency (GHz)',
                 ylabel=r'$|S_{21}|$ (dB)')

    # 2. Phase
    ax = axes[1]
    ax.plot(f_GHz, np.angle(S21,     deg=True),
            color=_COLORS['data'], lw=2.5, label='Data')
    ax.plot(f_GHz, np.angle(S21_fit, deg=True),
            color=_COLORS['fit'],  lw=2.0, ls='--', label='Fit')
    ax.axvline(f0_GHz, color=_COLORS['f0'], ls=':', lw=1.5, label=r'$f_0$')
    _apply_style(ax,
                 title=r'Phase of $S_{21}$',
                 xlabel=r'Frequency (GHz)',
                 ylabel=r'Phase (deg)')

    # 3. IQ circle
    ax = axes[2]
    ax.plot(S21.real,     S21.imag,
            color=_COLORS['data'], lw=2.5, label='Data')
    ax.plot(S21_fit.real, S21_fit.imag,
            color=_COLORS['fit'],  lw=2.0, ls='--', label='Fit')
    pk = np.argmin(np.abs(freqs - omega_0 / (2 * np.pi)))
    ax.plot(S21_fit[pk].real, S21_fit[pk].imag,
            'o', color=_COLORS['f0'], ms=10, label=r'$f_0$')
    ax.set_aspect('equal')
    _apply_style(ax,
                 title=r'IQ Plane',
                 xlabel=r'$\mathrm{Re}(S_{21})$',
                 ylabel=r'$\mathrm{Im}(S_{21})$')

    plt.tight_layout()
    plt.show()
    print(f"f0 = {omega_0 / (2 * np.pi) / 1e9:.6f} GHz")


def plot_complex_network_comparison(
    network1: rf.Network,
    network2: rf.Network,
    label1: str = "Network 1",
    label2: str = "Network 2",
    color1: str = '#1f77b4',
    color2: str = '#d62728',
):
    freqs1 = network1.f
    freqs2 = network2.f
    S21_1  = network1.s[:, 1, 0]
    S21_2  = network2.s[:, 1, 0]
    f_GHz1 = freqs1 / 1e9
    f_GHz2 = freqs2 / 1e9

    fig, axes = plt.subplots(1, 3, figsize=(16, 5), dpi=150)

    # 1. Magnitude
    ax = axes[0]
    ax.plot(f_GHz1, 20*np.log10(np.abs(S21_1)),
            color=color1, lw=2.5, label=label1)
    ax.plot(f_GHz2, 20*np.log10(np.abs(S21_2)),
            color=color2, lw=2.0, ls='--', label=label2)
    _apply_style(ax,
                 title=r'$|S_{21}|$',
                 xlabel=r'Frequency (GHz)',
                 ylabel=r'$|S_{21}|$ (dB)')

    # 2. Phase
    ax = axes[1]
    ax.plot(f_GHz1, np.angle(S21_1, deg=True),
            color=color1, lw=2.5, label=label1)
    ax.plot(f_GHz2, np.angle(S21_2, deg=True),
            color=color2, lw=2.0, ls='--', label=label2)
    _apply_style(ax,
                 title=r'Phase of $S_{21}$',
                 xlabel=r'Frequency (GHz)',
                 ylabel=r'Phase (deg)')

    # 3. IQ circle
    ax = axes[2]
    ax.plot(S21_1.real, S21_1.imag,
            color=color1, lw=2.5, label=label1)
    ax.plot(S21_2.real, S21_2.imag,
            color=color2, lw=2.0, ls='--', label=label2)
    ax.set_aspect('equal')
    _apply_style(ax,
                 title=r'IQ Plane',
                 xlabel=r'$\mathrm{Re}(S_{21})$',
                 ylabel=r'$\mathrm{Im}(S_{21})$')

    plt.tight_layout()
    plt.show()


def fancy_plot_all_models(
    networks: dict[str, rf.Network],
    m: int = 0,
    n: int = 0,
    quantity: str = "re",
    title: str = None,
) -> None:
    """
    Plot one S-parameter quantity for multiple networks on a single axes.
    Useful for comparing CPW, FosterFit, OptimizedFit, and AnalyticalFit
    in one call.

    Parameters
    ----------
    networks : dict[str, rf.Network]
        Mapping of label → network, e.g.
        {"CPW": cpw_net, "Foster": foster_net, "Optimized": opt_net}
    m, n : int
        S-parameter indices (0-based).
    quantity : {"re", "im", "db", "abs"}
        Which quantity to plot.
    title : str, optional
        Plot title.
    """
    # --- Talk-slide style overrides ---
    SLIDE_STYLE = {
        "font.size":        28,
        "axes.titlesize":   32,
        "axes.labelsize":   30,
        "xtick.labelsize":  24,
        "ytick.labelsize":  24,
        "legend.fontsize":  26,
        "lines.linewidth":   3,
        "figure.dpi":       300,
    }
    #plt.rcParams.update(_PLOT_STYLE)
    plt.rcParams.update(SLIDE_STYLE)   # slide overrides win

    _extractors = {
        "re":  lambda s: np.real(s),
        "im":  lambda s: np.imag(s),
        "db":  lambda s: 20 * np.log10(np.maximum(np.abs(s), 1e-300)),
        "abs": lambda s: np.abs(s),
    }
    _ylabels = {"re": "Re", "im": "Im", "db": "|S| (dB)", "abs": "|S|"}

    if quantity not in _extractors:
        raise ValueError("quantity must be one of: {}".format(list(_extractors)))

    extract = _extractors[quantity]

    # Wider figure so labels don't crowd
    fig, ax = plt.subplots(figsize=(14, 7))

    for label, net in networks.items():
        f = net.frequency.f / 1e9
        s = net.s[:, m, n]
        ax.plot(f, extract(s), label=label)

    ax.set_xlabel("Frequency (GHz)")
    ax.set_ylabel("{} (S{}{})".format(_ylabels[quantity], m + 1, n + 1))
    ax.legend(framealpha=0.9, edgecolor="gray")

    if title:
        ax.set_title(title)

    plt.tight_layout()
    plt.show()