from __future__ import annotations

"""
models/optimized.py
-------------------
OptimizedFit: numerical optimization over S11 + S22 residuals to find
the best-fit effective L and C for a CPW resonator.

Algorithm summary
-----------------
Stage 1 — coarse scan (fit_Ceff_Leff):
    Scan ω0 in a window around the measured resonance.  For each ω0,
    use Nelder-Mead to optimise Leff with Ceff forced by the resonance
    constraint ω0² * L * (C + Cc1 + Cc2) = 1.  Keep the best (L, C).

Stage 2 — windowed refinement (least_squares):
    Build separate frequency windows around the S11 and S22 resonances.
    Minimise the complex residuals of the 2-port LC network against the
    CPW reference network, weighted by resonance depth to prevent one
    shallow trace from dominating.
"""

import logging
from dataclasses import dataclass, field

import numpy as np
import skrf as rf
from scipy.optimize import minimize, least_squares

from simpleLOMs.models.base import BaseFit
from simpleLOMs.analysis import (
    resonance_from_res11,
    fwhm_from_res11,
    fwhm_from_trace_db,
    resonance,
    resonance_from_s_max,
)
from simpleLOMs.networks.lc import (
    lc_resonator_network,
    lc_resonator_network_2port,
    lc_resonator_network_with_grounds_2port,
)

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Helper dataclass for optimizer hyperparameters
# ---------------------------------------------------------------------------

@dataclass
class OptimizationConfig:
    """
    Tuning knobs for the two-stage OptimizedFit algorithm.

    You rarely need to change these — the defaults work well for typical
    CPW resonators in the 4–10 GHz range.  They are collected here so
    that a caller can pass a single object instead of many keyword args.

    Attributes
    ----------
    w0_window_frac : float
        Half-width of the ω0 scan window as a fraction of ω0_guess.
        e.g. 0.005 → scan ω0 ± 0.5 %.
    n_w0 : int
        Number of ω0 grid points in the coarse scan.
    n_dense : int
        Number of data points sampled around resonance in Stage 1.
    n_kappa : float
        Half-width of the dense sampling window in units of κ (linewidth).
    n_widths : float
        Half-width of each residual window in Stage 2, in units of κ.
    max_nfev : int
        Maximum function evaluations for the Stage 2 least_squares call.
    verbose : bool
        If True, least_squares prints iteration progress.
    """
    w0_window_frac: float = 0.00001
    n_w0:          int   = 20
    n_dense:       int   = 100
    n_kappa:       float = 0.75
    n_widths:      float = 1.0
    max_nfev:      int   = 100
    verbose:       bool  = False


# ---------------------------------------------------------------------------
# OptimizedFit class
# ---------------------------------------------------------------------------

