#!/usr/bin/env python3
"""Reproduce physical-bias scaling diagnostics from Dillavou et al. data.

This is a descriptive reanalysis of already published data, not a prospective
confirmation.  It reads only the small-network CSV files from Zenodo record
15692914 (release v1.0.1), records their hashes, and measures the low-period
power laws reported in the paper:

* nonzero combined-error plateau as the task-switching period decreases;
* squared cycle span proportional to approximately period squared;
* therefore a nonzero cycle-speed proxy span / period.

The script also audits the released two-dimensional bow-tie trajectories.
"""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
import re
from typing import Dict, List

import matplotlib

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


ZENODO_RECORD = "15692914"
ZENODO_DOI = "10.5281/zenodo.15692914"
RELEASE = "v1.0.1"
SOURCE_TREE = "maguzj-imperfect-learning-physical-systems-71b8d72"


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def log_slope(values: np.ndarray) -> float:
    if values.ndim != 2 or values.shape[1] != 2:
        raise ValueError("expected a two-column positive-valued table")
    if np.any(values <= 0):
        raise ValueError("log slope requires strictly positive values")
    return float(np.polyfit(np.log(values[:, 0]), np.log(values[:, 1]), 1)[0])


def load_table(path: Path) -> np.ndarray:
    values = np.loadtxt(path, delimiter=",")
    if values.ndim != 2 or values.shape[1] != 2:
        raise ValueError(f"unexpected table shape in {path}: {values.shape}")
    if not np.all(np.isfinite(values)):
        raise ValueError(f"nonfinite value in {path}")
    if np.any(np.diff(values[:, 0]) <= 0):
        raise ValueError(f"periods are not strictly increasing in {path}")
    return values


def experiment_tables(data_dir: Path, fit_points: int) -> Dict[str, dict]:
    records: Dict[str, dict] = {}
    for experiment in (1, 2, 3):
        mse_path = data_dir / f"measured-MSE-exp{experiment}.csv"
        span2_path = data_dir / f"measured-DG2-exp{experiment}.csv"
        mse = load_table(mse_path)
        span2 = load_table(span2_path)
        if not np.array_equal(mse[:, 0], span2[:, 0]):
            raise ValueError(f"period grids disagree for experiment {experiment}")
        if len(mse) < fit_points:
            raise ValueError("fit_points exceeds a source table")
        low_mse = mse[:fit_points]
        low_span2 = span2[:fit_points]
        speed = np.sqrt(low_span2[:, 1]) / low_span2[:, 0]
        records[str(experiment)] = {
            "period": mse[:, 0].tolist(),
            "combined_error": mse[:, 1].tolist(),
            "cycle_span_squared": span2[:, 1].tolist(),
            "low_period_points": int(fit_points),
            "low_period_combined_error_log_slope": log_slope(low_mse),
            "low_period_span_squared_log_slope": log_slope(low_span2),
            "low_period_span_log_slope": 0.5 * log_slope(low_span2),
            "low_period_speed_proxy": speed.tolist(),
            "low_period_speed_proxy_mean": float(np.mean(speed)),
            "low_period_speed_proxy_cv": float(np.std(speed) / np.mean(speed)),
            "source_files": {
                "combined_error": {
                    "path": str(mse_path),
                    "sha256": sha256(mse_path),
                },
                "cycle_span_squared": {
                    "path": str(span2_path),
                    "sha256": sha256(span2_path),
                },
            },
        }
    return records


def bowtie_records(bowtie_dir: Path) -> List[dict]:
    pattern = re.compile(r"bowtie-exp(?P<experiment>[123])-(?P<period>[0-9.]+)s\.csv")
    records = []
    for path in sorted(bowtie_dir.glob("bowtie-exp*-*s.csv")):
        match = pattern.fullmatch(path.name)
        if match is None:
            raise ValueError(f"unrecognized bow-tie filename: {path.name}")
        trajectory = np.loadtxt(path, delimiter=",")
        if trajectory.ndim != 2 or trajectory.shape[0] != 2:
            raise ValueError(f"unexpected bow-tie shape in {path}: {trajectory.shape}")
        pairwise = trajectory[:, :, None] - trajectory[:, None, :]
        diameter = float(np.sqrt(np.sum(pairwise * pairwise, axis=0)).max())
        midpoint = trajectory.shape[1] // 2
        half_cycle_span = float(np.linalg.norm(
            trajectory[:, 0] - trajectory[:, midpoint]))
        path_length = float(np.linalg.norm(
            np.diff(trajectory, axis=1), axis=0).sum())
        records.append({
            "experiment": int(match.group("experiment")),
            "period_seconds": float(match.group("period")),
            "samples": int(trajectory.shape[1]),
            "diameter": diameter,
            "half_cycle_span": half_cycle_span,
            "path_length": path_length,
            "trajectory": trajectory.tolist(),
            "source_file": {"path": str(path), "sha256": sha256(path)},
        })
    if len(records) != 10:
        raise ValueError(f"expected 10 bow-tie trajectories, found {len(records)}")
    return records


def representative_quadratic(period: np.ndarray, values: np.ndarray) -> np.ndarray:
    return values[0] * (period / period[0]) ** 2


