#!/usr/bin/env python3
"""Deterministic locality and numerical audit for the digital CLLN ladder."""

from __future__ import annotations

import inspect
import json
from pathlib import Path
import sys

import numpy as np

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

import sdil.coupled_ladder as ladder  # noqa: E402
from sdil.coupled_ladder import (  # noqa: E402
    DigitalTrainingConfig,
    make_scaled_grid,
    solve_linear_grid_state,
    tile_figure5_gates,
    train_digital_grid,
)
from sdil.physical_grid import (  # noqa: E402
    GridCircuit,
    GridSquareLawImperfection,
    RingClassificationDataset,
    edge_voltage_drops,
)


def maximum_unknown_kcl_residual(
    circuit: GridCircuit, gates: np.ndarray, voltages: np.ndarray
) -> float:
    conductances = circuit.conductance_scale * (
        gates - circuit.threshold_voltage)
    currents = np.zeros(circuit.node_count)
    for conductance, (first, second) in zip(
        conductances, circuit.edge_pairs
    ):
        current = conductance * (voltages[first] - voltages[second])
        currents[first] += current
        currents[second] -= current
    fixed = set(circuit.source_nodes)
    unknown = [
        node for node in range(circuit.node_count) if node not in fixed
    ]
    return float(np.max(np.abs(currents[unknown])))


def main() -> None:
    protocol_path = Path(
        "results/physical_bias/dillavou_fig5_protocol.json")
    protocol = json.loads(protocol_path.read_text())
    task = next(
        record for record in protocol["experiments"]
        if record["method"] == "standard"
    )
    base_gates = np.asarray(task["initial_gates_v"], dtype=float)
    dataset = RingClassificationDataset(
        inputs_v=np.asarray(task["inputs_v"], dtype=float).T,
        labels_v=(
            2.0 * np.asarray(task["classes"], dtype=float) - 1.0
        ) * 0.018,
    )

    side4 = make_scaled_grid(4)
    layout_exact = bool(
        side4.source_nodes == GridCircuit().source_nodes
        and side4.target_nodes == GridCircuit().target_nodes
        and side4.edge_pairs == GridCircuit().edge_pairs
    )
    tiling_exact = bool(np.array_equal(
        tile_figure5_gates(base_gates, 4), base_gates))

    side32 = make_scaled_grid(32)
    gates32 = tile_figure5_gates(base_gates, 32)
    state32 = solve_linear_grid_state(
        side32,
        gates32,
        side32.source_values(*dataset.inputs_v[0]),
    )
    kcl_residual = maximum_unknown_kcl_residual(
        side32, gates32, state32)

    edge_count = side4.edge_count
    rng = np.random.default_rng(20260829)
    free_drops = rng.normal(0.0, 0.1, edge_count)
    clamped_drops = rng.normal(0.0, 0.1, edge_count)
    common_offset = rng.normal(0.0, 2.3, edge_count)
    common_mode = GridSquareLawImperfection(
        free_gain=np.ones(edge_count),
        clamped_gain=np.ones(edge_count),
        free_input_offset_v=np.zeros(edge_count),
        clamped_input_offset_v=np.zeros(edge_count),
        multiplier_output_offset_v_per_s=common_offset,
    )
    ideal_rate = common_mode.ideal_rate(
        side4.measured_learning_rate, free_drops, clamped_drops)
    observed_rate = common_mode.observed_rate(
        side4.measured_learning_rate, free_drops, clamped_drops)
    neutral_rate = common_mode.neutral_bias(
        side4.measured_learning_rate, free_drops)
    common_mode_cancellation_error = float(np.max(np.abs(
        observed_rate - neutral_rate - ideal_rate)))

    state_dependent = GridSquareLawImperfection.sample_appendix_c(
        edge_count, 20260829)
    state_one = rng.normal(0.0, 0.05, edge_count)
    state_two = rng.normal(0.0, 0.15, edge_count)
    bias_one = state_dependent.neutral_bias(
        side4.measured_learning_rate, state_one)
    bias_two = state_dependent.neutral_bias(
        side4.measured_learning_rate, state_two)
    state_dependent_bias_rms = float(np.sqrt(np.mean(np.square(
        bias_one - bias_two))))

    ideal = GridSquareLawImperfection.ideal(edge_count)
    config = DigitalTrainingConfig(
        epochs=20, record_every=5, learning_time_seconds=0.01)
    trajectories = {
        method: train_digital_grid(
            side4,
            base_gates,
            dataset,
            ideal,
            method=method,
            config=config,
            noise_seed=20260829,
        )
        for method in ("clean", "matched_noise", "raw", "sdil")
    }
    reference_gates = np.asarray(trajectories["clean"]["final_gates_v"])
    ideal_path_max_gate_difference = float(max(
        np.max(np.abs(
            np.asarray(trajectories[method]["final_gates_v"])
            - reference_gates
        ))
        for method in ("matched_noise", "raw", "sdil")
    ))
    ideal_path_trace_identity = bool(all(
        trajectories[method]["trace"] == trajectories["clean"]["trace"]
        for method in ("matched_noise", "raw", "sdil")
    ))
    sdil_cost_identity = bool(
        trajectories["sdil"]["neutral_scalar_observations"]
        == trajectories["sdil"]["local_edge_updates"]
    )

    source = inspect.getsource(ladder)
    forbidden_source_tokens = {
        token: token in source
        for token in ("import torch", "autograd", ".backward(")
    }
    no_autodiff_dependency = not any(forbidden_source_tokens.values())

    checks = {
        "released_side4_layout_exact": layout_exact,
        "released_side4_gate_vector_exact": tiling_exact,
        "side32_kcl_residual_below_1e_12": kcl_residual < 1e-12,
        "common_mode_cancellation_below_1e_12": (
            common_mode_cancellation_error < 1e-12),
        "default_imperfection_is_state_dependent": (
            state_dependent_bias_rms > 1e-6),
        "ideal_paths_have_identical_gates": (
            ideal_path_max_gate_difference == 0.0),
        "ideal_paths_have_identical_traces": ideal_path_trace_identity,
        "one_neutral_scalar_observation_per_edge_update": (
            sdil_cost_identity),
        "no_autodiff_dependency": no_autodiff_dependency,
    }
    report = {
        "analysis": "digital_coupled_ladder_locality_audit",
        "autodiff_used": False,
        "source_protocol": str(protocol_path),
        "checks": checks,
        "measurements": {
            "side32_maximum_unknown_kcl_residual_a": kcl_residual,
            "common_mode_cancellation_max_abs_v_per_s": (
                common_mode_cancellation_error),
            "state_dependent_neutral_bias_difference_rms_v_per_s": (
                state_dependent_bias_rms),
            "ideal_path_max_gate_difference_v": (
                ideal_path_max_gate_difference),
            "sdil_local_edge_updates": (
                trajectories["sdil"]["local_edge_updates"]),
            "sdil_neutral_scalar_observations": (
                trajectories["sdil"]["neutral_scalar_observations"]),
            "forbidden_source_tokens": forbidden_source_tokens,
        },
        "gate": "pass" if all(checks.values()) else "fail",
    }
    output = Path("results/coupled_ladder/x0_locality_audit.json")
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))
    print(f"wrote {output}")
    if report["gate"] != "pass":
        raise SystemExit(1)


if __name__ == "__main__":
    main()