"""
system.py
---------
Top-level orchestration: runs all three fitting methods and compares them
against the CPW reference network.

"""

import logging
import itertools
import json

import numpy as np
import skrf as rf

from simpleLOMs.params import CPWParams
from simpleLOMs.analysis import (
    resonance,
    resonance_from_s_max,
    resonances_from_s_max,
    resonances_from_s,
    fwhm_from_trace_db,
    fwhm_from_res11,
)
from simpleLOMs.networks.cpw import (
    cpw_resonator_network,
    cpw_resonator_network_2port,
)
from simpleLOMs.models.foster_fit import FosterFit
from simpleLOMs.models.optimized_fit import OptimizedFit, OptimizationConfig
from simpleLOMs.models.analytical_fit import AnalyticalFit
from simpleLOMs.plotting import plot_re_im

logger = logging.getLogger(__name__)

_serialisable_keys = [
    "_sweep_point",
    "CPW f0 S11 (GHz)", "CPW f0 S22 (GHz)",
    "CPW kappa S11 (MHz)", "CPW kappa S22 (MHz)",
    "Foster f0 S11 (GHz)", "Foster f0 S22 (GHz)",
    "Foster kappa S11 (MHz)", "Foster kappa S22 (MHz)",
    "Foster f0 error S11 (%)", "Foster f0 error S22 (%)",
    "Foster kappa error S11 (%)", "Foster kappa error S22 (%)",
    "Optimized f0 S11 (GHz)", "Optimized f0 S22 (GHz)",
    "Optimized kappa S11 (MHz)", "Optimized kappa S22 (MHz)",
    "Optimized f0 error S11 (%)", "Optimized f0 error S22 (%)",
    "Optimized kappa error S11 (%)", "Optimized kappa error S22 (%)",
    "Analytical f0 S11 (GHz)", "Analytical f0 S22 (GHz)",
    "Analytical kappa S11 (MHz)", "Analytical kappa S22 (MHz)",
    "Analytical f0 error S11 (%)", "Analytical f0 error S22 (%)",
    "Analytical kappa error S11 (%)", "Analytical kappa error S22 (%)",
]


def _make_serialisable(r):
    if r is None:
        return None
    out = {}
    for k in _serialisable_keys:
        v = r.get(k)
        if isinstance(v, np.ndarray):
            v = v.tolist()
        out[k] = v
    return out