def plot_report(report: dict, output: Path) -> None:
    colors = ["#4477AA", "#EE6677", "#228833"]
    fig, axes = plt.subplots(2, 2, figsize=(9.0, 7.0))

    for index, (experiment, record) in enumerate(report["experiments"].items()):
        period = np.asarray(record["period"])
        error = np.asarray(record["combined_error"])
        span2 = np.asarray(record["cycle_span_squared"])
        label = f"physical experiment {experiment}"
        axes[0, 0].loglog(period, error, "o-", color=colors[index], label=label)
        axes[0, 1].loglog(period, span2, "o-", color=colors[index], label=label)
        low_n = int(record["low_period_points"])
        axes[0, 1].loglog(
            period[:low_n], representative_quadratic(period[:low_n], span2[:low_n]),
            "--", color=colors[index], alpha=0.55)
        speed = np.sqrt(span2) / period
        axes[1, 0].semilogx(period, speed, "o-", color=colors[index], label=label)

    axes[0, 0].set_title("A  Error remains nonzero under rapid switching")
    axes[0, 0].set_xlabel("task-switching period (s)")
    axes[0, 0].set_ylabel("combined error")
    axes[0, 0].legend(frameon=False, fontsize=8)

    axes[0, 1].set_title("B  Cycle span follows the bias-drift prediction")
    axes[0, 1].set_xlabel("task-switching period (s)")
    axes[0, 1].set_ylabel("cycle span squared")

    axes[1, 0].set_title("C  Span per unit period does not vanish")
    axes[1, 0].set_xlabel("task-switching period (s)")
    axes[1, 0].set_ylabel("sqrt(span squared) / period")

    selected = [
        item for item in report["bowtie_trajectories"]
        if item["experiment"] == 1
    ]
    selected.sort(key=lambda item: item["period_seconds"])
    for index, item in enumerate(selected):
        trajectory = np.asarray(item["trajectory"])
        axes[1, 1].plot(
            trajectory[0], trajectory[1], "o-", markersize=2.1,
            linewidth=1.0, color=colors[index],
            label=f"period={item['period_seconds']:g} s")
    axes[1, 1].set_title("D  Measured bias-driven parameter cycles")
    axes[1, 1].set_xlabel("gate voltage +")
    axes[1, 1].set_ylabel("gate voltage -")
    axes[1, 1].legend(frameon=False, fontsize=8)

    for axis in axes.flat:
        axis.grid(alpha=0.18)
    fig.suptitle(
        "Published physical coupled-learning data (Dillavou et al.; Zenodo 15692914)",
        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(
        "--artifact-root", type=Path, required=True,
        help="Root of the extracted Zenodo source tree")
    parser.add_argument("--fit-points", type=int, default=6)
    parser.add_argument(
        "--json", type=Path,
        default=Path("results/physical_bias/p0_summary.json"))
    parser.add_argument(
        "--figure", type=Path,
        default=Path("results/figs/physical_bias_p0.png"))
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.fit_points < 3:
        raise ValueError("fit_points must be at least three")
    root = args.artifact_root.resolve()
    expected_name = SOURCE_TREE
    if root.name != expected_name:
        raise ValueError(f"expected source root {expected_name}, received {root.name}")
    small = root / "small network"
    records = experiment_tables(small / "MSE-DG2-Data", args.fit_points)
    bowties = bowtie_records(small / "experimental_bowties")
    mse_slopes = np.asarray([
        value["low_period_combined_error_log_slope"]
        for value in records.values()])
    span2_slopes = np.asarray([
        value["low_period_span_squared_log_slope"]
        for value in records.values()])
    report = {
        "analysis": "published_physical_bias_descriptive_reproduction",
        "confirmatory": False,
        "provenance": {
            "zenodo_record": ZENODO_RECORD,
            "doi": ZENODO_DOI,
            "release": RELEASE,
            "source_tree": SOURCE_TREE,
        },
        "fit_definition": {
            "points": int(args.fit_points),
            "selection": "lowest task-switching periods in each published table",
            "regression": "ordinary least squares on log(period), log(metric)",
        },
        "experiments": records,
        "bowtie_trajectories": bowties,
        "summary": {
            "combined_error_log_slope_mean": float(np.mean(mse_slopes)),
            "combined_error_log_slopes": mse_slopes.tolist(),
            "span_squared_log_slope_mean": float(np.mean(span2_slopes)),
            "span_squared_log_slopes": span2_slopes.tolist(),
            "span_log_slope_mean": float(0.5 * np.mean(span2_slopes)),
            "descriptive_plateau_check_abs_error_slope_below_0p4": bool(
                np.all(np.abs(mse_slopes) < 0.4)),
            "descriptive_linear_drift_check_span2_slope_1p5_to_2p3": bool(
                np.all((span2_slopes > 1.5) & (span2_slopes < 2.3))),
        },
    }
    args.json.parent.mkdir(parents=True, exist_ok=True)
    args.json.write_text(json.dumps(report, indent=2) + "\n")
    plot_report(report, args.figure)
    print(json.dumps(report["summary"], indent=2))
    print(f"wrote {args.json}")
    print(f"wrote {args.figure}")


if __name__ == "__main__":
    main()