class OptimizedFit(BaseFit):
    """
    Lumped LC model via two-stage numerical optimisation.

    Parameters
    ----------
    config : OptimizationConfig, optional
        Hyperparameters for the fitting algorithm.
        Uses sensible defaults if not provided.

    Attributes
    ----------
    L : float or None
        Fitted effective inductance in Henries.  None before fit().
    C : float or None
        Fitted effective capacitance in Farads.  None before fit().
    scan_results : np.ndarray or None
        Diagnostic array from the Stage 1 ω0 scan, shape (n_w0, 4):
        columns are [ω0, Ceff, Leff, sse].

    Examples
    --------
    ::

        from simpleLOMs.models.optimized import OptimizedFit, OptimizationConfig
        import skrf as rf

        freq       = rf.Frequency(4e9, 10e9, 10_001, unit="Hz")
        data_ntw   = ...  # rf.Network of the CPW S11 measurement
        config     = OptimizationConfig(n_w0=30, w0_window_frac=0.01)

        model = OptimizedFit(config=config)
        model.fit(freq, data_ntw=data_ntw, Cc1=5e-15, Cc2=5e-15, Z0=50)

        net = model.get_network(freq, Cc1=5e-15, Cc2=5e-15, Z0=50)
    """

    def __init__(self, config: OptimizationConfig = None):
        self.config: OptimizationConfig    = config if config is not None else OptimizationConfig()
        self.L:            float | None    = None
        self.C:            float | None    = None
        self.scan_results: np.ndarray | None = None

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _build_sparse_data_points(
        self,
        data_ntw: rf.Network,
        n_dense: int,
        n_kappa: float,
    ) -> list[tuple[float, float, float]]:
        """
        Build a sparse, physically-motivated set of (ω, Re(S11), Im(S11)) tuples.

        Always includes the minimum of Re(S11) (resonance point).
        Optionally adds a dense cluster of points around the resonance.

        Parameters
        ----------
        data_ntw : rf.Network
            Single-port S11 measurement or model.
        n_dense : int
            How many extra points to add around resonance. 0 disables.
        n_kappa : float
            Half-width of the dense window in units of linewidth κ.

        Returns
        -------
        list of (omega, ReS11, ImS11)
        """
        f_hz  = data_ntw.frequency.f
        w     = 2 * np.pi * f_hz
        S11   = data_ntw.s[:, 0, 0]
        re_s11 = np.real(S11)
        im_s11 = np.imag(S11)

        # Resonance point (minimum of Re(S11))
        idx_re_min = np.argmin(re_s11)
        data_points = [(w[idx_re_min], re_s11[idx_re_min], im_s11[idx_re_min])]

        if n_dense > 0:
            f0    = f_hz[idx_re_min]
            kappa = fwhm_from_res11(data_ntw)

            f_dense = np.linspace(
                f0 - n_kappa * kappa,
                f0 + (n_kappa / 2) * kappa,
                n_dense + 2,
            )[1:-1]

            re_dense = np.interp(f_dense, f_hz, re_s11)
            im_dense = np.interp(f_dense, f_hz, im_s11)

            for f_pt, r, i in zip(f_dense, re_dense, im_dense):
                data_points.append((2 * np.pi * f_pt, r, i))

        return data_points

    def _fit_Ceff_Leff(
        self,
        w0_guess: float,
        Cc1: float,
        Cc2: float,
        Z0: float,
        data_points: list,
        freq: rf.Frequency,
    ) -> tuple[float, float, float, np.ndarray]:
        """
        Stage 1: scan ω0 and fit Leff by Nelder-Mead with Ceff constrained.

        Returns (Ceff_best, Leff_best, w0_best, scan_results).
        """
        cfg = self.config
        dp  = np.asarray(data_points, dtype=float)
        freqs  = dp[:, 0]
        target = dp[:, 1] + 1j * dp[:, 2]

        def ZCc1(w):   return 1.0 / (1j * w * Cc1)
        def ZCc2(w):   return 1.0 / (1j * w * Cc2)
        def ZL(w, Le): return 1j * w * Le
        def ZC(w, Ce): return 1.0 / (1j * w * Ce)
        def ZLC(w, Le, Ce):
            return 1.0 / (1.0 / ZL(w, Le) + 1.0 / ZC(w, Ce))
        def Zshunt(w, Le, Ce):
            return 1.0 / (1.0 / ZLC(w, Le, Ce) + 1.0 / ZCc2(w))
        def Zin(w, Le, Ce):
            return ZCc1(w) + Zshunt(w, Le, Ce)
        def S11(w, Le, Ce):
            z = Zin(w, Le, Ce)
            return (z - Z0) / (z + Z0)

        w0_min  = w0_guess * (1.0 - cfg.w0_window_frac)
        w0_max  = w0_guess * (1.0 + cfg.w0_window_frac)
        w0_grid = np.linspace(w0_min, w0_max, int(cfg.n_w0))

        best    = None
        results = []

        for w0 in w0_grid:
            def ceff_of_L(Le):
                return 1.0 / (Le * (w0 ** 2)) - (Cc1 + Cc2)

            Le_max   = 1.0 / ((Cc1 + Cc2) * (w0 ** 2))
            Le_upper = 0.999 * Le_max
            x0_local = 0.5 * Le_upper

            def obj(Le_arr):
                Le = float(Le_arr[0])
                if not (0.0 < Le < Le_upper):
                    dist = (0.0 - Le) if Le <= 0.0 else (Le - Le_upper)
                    return 1e6 * (1.0 + dist ** 2)
                Ce = ceff_of_L(Le)
                if Ce <= 0.0 or not np.isfinite(Ce):
                    return 1e6
                pred  = np.array([S11(w, Le, Ce) for w in freqs])
                resid = pred - target
                return float(np.sum(resid.real ** 2 + resid.imag ** 2))

            fit = minimize(
                obj,
                x0=np.array([x0_local]),
                method="Nelder-Mead",
                options=dict(maxiter=3000, xatol=1e-30, fatol=1e-30, disp=False),
            )

            Le_hat = float(np.clip(fit.x[0], np.nextafter(0.0, 1.0), Le_upper))
            Ce_hat = ceff_of_L(Le_hat)
            sse    = float(fit.fun)
            results.append((w0, Ce_hat, Le_hat, sse))

            if best is None or sse < best[-1]:
                best = (w0, Ce_hat, Le_hat, sse)

        w0_best, Ceff_best, Leff_best, _ = best
        return Ceff_best, Leff_best, w0_best, np.array(results)

    def _make_windowed_residuals(
        self,
        freq: rf.Frequency,
        data_ntw: rf.Network,
        Cc1: float,
        Cc2: float,
        f0_s11: float,
        kappa_s11: float,
        f0_s22: float,
        kappa_s22: float,
        Z0: float,
        eps: float = 1e-18,
    ):
        """
        Stage 2: build the windowed residual function for least_squares.

        Returns a callable residuals(x) where x = [Leff, Ceff].
        """
        cfg = self.config

        # Interpolate data onto the optimisation frequency grid
        f_old = data_ntw.frequency.f
        f_new = freq.f

        if len(f_old) == len(f_new) and np.allclose(f_old, f_new):
            data_aligned = data_ntw
        else:
            S_old = data_ntw.s
            S_new = np.zeros((len(f_new), 2, 2), dtype=complex)
            for i in range(2):
                for j in range(2):
                    yr = np.interp(f_new, f_old, np.real(S_old[:, i, j]))
                    yi = np.interp(f_new, f_old, np.imag(S_old[:, i, j]))
                    S_new[:, i, j] = yr + 1j * yi
            z0_old = data_ntw.z0
            if z0_old.ndim == 1:
                z0_new = np.tile(z0_old[0], (len(f_new), 2))
            else:
                z0_new = np.zeros((len(f_new), 2))
                for p in range(2):
                    z0_new[:, p] = np.interp(f_new, f_old, z0_old[:, p])
            data_aligned = rf.Network(frequency=freq, s=S_new, z0=z0_new)

        f = freq.f
        n = cfg.n_widths
        mask11 = (f > f0_s11 - n * kappa_s11) & (f < f0_s11 + n * kappa_s11)
        mask22 = (f > f0_s22 - n * kappa_s22) & (f < f0_s22 + n * kappa_s22)

        if not np.any(mask11):
            raise ValueError("S11 mask is empty: check f0_s11/kappa_s11/n_widths.")
        if not np.any(mask22):
            raise ValueError("S22 mask is empty: check f0_s22/kappa_s22/n_widths.")

        freq11   = rf.Frequency.from_f(f[mask11], unit="Hz")
        freq22   = rf.Frequency.from_f(f[mask22], unit="Hz")
        s11_data = data_aligned.s[mask11, 0, 0]
        s22_data = data_aligned.s[mask22, 1, 1]

        def _depth(s_complex):
            m = np.real(s_complex)
            return max(float(np.max(m) - np.min(m)), eps)

        S11_scale = _depth(s11_data)
        S22_scale = _depth(s22_data)
        depth_floor = max(S11_scale, S22_scale)
        S11_scale = max(S11_scale, depth_floor)
        S22_scale = max(S22_scale, depth_floor)

        logger.debug("S11 depth scale=%.4e, N11=%d", S11_scale, mask11.sum())
        logger.debug("S22 depth scale=%.4e, N22=%d", S22_scale, mask22.sum())

        def residuals(x):
            Le, Ce = float(x[0]), float(x[1])
            lom11  = lc_resonator_network_2port(Le, Ce, Cc1, Cc2, freq11, Z0)
            lom22  = lc_resonator_network_2port(Le, Ce, Cc1, Cc2, freq22, Z0)
            r11    = lom11.s[:, 0, 0] - s11_data
            r22    = lom22.s[:, 1, 1] - s22_data
            r11_sc = np.concatenate([r11.real, r11.imag]) / S11_scale
            r22_sc = np.concatenate([r22.real, r22.imag]) / S22_scale
            return np.concatenate([r11_sc, r22_sc])

        return residuals

    def _make_windowed_residuals2(
        self,
        freq: rf.Frequency,
        data_ntw: rf.Network,
        Cc1: float,
        Cc2: float,
        f0_s11: float,
        kappa_s11: float,
        f0_s22: float,
        kappa_s22: float,
        Z0: float,
        eps: float = 1e-18,
        phase=None,
    ):
        cfg = self.config
        f_old = data_ntw.frequency.f
        f_new = freq.f
        if len(f_old) == len(f_new) and np.allclose(f_old, f_new):
            data_aligned = data_ntw
        else:
            S_old = data_ntw.s
            S_new = np.zeros((len(f_new), 2, 2), dtype=complex)
            for i in range(2):
                for j in range(2):
                    yr = np.interp(f_new, f_old, np.real(S_old[:, i, j]))
                    yi = np.interp(f_new, f_old, np.imag(S_old[:, i, j]))
                    S_new[:, i, j] = yr + 1j * yi
            z0_old = data_ntw.z0
            if z0_old.ndim == 1:
                z0_new = np.tile(z0_old[0], (len(f_new), 2))
            else:
                z0_new = np.zeros((len(f_new), 2))
                for p in range(2):
                    z0_new[:, p] = np.interp(f_new, f_old, z0_old[:, p])
            data_aligned = rf.Network(frequency=freq, s=S_new, z0=z0_new)

        f = freq.f
        n = cfg.n_widths
        mask11 = (f > f0_s11 - n * kappa_s11) & (f < f0_s11 + n * kappa_s11)
        mask22 = (f > f0_s22 - n * kappa_s22) & (f < f0_s22 + n * kappa_s22)
        mask21 =(f > f0_s11 - n * kappa_s11) & (f < f0_s11 + n * kappa_s11)
        mask12 =(f > f0_s11 - n * kappa_s11) & (f < f0_s11 + n * kappa_s11)


        if not np.any(mask11):
            raise ValueError("S11 mask is empty: check f0_s11/kappa_s11/n_widths.")
        if not np.any(mask22):
            raise ValueError("S22 mask is empty: check f0_s22/kappa_s22/n_widths.")
        if not np.any(mask21):
            raise ValueError("S21 mask is empty: no frequencies outside S11/S22 windows.")

        freq11 = rf.Frequency.from_f(f[mask11], unit="Hz")
        freq22 = rf.Frequency.from_f(f[mask22], unit="Hz")
        freq21 = rf.Frequency.from_f(f[mask21], unit="Hz")
        freq12 = rf.Frequency.from_f(f[mask21], unit="Hz")

        s11_data = data_aligned.s[mask11, 0, 0]
        s22_data = data_aligned.s[mask22, 1, 1]
        s21_data = data_aligned.s[mask21, 1, 0]
        s12_data = data_aligned.s[mask12, 0, 1]

        def _depth(s_complex):
            m = np.abs(s_complex)
            return max(float(np.max(m) - np.min(m)), eps)

        S11_scale = _depth(s11_data)
        S22_scale = _depth(s22_data)
        depth_floor = max(S11_scale, S22_scale)
        S11_scale = max(S11_scale, depth_floor)
        S22_scale = max(S22_scale, depth_floor)
        S21_scale = max(_depth(s21_data), depth_floor)
        S12_scale = max(_depth(s12_data), depth_floor)

        logger.debug("S11 depth scale=%.4e, N11=%d", S11_scale, mask11.sum())
        logger.debug("S22 depth scale=%.4e, N22=%d", S22_scale, mask22.sum())
        logger.debug("S21 depth scale=%.4e, N21=%d", S21_scale, mask21.sum())

        def residuals(x):
            Le, Ce = float(x[0]), float(x[1])
            lom11 = lc_resonator_network_2port(Le, Ce, Cc1, Cc2, freq11, Z0)
            lom22 = lc_resonator_network_2port(Le, Ce, Cc1, Cc2, freq22, Z0)
            lom21 = lc_resonator_network_2port(Le, Ce, Cc1, Cc2, freq21, Z0)
            lom12 = lc_resonator_network_2port(Le, Ce, Cc1, Cc2, freq12, Z0)
            if phase is not None:
                r11 = lom11.s[:, 0, 0] * phase - s11_data
                r22 = lom22.s[:, 1, 1] * phase - s22_data
                r21 = lom21.s[:, 1, 0] * phase - s21_data
                r12 = lom12.s[:, 0, 1] * phase - s12_data

                r11_sc = np.concatenate([r11.real, r11.imag]) / S11_scale
                r22_sc = np.concatenate([r22.real, r22.imag]) / S22_scale
                r21_sc = np.concatenate([r21.real, r21.imag]) / S21_scale
                r12_sc = np.concatenate([r12.real, r12.imag]) / S12_scale
            else: 
                r11 = lom11.s[:, 0, 0] - s11_data
                r22 = lom22.s[:, 1, 1] - s22_data
                r21 = lom21.s[:, 1, 0] - s21_data
                r12 = lom12.s[:, 0, 1] - s12_data

                r11_sc = np.concatenate([r11.real, r11.imag]) / S11_scale
                r22_sc = np.concatenate([r22.real, r22.imag]) / S22_scale
                r21_sc = np.concatenate([r21.real, r21.imag]) / S21_scale
                r12_sc = np.concatenate([r12.real, r12.imag]) / S21_scale

            return np.concatenate([r11_sc, r22_sc, r21_sc, r12_sc])
            # lom21  = lc_resonator_network_2port(Le, Ce, Cc1, Cc2, freq21, Z0)
            # r21_sc = np.concatenate([(lom21.s[:, 1, 0] - s21_data).real,
            #                         (lom21.s[:, 1, 0] - s21_data).imag]) / S21_scale
            # r12_sc = np.concatenate([(lom21.s[:, 0, 1] - s12_data).real,
            #                         (lom21.s[:, 0, 1] - s12_data).imag]) / S21_scale
            # return np.concatenate([r11_sc, r22_sc, r21_sc, r12_sc])

        return residuals
    
    def _make_windowed_residuals_complex(
        self,
        freq: rf.Frequency,
        data_ntw: rf.Network,
        Cc1: float,
        Cc2: float,
        f0_s11: float,
        kappa_s11: float,
        f0_s22: float,
        kappa_s22: float,
        f0_s12: float,
        f0_s21: float,
        Z0: float,
        eps: float = 1e-18,
    ):
        cfg = self.config
        f_old = data_ntw.frequency.f
        f_new = freq.f
        if len(f_old) == len(f_new) and np.allclose(f_old, f_new):
            data_aligned = data_ntw
        else:
            S_old = data_ntw.s
            S_new = np.zeros((len(f_new), 2, 2), dtype=complex)
            for i in range(2):
                for j in range(2):
                    yr = np.interp(f_new, f_old, np.real(S_old[:, i, j]))
                    yi = np.interp(f_new, f_old, np.imag(S_old[:, i, j]))
                    S_new[:, i, j] = yr + 1j * yi
            z0_old = data_ntw.z0
            if z0_old.ndim == 1:
                z0_new = np.tile(z0_old[0], (len(f_new), 2))
            else:
                z0_new = np.zeros((len(f_new), 2))
                for p in range(2):
                    z0_new[:, p] = np.interp(f_new, f_old, z0_old[:, p])
            data_aligned = rf.Network(frequency=freq, s=S_new, z0=z0_new)

        f = freq.f
        n = cfg.n_widths
        mask11 = (f > f0_s11 - n * kappa_s11) & (f < f0_s11 + n * kappa_s11)
        mask22 = (f > f0_s22 - n * kappa_s22) & (f < f0_s22 + n * kappa_s22)
        mask21 = (f > f0_s22 - n * kappa_s22) & (f < f0_s22 + n * kappa_s22)

        if not np.any(mask11):
            raise ValueError("S11 mask is empty: check f0_s11/kappa_s11/n_widths.")
        if not np.any(mask22):
            raise ValueError("S22 mask is empty: check f0_s22/kappa_s22/n_widths.")
        if not np.any(mask21):
            raise ValueError("S21 mask is empty: no frequencies outside S11/S22 windows.")

        freq11 = rf.Frequency.from_f(f[mask11], unit="Hz")
        freq22 = rf.Frequency.from_f(f[mask22], unit="Hz")
        freq21 = rf.Frequency.from_f(f[mask21], unit="Hz")

        s11_data = data_aligned.s[mask11, 0, 0]
        s22_data = data_aligned.s[mask22, 1, 1]
        s21_data = data_aligned.s[mask21, 1, 0]
        s12_data = data_aligned.s[mask21, 0, 1]

        def _depth(s_complex):
            m = np.abs(s_complex)
            return max(float(np.max(m) - np.min(m)), eps)

        S11_scale = _depth(s11_data)
        S22_scale = _depth(s22_data)
        depth_floor = max(S11_scale, S22_scale)
        S11_scale = max(S11_scale, depth_floor)
        S22_scale = max(S22_scale, depth_floor)
        S21_scale = max(_depth(s21_data), depth_floor)

        logger.debug("S11 depth scale=%.4e, N11=%d", S11_scale, mask11.sum())
        logger.debug("S22 depth scale=%.4e, N22=%d", S22_scale, mask22.sum())
        logger.debug("S21 depth scale=%.4e, N21=%d", S21_scale, mask21.sum())

        def residuals(x):
            Le, Ce, phi = float(x[0]), float(x[1]), float(x[2])
            phase = np.exp(1j * phi)

            lom11 = lc_resonator_network_2port(Le, Ce, Cc1, Cc2, freq11, Z0)
            lom22 = lc_resonator_network_2port(Le, Ce, Cc1, Cc2, freq22, Z0)
            lom21 = lc_resonator_network_2port(Le, Ce, Cc1, Cc2, freq21, Z0)

            r11 = lom11.s[:, 0, 0] * phase - s11_data
            r22 = lom22.s[:, 1, 1] * phase - s22_data
            r21 = lom21.s[:, 1, 0] * phase - s21_data
            #r12 = lom21.s[:, 0, 1] * phase - s12_data

            r11_sc = np.concatenate([r11.real, r11.imag]) / S11_scale
            r22_sc = np.concatenate([r22.real, r22.imag]) / S22_scale
            r21_sc = np.concatenate([r21.real, r21.imag]) / S21_scale
            #r12_sc = np.concatenate([r12.real, r12.imag]) / S21_scale

            return np.concatenate([r11_sc, r22_sc, r21_sc])
            #return np.concatenate([r11_sc, r22_sc, r21_sc, r12_sc])

        return residuals
    
    def _make_windowed_residuals_z(
        self,
        freq: rf.Frequency,
        data_ntw: rf.Network,
        Cc1: float,
        Cc2: float,
        f0_s11: float,
        kappa_s11: float,
        f0_s22: float,
        kappa_s22: float,
        Z0: float,
        eps: float = 1e-18,
        ):
        cfg = self.config
        f_old = data_ntw.frequency.f
        f_new = freq.f
        if len(f_old) == len(f_new) and np.allclose(f_old, f_new):
            data_aligned = data_ntw
        else:
            S_old = data_ntw.s
            S_new = np.zeros((len(f_new), 2, 2), dtype=complex)
            for i in range(2):
                for j in range(2):
                    yr = np.interp(f_new, f_old, np.real(S_old[:, i, j]))
                    yi = np.interp(f_new, f_old, np.imag(S_old[:, i, j]))
                    S_new[:, i, j] = yr + 1j * yi
            z0_old = data_ntw.z0
            if z0_old.ndim == 1:
                z0_new = np.tile(z0_old[0], (len(f_new), 2))
            else:
                z0_new = np.zeros((len(f_new), 2))
                for p in range(2):
                    z0_new[:, p] = np.interp(f_new, f_old, z0_old[:, p])
            data_aligned = rf.Network(frequency=freq, s=S_new, z0=z0_new)

        f = freq.f
        n = cfg.n_widths
        mask11 = (f > f0_s11 - n * kappa_s11) & (f < f0_s11 + n * kappa_s11)
        mask22 = (f > f0_s22 - n * kappa_s22) & (f < f0_s22 + n * kappa_s22)

        if not np.any(mask11):
            raise ValueError("Z11 mask is empty: check f0_s11/kappa_s11/n_widths.")
        if not np.any(mask22):
            raise ValueError("Z22 mask is empty: check f0_s22/kappa_s22/n_widths.")

        freq11 = rf.Frequency.from_f(f[mask11], unit="Hz")
        freq22 = rf.Frequency.from_f(f[mask22], unit="Hz")

        # Use Z-parameters directly
        z11_data = data_aligned.z[mask11, 0, 0]
        z22_data = data_aligned.z[mask22, 1, 1]

        def _depth(z_complex):
            m = np.abs(z_complex)
            return max(float(np.max(m) - np.min(m)), eps)

        Z11_scale = _depth(z11_data)
        Z22_scale = _depth(z22_data)
        depth_floor = max(Z11_scale, Z22_scale)
        Z11_scale = max(Z11_scale, depth_floor)
        Z22_scale = max(Z22_scale, depth_floor)

        logger.debug("Z11 depth scale=%.4e, N11=%d", Z11_scale, mask11.sum())
        logger.debug("Z22 depth scale=%.4e, N22=%d", Z22_scale, mask22.sum())

        def residuals(x):
            Le, Ce = float(x[0]), float(x[1])
            lom11 = lc_resonator_network_2port(Le, Ce, Cc1, Cc2, freq11, Z0)
            lom22 = lc_resonator_network_2port(Le, Ce, Cc1, Cc2, freq22, Z0)

            r11 = lom11.z[:, 0, 0] - z11_data
            r22 = lom22.z[:, 1, 1] - z22_data

            r11_sc = np.concatenate([r11.real, r11.imag]) / Z11_scale
            r22_sc = np.concatenate([r22.real, r22.imag]) / Z22_scale

            return np.concatenate([r11_sc, r22_sc])

        return residuals
        # ------------------------------------------------------------------
    # fit()
    # ------------------------------------------------------------------

    def fit(
        self,
        freq: rf.Frequency,
        data_ntw: rf.Network,
        Cc1: float,
        Cc2: float,
        Ctog1: float,
        Ctog2: float,
        d: float,
        cpw_params: CPWParams,
        Z0: float = 50.0,
        **kwargs,
    ) -> None:
        """
        Two-stage optimisation to fit Leff and Ceff to data_ntw.

        Stage 1: coarse Nelder-Mead scan over ω0, constrained Ceff.
            Uses an explicitly constructed 1-port CPW network so that the
            Stage 1 data points come from a true single-port reflection
            measurement rather than a 2-port S11 view.
        Stage 2: windowed least_squares refinement on S11 + S22.

        Parameters
        ----------
        freq : rf.Frequency
            Frequency grid for the optimisation (should be centred on the
            resonance with enough bandwidth to capture the linewidth).
        data_ntw : rf.Network
            2-port reference network (CPW model or measured data).
            Used for Stage 2 S11 + S22 residuals.
        Cc1, Cc2 : float
            Coupling capacitances in Farads.
        Ctog1, Ctog2 : float
            Shunt-to-ground capacitances in Farads.
        d : float
            Resonator length in metres — used to build the 1-port CPW network.
        cpw_params : CPWParams
            CPW geometry — used to build the 1-port CPW network.
        Z0 : float
            Reference impedance in Ohms.
        """
        from simpleLOMs.networks.cpw import cpw_resonator_network

        cfg = self.config

        # --- Stage 1: explicit 1-port CPW network ---
        single_port_ntw = cpw_resonator_network(
            freq, d, Cc1, Cc2, Ctog1, Ctog2,
            cpw_params=cpw_params, Z0=Z0,
        )

        data_points = self._build_sparse_data_points(
            single_port_ntw,
            n_dense=cfg.n_dense,
            n_kappa=cfg.n_kappa,
        )
        w0_guess = resonance_from_res11(single_port_ntw) * 2 * np.pi

        Ceff_guess, Leff_guess, w0_best, scan_results = self._fit_Ceff_Leff(
            w0_guess=w0_guess,
            Cc1=Cc1,
            Cc2=Cc2,
            Z0=Z0,
            data_points=data_points,
            freq=freq,
        )
        self.scan_results = scan_results
        logger.debug("Stage 1: L_guess=%.4e H, C_guess=%.4e F", Leff_guess, Ceff_guess)

        # --- Stage 2: 2-port windowed least_squares ---
        kappa_s11 = fwhm_from_trace_db(data_ntw, m=0, n=0, kind="dip")
        kappa_s22 = fwhm_from_trace_db(data_ntw, m=1, n=1, kind="dip")
        f0_s11    = resonance(data_ntw, m=0, n=0)
        f0_s22    = resonance(data_ntw, m=1, n=1)
        f0_s21    = resonance_from_s_max(data_ntw, m=1, n=0)
        f0_s12    = resonance_from_s_max(data_ntw, m=0, n=1)

        residuals_fn = self._make_windowed_residuals(
            freq=freq,
            data_ntw=data_ntw,
            Cc1=Cc1,
            Cc2=Cc2,
            f0_s11=f0_s11,
            kappa_s11=kappa_s11,
            f0_s22=f0_s22,
            kappa_s22=kappa_s22,
            Z0=Z0,
            # f0_s12 = f0_s12,
            # f0_s21 = f0_s21,
        )

        res = least_squares(
            residuals_fn,
            x0=(Leff_guess, Ceff_guess),
            bounds=(
                [Leff_guess - 2e-10, Ceff_guess - 2e-13],
                [Leff_guess + 2e-10, Ceff_guess + 2e-13],
            ),
            method="trf",
            jac="3-point",
            diff_step=1e-4,
            x_scale="jac",
            xtol=1e-15,
            ftol=1e-15,
            gtol=1e-15,
            max_nfev=cfg.max_nfev,
            verbose=2 if cfg.verbose else 0,
        )

        self.L = float(res.x[0])
        self.C = float(res.x[1])
        logger.debug("Stage 2 result: L=%.4e H, C=%.4e F", self.L, self.C)

    # ------------------------------------------------------------------
    # get_network()
    # ------------------------------------------------------------------

    def get_network(
        self,
        freq: rf.Frequency,
        Cc1: float,
        Cc2: float,
        Z0: float = 50.0,
        with_grounds: bool = False,
        Ctog1: float = None,
        Ctog2: float = None,
        **kwargs,
    ) -> rf.Network:
        """
        Build the 2-port LC LOM network from the optimised L and C.

        Parameters
        ----------
        freq : rf.Frequency
        Cc1, Cc2 : float
            Coupling capacitances in Farads.
        Z0 : float
            Reference impedance in Ohms.
        with_grounds : bool
            If True, include Ctog1/Ctog2 shunt caps (requires them to
            be provided).
        Ctog1, Ctog2 : float or None
            Required when with_grounds=True.

        Returns
        -------
        rf.Network
        """
        self._require_fitted()

        if with_grounds:
            if Ctog1 is None or Ctog2 is None:
                raise ValueError("Ctog1 and Ctog2 must be provided when with_grounds=True.")
            return lc_resonator_network_with_grounds_2port(
                Leff=self.L, Ceff=self.C,
                Cc1=Cc1, Cc2=Cc2,
                Ctog1=Ctog1, Ctog2=Ctog2,
                freq=freq, Z0=Z0,
            )

        return lc_resonator_network_2port(
            Leff=self.L, Ceff=self.C,
            Cc1=Cc1, Cc2=Cc2,
            freq=freq, Z0=Z0,
        )

    # ------------------------------------------------------------------
    # get_params()
    # ------------------------------------------------------------------

    def get_params(self) -> dict:
        """
        Return fitted parameters and scan diagnostics.

        Returns
        -------
        dict
            Keys: "L" (H), "C" (F), "scan_results" (ndarray or None).
        """
        self._require_fitted()
        return {
            "L":            self.L,
            "C":            self.C,
            "scan_results": self.scan_results,
        }

    def __repr__(self) -> str:
        if self.is_fitted:
            return "OptimizedFit(L={:.4e} H, C={:.4e} F)".format(self.L, self.C)
        return "OptimizedFit(unfitted)"