def analyze_system(
    freq: rf.Frequency,
    d: float,
    Cc1: float,
    Cc2: float,
    Ctog1: float,
    Ctog2: float,
    cpw_params: CPWParams = None,
    Z0: float = 50.0,
    analytical_f_r: float = None,
    analytical_Z0: float = None,
    opt_config: OptimizationConfig = None,
    verbose: bool = False,
) -> dict:
    """
    Compare FosterFit, OptimizedFit, and AnalyticalFit against a CPW reference.

    Parameters
    ----------
    freq : rf.Frequency
        Coarse frequency sweep for initial network construction and Foster fit.
    d : float
        Resonator length in metres.
    Cc1, Cc2 : float
        Coupling capacitances in Farads.
    Ctog1, Ctog2 : float
        Shunt-to-ground capacitances in Farads.
    cpw_params : CPWParams, optional
        CPW geometry. Uses default CPWParams() if not provided.
    Z0 : float
        Reference impedance in Ohms (default 50 Ω).
    analytical_f_r : float, optional
        Resonance frequency (Hz) to use for AnalyticalFit.
        If None, the CPW S11 resonance is used automatically.
    analytical_Z0 : float, optional
        Characteristic impedance for AnalyticalFit formula.
        If None, Z0 is used.
    opt_config : OptimizationConfig, optional
        Hyperparameters for OptimizedFit. Uses defaults if not provided.
    verbose : bool
        If True, passes verbose=True to least_squares in OptimizedFit.

    Returns
    -------
    dict
        Flat dictionary with resonance frequencies, linewidths, and percentage
        errors for all four methods (CPW, Foster, Optimized, Analytical).

    Notes
    -----
    The three fitted models and the CPW network are stored in the output dict
    under keys "cpw_network", "foster_model", "optimized_model",
    "analytical_model" so callers can do further analysis or plotting.
    """
    if cpw_params is None:
        cpw_params = CPWParams()
    if opt_config is None:
        opt_config = OptimizationConfig(verbose=verbose, n_widths=4)

    # ----------------------------------------------------------------
    # 1) Build single-port CPW for initial LOM guess (OptimizedFit)
    # ----------------------------------------------------------------
    cpw_single = cpw_resonator_network(
        freq, d, Cc1, Cc2, Ctog1, Ctog2,
        cpw_params=cpw_params, Z0=Z0,
    )

    # ----------------------------------------------------------------
    # 2) Build 2-port CPW; refine frequency window around resonance
    # ----------------------------------------------------------------
    f0_coarse    = resonance(cpw_resonator_network_2port(
        freq, d, Cc1, Cc2, Ctog1, Ctog2, cpw_params=cpw_params, Z0=Z0
    ), m=0, n=0)
    kappa_coarse = fwhm_from_trace_db(cpw_resonator_network_2port(
        freq, d, Cc1, Cc2, Ctog1, Ctog2, cpw_params=cpw_params, Z0=Z0
    ), m=0, n=0, kind="dip")

    if Cc1 == Cc2:
        refined_freq = rf.Frequency(f0_coarse - 0.5e9, f0_coarse + 0.5e9, 500_001, unit="Hz")
    else:
        refined_freq = rf.Frequency(
            f0_coarse - 5 * kappa_coarse,
            f0_coarse + 5 * kappa_coarse,
            200_001,
            unit="Hz",
        )

    cpw_network = cpw_resonator_network_2port(
        refined_freq, d, Cc1, Cc2, Ctog1, Ctog2,
        cpw_params=cpw_params, Z0=Z0,
    )

    plot_re_im(cpw_network)

    # CPW reference resonances and linewidths
    f0_cpw_s11 = resonance(cpw_network, m=0, n=0) / 1e9
    f0_cpw_s22 = resonance(cpw_network, m=1, n=1) / 1e9
    k_cpw_s11  = fwhm_from_trace_db(cpw_network, m=0, n=0, kind="dip")
    k_cpw_s22  = fwhm_from_trace_db(cpw_network, m=1, n=1, kind="dip")
    f0_cpw_s12 = resonance_from_s_max(cpw_network, m=0, n=1)
    logger.debug("CPW S11: f0=%.4f GHz, κ=%.4f MHz", f0_cpw_s11, k_cpw_s11 / 1e6)

    # ----------------------------------------------------------------
    # 3) FosterFit
    # ----------------------------------------------------------------
    foster_model = FosterFit(cpw_params=cpw_params)
    foster_model.fit(freq, d=d)

    foster_network = foster_model.get_network(
        refined_freq, Cc1=Cc1, Cc2=Cc2, Ctog1=Ctog1, Ctog2=Ctog2, Z0=Z0,
    )

    # ----------------------------------------------------------------
    # 4) OptimizedFit
    # ----------------------------------------------------------------
    optimized_model = OptimizedFit(config=opt_config)
    optimized_model.fit(
        refined_freq,
        data_ntw=cpw_network,
        Cc1=Cc1,
        Cc2=Cc2,
        Ctog1=Ctog1,
        Ctog2=Ctog2,
        d=d,
        cpw_params=cpw_params,
        Z0=Z0,
    )
    optimized_network = optimized_model.get_network(
        refined_freq, Cc1=Cc1, Cc2=Cc2, Z0=Z0,
    )

    # ----------------------------------------------------------------
    # 5) AnalyticalFit
    # ----------------------------------------------------------------
    f_r_for_analytical = analytical_f_r if analytical_f_r is not None else (f0_cpw_s11 * 1e9)
    Z0_for_analytical  = analytical_Z0  if analytical_Z0  is not None else Z0

    analytical_model = AnalyticalFit()
    analytical_model.fit(refined_freq, f_r=8.58e9, Z0=Z0_for_analytical)

    analytical_network = analytical_model.get_network(
        refined_freq, Cc1=Cc1, Cc2=Cc2, Ctog1=Ctog1, Ctog2=Ctog2, Z0=Z0,
    )

    # ----------------------------------------------------------------
    # 6) S11/S22 resonance and linewidth comparison
    # ----------------------------------------------------------------
    def _metrics(net, label):
        return {
            "f0_s11": resonance(net, 0, 0) / 1e9,
            "f0_s22": resonance(net, 1, 1) / 1e9,
            "k_s11":  fwhm_from_trace_db(net, 0, 0, kind="dip"),
            "k_s22":  fwhm_from_trace_db(net, 1, 1, kind="dip"),
            "f0_s21": resonance_from_s_max(net, 1, 0),
        }

    m_foster     = _metrics(foster_network,     "Foster")
    m_opt        = _metrics(optimized_network,  "Optimized")
    m_analytical = _metrics(analytical_network, "Analytical")

    def _pct_err(ref, val):
        return abs(ref - val) / abs(ref) * 100 if ref != 0 else float("nan")

    # ----------------------------------------------------------------
    # 7) Assemble output dict
    # ----------------------------------------------------------------
    output = {
        # --- models (for further analysis / plotting) ---
        "cpw_network":        cpw_network,
        "foster_model":       foster_model,
        "optimized_model":    optimized_model,
        "analytical_model":   analytical_model,
        "foster_network":     foster_network,
        "optimized_network":  optimized_network,
        "analytical_network": analytical_network,

        # --- circuit parameters ---
        "Z0 (Ohms)":  Z0,
        "Cc1 (F)":    Cc1,
        "Cc2 (F)":    Cc2,
        "Ctog1 (F)":  Ctog1,
        "Ctog2 (F)":  Ctog2,

        # --- fitted L/C ---
        "Optimized L (H)":   optimized_model.L,
        "Optimized C (F)":   optimized_model.C,
        "Foster L (H)":      foster_model.L,
        "Foster C (F)":      foster_model.C,
        "Analytical L (H)":  analytical_model.L,
        "Analytical C (F)":  analytical_model.C,

        # --- S11 resonance / linewidth ---
        "CPW f0 S11 (GHz)":              f0_cpw_s11,
        "Optimized f0 S11 (GHz)":        m_opt["f0_s11"],
        "Foster f0 S11 (GHz)":           m_foster["f0_s11"],
        "Analytical f0 S11 (GHz)":       m_analytical["f0_s11"],
        "Optimized f0 error S11 (%)":    _pct_err(f0_cpw_s11, m_opt["f0_s11"]),
        "Foster f0 error S11 (%)":       _pct_err(f0_cpw_s11, m_foster["f0_s11"]),
        "Analytical f0 error S11 (%)":   _pct_err(f0_cpw_s11, m_analytical["f0_s11"]),

        "CPW kappa S11 (MHz)":           k_cpw_s11 / 1e6,
        "Optimized kappa S11 (MHz)":     m_opt["k_s11"] / 1e6,
        "Foster kappa S11 (MHz)":        m_foster["k_s11"] / 1e6,
        "Analytical kappa S11 (MHz)":    m_analytical["k_s11"] / 1e6,
        "Optimized kappa error S11 (%)": _pct_err(k_cpw_s11, m_opt["k_s11"]),
        "Foster kappa error S11 (%)":    _pct_err(k_cpw_s11, m_foster["k_s11"]),
        "Analytical kappa error S11 (%)":_pct_err(k_cpw_s11, m_analytical["k_s11"]),

        # --- S22 resonance / linewidth ---
        "CPW f0 S22 (GHz)":              f0_cpw_s22,
        "Optimized f0 S22 (GHz)":        m_opt["f0_s22"],
        "Foster f0 S22 (GHz)":           m_foster["f0_s22"],
        "Analytical f0 S22 (GHz)":       m_analytical["f0_s22"],
        "Optimized f0 error S22 (%)":    _pct_err(f0_cpw_s22, m_opt["f0_s22"]),
        "Foster f0 error S22 (%)":       _pct_err(f0_cpw_s22, m_foster["f0_s22"]),
        "Analytical f0 error S22 (%)":   _pct_err(f0_cpw_s22, m_analytical["f0_s22"]),

        "CPW kappa S22 (MHz)":           k_cpw_s22 / 1e6,
        "Optimized kappa S22 (MHz)":     m_opt["k_s22"] / 1e6,
        "Foster kappa S22 (MHz)":        m_foster["k_s22"] / 1e6,
        "Analytical kappa S22 (MHz)":    m_analytical["k_s22"] / 1e6,
        "Optimized kappa error S22 (%)": _pct_err(k_cpw_s22, m_opt["k_s22"]),
        "Foster kappa error S22 (%)":    _pct_err(k_cpw_s22, m_foster["k_s22"]),
        "Analytical kappa error S22 (%)":_pct_err(k_cpw_s22, m_analytical["k_s22"]),
    }

    return output


