#!/usr/bin/env python3
"""P2: test whether SDIL reduces the overclamping voltage budget.

The bias field comes from the held-out affine fits to the released Dillavou
drift traces.  The circuit and overclamping dynamics follow Appendix D/F of
arXiv:2505.22887v2.  This remains a measured-data surrogate rather than a new
hardware experiment.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import sys

import matplotlib

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

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from experiments.physical_bias_p1 import load_pair  # noqa: E402
from sdil.physical_coupled import simulate_alternating_tasks  # noqa: E402


OVERCLAMP_METHODS = ("overclamp", "overclamp_sdil")
STANDARD_METHODS = ("raw", "frozen_constant", "frozen_sdil", "oracle")
LABELS = {
    "overclamp": "overclamping",
    "overclamp_sdil": "SDIL + overclamping",
}
COLORS = {
    "overclamp": "#CC3311",
    "overclamp_sdil": "#0077BB",
}


def run_overclamp(
    pair: dict,
    method: str,
    period: float,
    cycles: int,
    eta: float,
    seed: int,
) -> dict:
    predictor = pair["sdil"] if method == "overclamp_sdil" else None
    result = simulate_alternating_tasks(
        pair["circuit"],
        pair["tasks"],
        pair["field"],
        method=method,
        period_seconds=period,
        cycles=cycles,
        initial_gates=np.asarray((4.0, 4.0)),
        bias_strength=1.0,
        predictor=predictor,
        seed=seed,
        overclamp_nudging=eta,
        overclamp_magnitude=pair["circuit"].high,
        overclamp_exact_target=True,
    )
    result["overclamp_eta"] = eta
    return result


def run_standard(
    pair: dict,
    method: str,
    period: float,
    cycles: int,
    seed: int,
) -> dict:
    predictor = None
    if method == "frozen_constant":
        predictor = pair["constant"]
    elif method == "frozen_sdil":
        predictor = pair["sdil"]
    return simulate_alternating_tasks(
        pair["circuit"],
        pair["tasks"],
        pair["field"],
        method=method,
        period_seconds=period,
        cycles=cycles,
        initial_gates=np.asarray((4.0, 4.0)),
        bias_strength=1.0,
        predictor=predictor,
        seed=seed,
    )


def select_minimum_eta(
    records: list[dict], method: str, periods: list[float], threshold: float
) -> list[dict]:
    selections = []
    for period in periods:
        candidates = sorted(
            (
                record for record in records
                if record["method"] == method
                and record["period_seconds"] == period
                and record["mean_combined_error"] <= threshold
            ),
            key=lambda record: record["overclamp_eta"],
        )
        if candidates:
            selected = candidates[0]
            selections.append({
                "period_seconds": period,
                "minimum_passing_eta": selected["overclamp_eta"],
                "combined_error": selected["mean_combined_error"],
                "clamp_displacement_l2_time_v2_s": selected[
                    "clamp_displacement_l2_time_v2_s"],
                "max_abs_clamp_displacement_v": selected[
                    "max_abs_clamp_displacement_v"],
            })
        else:
            selections.append({
                "period_seconds": period,
                "minimum_passing_eta": None,
                "combined_error": None,
                "clamp_displacement_l2_time_v2_s": None,
                "max_abs_clamp_displacement_v": None,
            })
    return selections


def summarize_pair(
    records: list[dict], periods: list[float], threshold: float
) -> dict:
    selections = {
        method: select_minimum_eta(records, method, periods, threshold)
        for method in OVERCLAMP_METHODS
    }
    paired_ratios = []
    eta_reductions = []
    for baseline, combined in zip(
        selections["overclamp"], selections["overclamp_sdil"]
    ):
        baseline_eta = baseline["minimum_passing_eta"]
        combined_eta = combined["minimum_passing_eta"]
        if baseline_eta is not None and combined_eta is not None:
            eta_reductions.append(baseline_eta / combined_eta)
            paired_ratios.append(
                baseline["clamp_displacement_l2_time_v2_s"]
                / combined["clamp_displacement_l2_time_v2_s"]
            )
    return {
        "precision_threshold_v2": threshold,
        "minimum_passing_eta": selections,
        "all_periods_passed": {
            method: all(
                item["minimum_passing_eta"] is not None
                for item in method_selections
            )
            for method, method_selections in selections.items()
        },
        "sdil_combination_never_requires_larger_eta": all(
            combined["minimum_passing_eta"] is not None
            and (
                baseline["minimum_passing_eta"] is None
                or combined["minimum_passing_eta"]
                <= baseline["minimum_passing_eta"]
            )
            for baseline, combined in zip(
                selections["overclamp"], selections["overclamp_sdil"]
            )
        ),
        "median_eta_reduction_at_passing_endpoint": (
            None if not eta_reductions else float(np.median(eta_reductions))
        ),
        "median_clamp_l2_exposure_reduction_at_passing_endpoint": (
            None if not paired_ratios else float(np.median(paired_ratios))
        ),
    }


def aggregate(records: list[dict], method: str, etas: list[float], key: str):
    values = []
    for eta in etas:
        selected = np.asarray([
            record[key] for record in records
            if record["method"] == method and record["overclamp_eta"] == eta
        ])
        values.append((
            float(np.median(selected)),
            float(np.min(selected)),
            float(np.max(selected)),
        ))
    return np.asarray(values)


def plot_report(report: dict, output: Path) -> None:
    etas = report["protocol"]["overclamp_eta_grid"]
    fig, axes = plt.subplots(2, 2, figsize=(9.2, 7.0), sharex="col")
    for column, name in enumerate(("experiment_1", "experiment_2")):
        records = report["pairs"][name]["overclamp_sweep"]
        for method in OVERCLAMP_METHODS:
            error = aggregate(records, method, etas, "mean_combined_error")
            exposure = aggregate(
                records, method, etas,
                "clamp_displacement_l2_time_v2_s")
            axes[0, column].loglog(
                etas, error[:, 0], "o-", color=COLORS[method],
                label=LABELS[method])
            axes[0, column].fill_between(
                etas, np.maximum(error[:, 1], 1e-32), error[:, 2],
                color=COLORS[method], alpha=0.12)
            axes[1, column].loglog(
                etas, exposure[:, 0], "o-", color=COLORS[method],
                label=LABELS[method])
        axes[0, column].axhline(
            report["protocol"]["precision_threshold_v2"],
            color="#666666", linestyle="--", linewidth=1.0)
        axes[0, column].set_title(
            f"{chr(ord('A') + column)}  {name.replace('_', ' ')}: error")
        axes[0, column].set_ylabel("combined task error")
        axes[1, column].set_title(
            f"{chr(ord('C') + column)}  clamp exposure")
        axes[1, column].set_xlabel("overclamping nudging strength η")
        axes[1, column].set_ylabel("Σ duration × displacement² (V²s)")
        for row in range(2):
            axes[row, column].grid(alpha=0.18)
    axes[0, 0].legend(frameon=False, fontsize=8)
    fig.suptitle(
        "SDIL reduces the clamping required under measured state-dependent bias",
        fontsize=11)
    fig.tight_layout()
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=180)
    plt.close(fig)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--state-dependence-json", type=Path,
        default=Path("results/physical_bias/p0_state_dependence.json"))
    parser.add_argument(
        "--output", type=Path,
        default=Path("results/physical_bias/p2_clamp_budget.json"))
    parser.add_argument(
        "--figure", type=Path,
        default=Path("results/figs/physical_bias_p2_clamp_budget.png"))
    parser.add_argument("--minimum-cycles", type=int, default=120)
    parser.add_argument("--total-nominal-time", type=float, default=6.0)
    parser.add_argument(
        "--periods", type=float, nargs="+",
        default=(0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2))
    parser.add_argument(
        "--etas", type=float, nargs="+",
        default=(0.005, 0.01, 0.025, 0.05, 0.1, 0.25))
    parser.add_argument("--precision-threshold-v2", type=float, default=1e-8)
    parser.add_argument("--seed", type=int, default=20260829)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    source = json.loads(args.state_dependence_json.read_text())
    periods = list(args.periods)
    etas = list(args.etas)
    report = {
        "analysis": "physical_measured_state_dependent_clamp_budget_p2",
        "confirmatory": False,
        "physical_hardware_demonstration": False,
        "autodiff_used": False,
        "source_analysis": str(args.state_dependence_json),
        "protocol": {
            "periods_seconds": periods,
            "minimum_cycles": args.minimum_cycles,
            "total_nominal_time_seconds": args.total_nominal_time,
            "overclamp_eta_grid": etas,
            "overclamp_target_magnitude_v": 0.4351,
            "overclamp_equation": "Appendix F Eq. F5 with th proportional to error",
            "precision_threshold_v2": args.precision_threshold_v2,
            "precision_threshold_interpretation": "0.1 mV output RMSE",
            "neutral_observation_protocol": (
                "same upfront released-trace observations for constant and affine predictors"
            ),
        },
        "pairs": {},
    }
    for pair_index, name in enumerate(("experiment_1", "experiment_2")):
        pair = load_pair(source, name, 1.0)
        overclamp_records = []
        standard_records = []
        for period_index, period in enumerate(periods):
            cycles = max(
                args.minimum_cycles,
                int(np.ceil(args.total_nominal_time / period)),
            )
            for method_index, method in enumerate(STANDARD_METHODS):
                standard_records.append(run_standard(
                    pair,
                    method,
                    period,
                    cycles,
                    args.seed + 10000 * pair_index
                    + 100 * period_index + method_index,
                ))
            for eta_index, eta in enumerate(etas):
                for method_index, method in enumerate(OVERCLAMP_METHODS):
                    overclamp_records.append(run_overclamp(
                        pair,
                        method,
                        period,
                        cycles,
                        eta,
                        args.seed + 10000 * pair_index
                        + 100 * period_index + 10 * eta_index + method_index,
                    ))
        report["pairs"][name] = {
            "calibration": pair["calibration"],
            "standard_clamping": standard_records,
            "overclamp_sweep": overclamp_records,
            "summary": summarize_pair(
                overclamp_records,
                periods,
                args.precision_threshold_v2,
            ),
        }
    report["summary"] = {
        "combination_never_requires_larger_eta_both_pairs": all(
            pair["summary"]["sdil_combination_never_requires_larger_eta"]
            for pair in report["pairs"].values()
        ),
        "median_eta_reduction_by_pair": {
            name: pair["summary"][
                "median_eta_reduction_at_passing_endpoint"]
            for name, pair in report["pairs"].items()
        },
        "median_clamp_l2_exposure_reduction_by_pair": {
            name: pair["summary"][
                "median_clamp_l2_exposure_reduction_at_passing_endpoint"]
            for name, pair in report["pairs"].items()
        },
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n")
    plot_report(report, args.figure)
    print(json.dumps(report["summary"], indent=2))
    print(f"wrote {args.output}")
    print(f"wrote {args.figure}")


if __name__ == "__main__":
    main()