#!/usr/bin/env python3
"""Test whether released physical drift is state dependent.

Dillavou et al. approximate the two-edge physical network with a constant
bias vector.  Their appendix also notes that the experimental drift rate is
not constant.  This script tests that statement directly from the four
released DriftTests traces.

Once a trace is close to a one-task solution line, the clean learning force is
normal to that line.  Its measured tangential velocity therefore estimates
the tangential projection of the physical bias.  We compare:

    constant:       B(G) = b
    local affine:   B_i(G_i) = b_i + k_i (G_i - G_ref_i)

The second model uses exactly the per-edge local information available to an
affine SDIL predictor.  It is a measured-data identifiability diagnostic, not
yet evidence that SDIL improves task learning.
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass
import hashlib
import json
from pathlib import Path
from typing import Dict, List, Sequence

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.io import loadmat


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


@dataclass(frozen=True)
class TraceSpec:
    filename: str
    task: str
    solution_slope: float
    solution_intercept: float


@dataclass(frozen=True)
class PairSpec:
    name: str
    reference_gate: tuple[float, float]
    published_bias_at_reference: tuple[float, float]
    traces: tuple[TraceSpec, TraceSpec]


PAIR_SPECS = (
    PairSpec(
        name="experiment_1",
        reference_gate=(3.0, 3.5),
        published_bias_at_reference=(2.9, 4.7),
        traces=(
            TraceSpec("exp1DriftTest_1.mat", "beta", 0.43914, 1.6061),
            TraceSpec("exp1DriftTest_2.mat", "alpha", 2.7257, -3.3003),
        ),
    ),
    PairSpec(
        name="experiment_2",
        reference_gate=(3.1, 4.3),
        published_bias_at_reference=(-0.3, 2.0),
        traces=(
            TraceSpec("exp2DriftTest_1.mat", "beta", 0.74821, 1.9224),
            TraceSpec("exp2DriftTest_2.mat", "alpha", 2.5678, -2.901),
        ),
    ),
)


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 rmse(observed: np.ndarray, predicted: np.ndarray) -> float:
    return float(np.sqrt(np.mean((observed - predicted) ** 2)))


def r_squared(observed: np.ndarray, predicted: np.ndarray) -> float:
    denominator = float(np.sum((observed - np.mean(observed)) ** 2))
    if denominator == 0.0:
        return float("nan")
    return float(1.0 - np.sum((observed - predicted) ** 2) / denominator)


def load_trace(
    drift_dir: Path,
    spec: TraceSpec,
    reference_gate: Sequence[float],
    *,
    minimum_time: float,
    terminal_headroom: float,
    maximum_line_distance: float,
    trim_selected_ends: int,
) -> dict:
    path = drift_dir / spec.filename
    raw = loadmat(path, squeeze_me=True)
    time = np.asarray(raw["time"], dtype=float)
    gate_minus = np.asarray(raw["GMINUS"], dtype=float)
    gate_plus = np.asarray(raw["GPLUS"], dtype=float)
    if time.ndim != 1 or len(time) != 300:
        raise ValueError(f"unexpected time array in {path}: {time.shape}")
    if not (time.shape == gate_minus.shape == gate_plus.shape):
        raise ValueError(f"array shape disagreement in {path}")
    if not np.all(np.isfinite(np.column_stack((time, gate_minus, gate_plus)))):
        raise ValueError(f"nonfinite source value in {path}")

    # The first two entries straddle the experiment initialization and a
    # roughly 0.9 s gap.  The remaining samples have a stable ~30 ms cadence.
    time = time[2:]
    gate_minus = gate_minus[2:]
    gate_plus = gate_plus[2:]
    sample_period = float(np.median(np.diff(time)))
    velocity_minus = np.gradient(gate_minus, time)
    velocity_plus = np.gradient(gate_plus, time)

    slope = spec.solution_slope
    intercept = spec.solution_intercept
    tangent = np.asarray((1.0, slope), dtype=float)
    tangent /= np.linalg.norm(tangent)
    signed_distance = (
        gate_plus - slope * gate_minus - intercept
    ) / np.sqrt(1.0 + slope * slope)
    # Every released trace terminates on a clear gate-voltage plateau.  Exclude
    # the final headroom before either observed plateau so a hardware rail is
    # not mislabeled as a state-dependent bias.  Trimming the selected ends
    # then ensures that centered differences use only pre-plateau samples.
    selected = (
        (time >= minimum_time)
        & (gate_minus < np.max(gate_minus) - terminal_headroom)
        & (gate_plus < np.max(gate_plus) - terminal_headroom)
        & (np.abs(signed_distance) < maximum_line_distance)
    )
    indices = np.flatnonzero(selected)
    if len(indices) <= 2 * trim_selected_ends:
        raise ValueError(f"too few usable solution-line samples in {path}")
    if np.any(np.diff(indices) != 1):
        raise ValueError(f"selected solution-line window is not contiguous in {path}")
    indices = indices[trim_selected_ends:-trim_selected_ends]

    gates = np.column_stack((gate_minus[indices], gate_plus[indices]))
    velocities = np.column_stack(
        (velocity_minus[indices], velocity_plus[indices]))
    reference = np.asarray(reference_gate, dtype=float)
    design = np.column_stack((
        np.full(len(indices), tangent[0]),
        np.full(len(indices), tangent[1]),
        tangent[0] * (gates[:, 0] - reference[0]),
        tangent[1] * (gates[:, 1] - reference[1]),
    ))
    projected_velocity = velocities @ tangent
    path_coordinate = gates @ tangent
    if len(indices) < 6:
        raise ValueError(f"need at least six retained samples in {path}")

    return {
        "spec": spec,
        "source_file": {"path": str(path), "sha256": sha256(path)},
        "sample_period_seconds": sample_period,
        "raw_gate_minus": gate_minus,
        "raw_gate_plus": gate_plus,
        "time": time,
        "selected_indices": indices,
        "selected_time": time[indices],
        "selected_gates": gates,
        "selected_line_distance": signed_distance[indices],
        "tangent": tangent,
        "path_coordinate": path_coordinate,
        "projected_velocity": projected_velocity,
        "design": design,
    }


def fit_columns(design: np.ndarray, target: np.ndarray, columns: List[int]) -> dict:
    selected = design[:, columns]
    parameters, _, _, _ = np.linalg.lstsq(selected, target, rcond=None)
    predicted = selected @ parameters
    return {
        "columns": columns,
        "parameters": parameters,
        "predicted": predicted,
        "rank": int(np.linalg.matrix_rank(selected)),
        "condition_number": float(np.linalg.cond(selected)),
        "rmse": rmse(target, predicted),
        "r_squared": r_squared(target, predicted),
    }


def forward_block_test(traces: Sequence[dict], columns: List[int]) -> dict:
    train_design = []
    train_target = []
    test_design = []
    test_target = []
    split_records = []
    for trace in traces:
        count = len(trace["projected_velocity"])
        train_count = max(2, int(np.floor(0.6 * count)))
        if count - train_count < 2:
            raise ValueError("forward block test needs at least two held-out samples")
        train_design.append(trace["design"][:train_count, columns])
        train_target.append(trace["projected_velocity"][:train_count])
        test_design.append(trace["design"][train_count:, columns])
        test_target.append(trace["projected_velocity"][train_count:])
        split_records.append({
            "source": trace["spec"].filename,
            "train_samples": train_count,
            "test_samples": count - train_count,
        })
    train_design_array = np.vstack(train_design)
    train_target_array = np.concatenate(train_target)
    test_design_array = np.vstack(test_design)
    test_target_array = np.concatenate(test_target)
    parameters, _, _, _ = np.linalg.lstsq(
        train_design_array, train_target_array, rcond=None)
    train_prediction = train_design_array @ parameters
    test_prediction = test_design_array @ parameters
    return {
        "split": "first 60% of each retained trace trains; final 40% tests",
        "trace_splits": split_records,
        "parameters": parameters.tolist(),
        "train_rank": int(np.linalg.matrix_rank(train_design_array)),
        "train_condition_number": float(np.linalg.cond(train_design_array)),
        "train_rmse": rmse(train_target_array, train_prediction),
        "test_rmse": rmse(test_target_array, test_prediction),
        "test_r_squared": r_squared(test_target_array, test_prediction),
    }


def serializable_trace(trace: dict) -> dict:
    spec = trace["spec"]
    return {
        "source_file": trace["source_file"],
        "task": spec.task,
        "solution_line": {
            "gate_plus_slope": spec.solution_slope,
            "gate_plus_intercept": spec.solution_intercept,
        },
        "sample_period_seconds": trace["sample_period_seconds"],
        "retained_samples": int(len(trace["projected_velocity"])),
        "retained_time_seconds": trace["selected_time"].tolist(),
        "retained_gate_minus": trace["selected_gates"][:, 0].tolist(),
        "retained_gate_plus": trace["selected_gates"][:, 1].tolist(),
        "retained_line_distance": trace["selected_line_distance"].tolist(),
        "path_coordinate": trace["path_coordinate"].tolist(),
        "projected_drift_velocity_v_per_s": trace["projected_velocity"].tolist(),
    }


def analyze_pair(drift_dir: Path, spec: PairSpec, settings: dict) -> dict:
    traces = [
        load_trace(
            drift_dir,
            trace_spec,
            spec.reference_gate,
            **settings,
        )
        for trace_spec in spec.traces
    ]
    design = np.vstack([trace["design"] for trace in traces])
    target = np.concatenate([trace["projected_velocity"] for trace in traces])
    constant = fit_columns(design, target, [0, 1])
    local_affine = fit_columns(design, target, [0, 1, 2, 3])
    constant_forward = forward_block_test(traces, [0, 1])
    affine_forward = forward_block_test(traces, [0, 1, 2, 3])
    affine_parameters = local_affine["parameters"]
    published = np.asarray(spec.published_bias_at_reference)
    intercept = affine_parameters[:2]
    return {
        "reference_gate": list(spec.reference_gate),
        "published_constant_bias_at_reference_v_per_s": published.tolist(),
        "traces": [serializable_trace(trace) for trace in traces],
        "constant_model": {
            "definition": "B(G) = b",
            "bias_v_per_s": constant["parameters"].tolist(),
            "rank": constant["rank"],
            "condition_number": constant["condition_number"],
            "in_sample_rmse_v_per_s": constant["rmse"],
            "in_sample_r_squared": constant["r_squared"],
            "forward_block": constant_forward,
        },
        "local_affine_model": {
            "definition": "B_i(G_i) = b_i + k_i (G_i - G_ref_i)",
            "bias_at_reference_v_per_s": intercept.tolist(),
            "local_slopes_per_s": affine_parameters[2:].tolist(),
            "rank": local_affine["rank"],
            "condition_number": local_affine["condition_number"],
            "in_sample_rmse_v_per_s": local_affine["rmse"],
            "in_sample_r_squared": local_affine["r_squared"],
            "forward_block": affine_forward,
            "reference_bias_error_v_per_s": (intercept - published).tolist(),
        },
        "diagnostic": {
            "in_sample_rmse_ratio_affine_over_constant": (
                local_affine["rmse"] / constant["rmse"]),
            "forward_test_rmse_ratio_affine_over_constant": (
                affine_forward["test_rmse"] / constant_forward["test_rmse"]),
            "affine_improves_forward_test": bool(
                affine_forward["test_rmse"] < constant_forward["test_rmse"]),
        },
        "_plot": {
            "traces": traces,
            "constant_parameters": constant["parameters"],
            "affine_parameters": affine_parameters,
        },
    }


def headroom_sensitivity(drift_dir: Path, base_settings: dict) -> List[dict]:
    records = []
    for headroom in (0.03, 0.05, 0.08, 0.10):
        settings = dict(base_settings)
        settings["terminal_headroom"] = headroom
        pair_ratios = {}
        for spec in PAIR_SPECS:
            pair = analyze_pair(drift_dir, spec, settings)
            pair_ratios[spec.name] = pair["diagnostic"][
                "forward_test_rmse_ratio_affine_over_constant"]
        records.append({
            "terminal_headroom_v": headroom,
            "forward_test_rmse_ratios_affine_over_constant": pair_ratios,
            "affine_improves_both": bool(all(
                ratio < 1.0 for ratio in pair_ratios.values())),
        })
    return records


def plot_report(report: dict, output: Path) -> None:
    colors = {"alpha": "#4477AA", "beta": "#EE6677"}
    fig, axes = plt.subplots(2, 2, figsize=(9.2, 7.0))
    for column, pair_name in enumerate(("experiment_1", "experiment_2")):
        pair = report["pairs"][pair_name]
        plot_data = pair["_plot"]
        path_axis = axes[0, column]
        velocity_axis = axes[1, column]
        for trace in plot_data["traces"]:
            spec = trace["spec"]
            color = colors[spec.task]
            path_axis.plot(
                trace["raw_gate_minus"], trace["raw_gate_plus"],
                color=color, alpha=0.24, linewidth=1.0)
            gates = trace["selected_gates"]
            path_axis.plot(
                gates[:, 0], gates[:, 1], "o-", color=color,
                markersize=3.0, linewidth=1.2, label=f"task {spec.task}")
            x_line = np.linspace(gates[:, 0].min(), gates[:, 0].max(), 100)
            path_axis.plot(
                x_line, spec.solution_slope * x_line + spec.solution_intercept,
                "--", color=color, linewidth=1.0)

            order = np.argsort(trace["path_coordinate"])
            coordinate = trace["path_coordinate"][order]
            observed = trace["projected_velocity"][order]
            design = trace["design"][order]
            constant_prediction = design[:, :2] @ plot_data["constant_parameters"]
            affine_prediction = design @ plot_data["affine_parameters"]
            velocity_axis.plot(
                coordinate, observed, "o", color=color, markersize=3.5,
                label=f"task {spec.task} measured")
            velocity_axis.plot(
                coordinate, constant_prediction, ":", color=color,
                linewidth=1.3)
            velocity_axis.plot(
                coordinate, affine_prediction, "-", color=color,
                linewidth=1.5)

        path_axis.set_title(
            f"{chr(ord('A') + column)}  Published drift: {pair_name.replace('_', ' ')}")
        path_axis.set_xlabel("gate voltage minus (V)")
        path_axis.set_ylabel("gate voltage plus (V)")
        path_axis.legend(frameon=False, fontsize=8)
        path_axis.grid(alpha=0.18)

        ratio = pair["diagnostic"]["forward_test_rmse_ratio_affine_over_constant"]
        velocity_axis.set_title(
            f"{chr(ord('C') + column)}  Held-out RMSE ratio affine/constant = {ratio:.2f}")
        velocity_axis.set_xlabel("position along solution line (V)")
        velocity_axis.set_ylabel("tangential drift (V/s)")
        velocity_axis.legend(frameon=False, fontsize=7)
        velocity_axis.grid(alpha=0.18)

    fig.suptitle(
        "Released physical drift varies with local gate state\n"
        "dots: measured projection; dotted: constant bias; solid: per-edge affine bias",
        fontsize=11)
    fig.tight_layout()
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=180)
    plt.close(fig)


def strip_plot_data(report: dict) -> dict:
    clean = dict(report)
    clean["pairs"] = {}
    for name, pair in report["pairs"].items():
        clean["pairs"][name] = {key: value for key, value in pair.items() if key != "_plot"}
    return clean


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(
        "--json", type=Path,
        default=Path("results/physical_bias/p0_state_dependence.json"))
    parser.add_argument(
        "--figure", type=Path,
        default=Path("results/figs/physical_bias_state_dependence.png"))
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    root = args.artifact_root.resolve()
    if root.name != SOURCE_TREE:
        raise ValueError(f"expected source root {SOURCE_TREE}, received {root.name}")
    drift_dir = root / "small network" / "DriftTests"
    settings = {
        "minimum_time": 1.05,
        "terminal_headroom": 0.05,
        "maximum_line_distance": 0.03,
        "trim_selected_ends": 1,
    }
    report = {
        "analysis": "published_physical_drift_state_dependence_diagnostic",
        "confirmatory": False,
        "supports_task_learning_improvement": False,
        "provenance": {
            "zenodo_record": ZENODO_RECORD,
            "doi": ZENODO_DOI,
            "release": RELEASE,
            "source_tree": SOURCE_TREE,
        },
        "identification": {
            "observable": "tangential gate velocity after reaching a one-task solution line",
            "assumption": (
                "the clean force has zero tangential projection on the published solution line"
            ),
            "model_restriction": "each edge bias is affine in its own local gate voltage",
            "limitation": (
                "the traces identify a local affine surrogate over their visited states, not the "
                "complete physical vector field"
            ),
        },
        "preprocessing": settings,
        "pairs": {},
    }
    for spec in PAIR_SPECS:
        report["pairs"][spec.name] = analyze_pair(drift_dir, spec, settings)
    report["headroom_sensitivity"] = headroom_sensitivity(drift_dir, settings)
    report["summary"] = {
        "forward_test_rmse_ratios_affine_over_constant": [
            report["pairs"][spec.name]["diagnostic"][
                "forward_test_rmse_ratio_affine_over_constant"]
            for spec in PAIR_SPECS
        ],
        "affine_improves_both_forward_tests": bool(all(
            report["pairs"][spec.name]["diagnostic"][
                "affine_improves_forward_test"]
            for spec in PAIR_SPECS
        )),
        "affine_improves_both_for_all_headroom_thresholds": bool(all(
            record["affine_improves_both"]
            for record in report["headroom_sensitivity"]
        )),
    }
    plot_report(report, args.figure)
    serializable = strip_plot_data(report)
    args.json.parent.mkdir(parents=True, exist_ok=True)
    args.json.write_text(json.dumps(serializable, indent=2) + "\n")
    print(json.dumps(serializable["summary"], indent=2))
    print(f"wrote {args.json}")
    print(f"wrote {args.figure}")


if __name__ == "__main__":
    main()