import os


def run_accuracy_sweep(
    sweep_params: dict,
    fixed_params: dict,
    save_path: str = None,
    skip_errors: bool = True,
    verbose: bool = True,
) -> list:
    """
    Run analyze_system() over a grid of parameter combinations and collect
    results.

    Parameters
    ----------
    sweep_params : dict[str, array-like]
        Parameters to sweep, as {param_name: values_array}.
        Any number of parameters can be swept — the function builds a full
        Cartesian grid.  Parameter names must match analyze_system() kwargs.

        Examples::

            {"Cc1": np.linspace(1e-15, 20e-15, 10),
             "Cc2": np.linspace(1e-15, 20e-15, 10)}

            {"d": np.linspace(5e-3, 10e-3, 8),
             "Ctog1": np.linspace(1e-15, 50e-15, 8)}

    fixed_params : dict
        All remaining analyze_system() arguments held constant across the
        sweep.  Must include at minimum: freq, d (or sweep it), Cc1, Cc2,
        Ctog1, Ctog2, cpw_params, Z0.

    save_path : str, optional
        If provided, saves results metadata to this path after the sweep.
        - Ends in ".csv":  saves a flat summary DataFrame (one row per grid
          point, columns = sweep params + scalar results like f0/kappa/errors).
        - Ends in ".json": saves the full nested results including shift
          arrays, as a JSON file.
        Does not save the rf.Network objects, since they are not serializable. 

    skip_errors : bool
        If False, the first exception stops the sweep.

    verbose : bool
        If True (default), shows a tqdm progress bar with the current
        parameter combination displayed.

    Returns
    -------
    list
        N-dimensional nested list of analyze_system() result dicts (or None
        for failed points if skip_errors=True).  The nesting order matches
        the order of keys in sweep_params.

        For a 2-parameter sweep over params A and B:
            results_grid[i][j]  →  result at (A[i], B[j])

        For a 3-parameter sweep over A, B, C:
            results_grid[i][j][k]  →  result at (A[i], B[j], C[k])

    Examples
    --------
    Basic 2D sweep::

        results_grid = run_accuracy_sweep(
            sweep_params={
                "Cc1": np.linspace(1e-15, 20e-15, 10),
                "Cc2": np.linspace(1e-15, 20e-15, 10),
            },
            fixed_params=dict(
                freq=freq, d=7e-3,
                Ctog1=1e-14, Ctog2=1e-14,
                cpw_params=cpw_params, Z0=50.0,
                analytical_Z0=45.926,
            ),
            save_path="results/cc_sweep.json",
        )

    Then plot::

        plot_error_heatmap_trio(
            results_grid,
            param1_values=np.linspace(1e-15, 20e-15, 10),
            param2_values=np.linspace(1e-15, 20e-15, 10),
            param1_label="Cc1", param2_label="Cc2",
            metric="kappa_s11",
            param1_scale=1e15, param2_scale=1e15,
            param1_unit="fF",  param2_unit="fF",
        )
    """

    try:
        from tqdm import tqdm
        _tqdm_available = True
    except ImportError:
        _tqdm_available = False

    def _dummy_tqdm(iterable, **kwargs):
        desc = kwargs.get("desc", "Running")
        print(f"{desc} ({len(list(iterable))} points)...")
        return iter(iterable)

    param_names  = list(sweep_params.keys())
    param_values = [np.asarray(v) for v in sweep_params.values()]
    grid_points  = list(itertools.product(*param_values))
    n_total      = len(grid_points)
    n_failed     = 0

    if verbose:
        print(f"Sweep grid: {' × '.join(f'{k}({len(v)})' for k, v in sweep_params.items())} "
              f"= {n_total} points")

    # ── Run the sweep ─────────────────────────────────────────────────────────
    flat_results = []

    _tqdm_func = tqdm if _tqdm_available else _dummy_tqdm

    pbar = _tqdm_func(grid_points, desc="Sweeping", unit="pt", disable=not verbose)
    _json_file = None
    if save_path and save_path.endswith(".json"):
        os.makedirs(os.path.dirname(save_path), exist_ok=True) if os.path.dirname(save_path) else None
        _json_file = open(save_path, "w")
        _json_file.write("[\n")
        _first_entry = [True]  # use list so closure can mutate it

    for point in pbar:
        combo = dict(zip(param_names, point))

        if verbose and _tqdm_available:
            combo_str = "  ".join(f"{k}={v:.3g}" for k, v in combo.items())
            pbar.set_postfix_str(combo_str, refresh=False)

        call_kwargs = {**fixed_params, **combo}

        try:
            result = analyze_system(**call_kwargs)
            result["_sweep_point"] = combo
            flat_results.append(result)
            # Write incrementally if JSON
            if _json_file is not None:
                entry = _make_serialisable(result)
                prefix = "" if _first_entry[0] else ","
                _json_file.write(prefix + json.dumps(entry, indent=2, default=str) + "\n")
                _json_file.flush()
                _first_entry[0] = False
        except Exception as exc:
            if skip_errors:
                n_failed += 1
                logger.warning("analyze_system failed at %s: %s", combo, exc)
                flat_results.append(None)
            else:
                raise

    if verbose and n_failed > 0:
        print(f"  ⚠  {n_failed}/{n_total} points failed and were skipped.")

    # ── Reshape flat list into N-dimensional nested list ──────────────────────
    def _reshape(flat, shapes):
        if len(shapes) == 1:
            return flat
        stride = int(np.prod(shapes[1:]))
        return [_reshape(flat[i*stride:(i+1)*stride], shapes[1:])
                for i in range(shapes[0])]

    shapes       = [len(v) for v in param_values]
    results_grid = _reshape(flat_results, shapes)

    # ── Save ──────────────────────────────────────────────────────────────────
    if save_path:
        serialisable_flat = [_make_serialisable(r) for r in flat_results]

        if save_path.endswith(".json"):
            payload = {
                "sweep_params":  {k: v.tolist() for k, v in zip(param_names, param_values)},
                "fixed_params":  {k: v for k, v in fixed_params.items()
                                  if isinstance(v, (int, float, str, bool))},
                "results":       serialisable_flat,
            }
            with open(save_path, "w") as f:
                json.dump(payload, f, indent=2, default=str)
            if verbose:
                print(f"Saved JSON: {save_path}")

        elif save_path.endswith(".csv"):
            rows = []
            for r in serialisable_flat:
                if r is None:
                    rows.append({**{k: np.nan for k in param_names}})
                    continue
                row = {**r["_sweep_point"]}
                for k, v in r.items():
                    if k == "_sweep_point":
                        continue
                    if isinstance(v, list):
                        for idx, val in enumerate(v):
                            row[f"{k}[{idx}]"] = val
                    elif isinstance(v, (int, float, str, type(None))):
                        row[k] = v
                rows.append(row)
            import pandas as pd
            pd.DataFrame(rows).to_csv(save_path, index=False)
            if verbose:
                print(f"Saved CSV: {save_path}")

    return results_grid