#!/usr/bin/env python3
"""Deterministic contract checks for the physical coupled-learning adapter."""

from __future__ import annotations

from pathlib import Path
import sys

import numpy as np

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

from sdil.physical_coupled import (
    Circuit,
    LocalAffineBias,
    LocalPredictor,
    Task,
    calibrate_predictor,
    free_output,
    joint_solution,
    local_replay_update,
    overclamped_clean_rate,
    simulate_alternating_tasks,
    standard_clean_rate,
)


def main() -> None:
    circuit = Circuit()
    tasks = (
        Task("alpha", circuit.high, 0.31),
        Task("beta", circuit.low, 0.14),
    )
    joint = joint_solution(circuit, tasks)
    for task in tasks:
        assert abs(free_output(circuit, joint, task.input_voltage) - task.label_voltage) < 1e-12
        rate, _, _ = standard_clean_rate(circuit, joint, task)
        assert np.linalg.norm(rate) < 1e-10

    field = LocalAffineBias(
        reference_gate=np.asarray((3.0, 3.5)),
        bias_at_reference=np.asarray((2.9, 4.7)),
        local_slopes=np.asarray((0.6, -0.2)),
    )
    states = np.column_stack((
        np.linspace(2.2, 4.8, 40),
        np.linspace(2.8, 5.0, 40),
    ))
    scale = np.ptp(states, axis=0)
    affine = LocalPredictor.zeros(field.reference_gate, scale, affine=True)
    constant = LocalPredictor.zeros(field.reference_gate, scale, affine=False)
    observations_affine = calibrate_predictor(
        affine, states, field, epochs=30, learning_rate=0.2)
    observations_constant = calibrate_predictor(
        constant, states, field, epochs=30, learning_rate=0.2)
    assert observations_affine == observations_constant
    affine_error = np.mean([
        np.linalg.norm(field(state) - affine.predict(state)) for state in states
    ])
    constant_error = np.mean([
        np.linalg.norm(field(state) - constant.predict(state)) for state in states
    ])
    assert affine_error < 1e-4
    assert constant_error > 0.1

    gates = states[7]
    teaching = np.asarray((3.2, -1.7))
    eligibility = np.asarray((0.4, 0.8))
    update_a = local_replay_update(affine, gates, teaching, eligibility, 0.03)
    update_b = 0.03 * (teaching - affine.predict(gates)) * eligibility
    assert np.array_equal(update_a, update_b)

    downstream = np.random.default_rng(19).normal(size=(100, 100))
    downstream[:] = np.random.default_rng(29).normal(size=downstream.shape)
    update_c = local_replay_update(affine, gates, teaching, eligibility, 0.03)
    assert np.array_equal(update_a, update_c)

    _, output_free, output_clamped = overclamped_clean_rate(
        circuit,
        joint,
        tasks[0],
        nudging=0.25,
        clamp_magnitude=circuit.high,
        exact_target=True,
    )
    expected_clamped = output_free + 0.25 * (
        circuit.high * np.sign(tasks[0].label_voltage - output_free)
        - output_free
    )
    assert output_clamped == expected_clamped

    exact_predictor = LocalPredictor.zeros(
        field.reference_gate, scale, affine=True)
    exact_predictor.coefficients[:, 0] = field.bias_at_reference
    exact_predictor.coefficients[:, 1] = field.local_slopes * scale
    shared = {
        "period_seconds": 0.01,
        "cycles": 8,
        "initial_gates": np.asarray((4.0, 4.0)),
        "overclamp_nudging": 0.25,
        "overclamp_magnitude": circuit.high,
        "overclamp_exact_target": True,
    }
    clean_overclamp = simulate_alternating_tasks(
        circuit,
        tasks,
        field,
        method="overclamp",
        bias_strength=0.0,
        **shared,
    )
    debiased_overclamp = simulate_alternating_tasks(
        circuit,
        tasks,
        field,
        method="overclamp_sdil",
        bias_strength=1.0,
        predictor=exact_predictor,
        **shared,
    )
    assert np.allclose(
        clean_overclamp["final_gates"],
        debiased_overclamp["final_gates"],
        atol=1e-12,
        rtol=0.0,
    )
    assert debiased_overclamp["max_abs_clamp_displacement_v"] > 0.0
    print({
        "joint_solution": joint.tolist(),
        "affine_calibration_mae": float(affine_error),
        "constant_calibration_mae": float(constant_error),
        "neutral_observations_each": observations_affine,
        "local_replay_exact": True,
        "sdil_overclamp_composition_exact": True,
        "autodiff_used": False,
    })


if __name__ == "__main__":
    main()