#!/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,
    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)
    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,
        "autodiff_used": False,
    })


if __name__ == "__main__":
    main()