#!/usr/bin/env python3
"""Mechanics checks for the Appendix-C physical imperfection model."""

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, Task  # noqa: E402
from sdil.physical_imperfection import (  # noqa: E402
    DifferentialSquareLawImperfection,
    LocalPolynomialPredictor,
    edge_voltage_drops,
    fit_polynomial_predictor,
    simulate_imperfect_alternating_tasks,
)


def main() -> None:
    circuit = Circuit()
    ideal = DifferentialSquareLawImperfection.ideal()
    output_free = 0.23
    output_clamped = 0.26
    assert np.array_equal(
        ideal.observed_rate(circuit, output_free, output_clamped),
        ideal.ideal_rate(circuit, output_free, output_clamped),
    )
    assert np.array_equal(
        ideal.neutral_bias(circuit, output_free), np.zeros(2))

    hardware = DifferentialSquareLawImperfection.sample_appendix_c(20260829)
    outputs = np.linspace(circuit.low + 0.01, circuit.high - 0.01, 81)
    local_states = np.asarray([
        edge_voltage_drops(circuit, output) for output in outputs])
    neutral = np.asarray([
        hardware.neutral_bias(circuit, output) for output in outputs])
    center = np.mean(local_states, axis=0)
    scale = np.ptp(local_states, axis=0)
    constant = LocalPolynomialPredictor.zeros(center, scale, degree=0)
    quadratic = LocalPolynomialPredictor.zeros(center, scale, degree=2)
    constant_count = fit_polynomial_predictor(
        constant,
        local_states,
        neutral,
    )
    quadratic_count = fit_polynomial_predictor(
        quadratic,
        local_states,
        neutral,
    )
    assert constant_count == quadratic_count
    constant_rmse = float(np.sqrt(np.mean([
        np.square(measurement - constant.predict(state))
        for state, measurement in zip(local_states, neutral)
    ])))
    quadratic_rmse = float(np.sqrt(np.mean([
        np.square(measurement - quadratic.predict(state))
        for state, measurement in zip(local_states, neutral)
    ])))
    assert quadratic_rmse < 1e-6
    assert constant_rmse > 1e-3

    tasks = (
        Task("alpha", circuit.high, 0.31),
        Task("beta", circuit.low, 0.14),
    )
    shared = {
        "circuit": circuit,
        "tasks": tasks,
        "imperfection": hardware,
        "period_seconds": 0.01,
        "cycles": 12,
        "initial_gates": np.asarray((4.0, 4.0)),
        "overclamp_nudging": 0.05,
        "overclamp_magnitude": circuit.high,
    }
    oracle = simulate_imperfect_alternating_tasks(
        method="overclamp_oracle_neutral", **shared)
    sdil = simulate_imperfect_alternating_tasks(
        method="overclamp_sdil", predictor=quadratic, **shared)
    assert np.allclose(
        oracle["final_gates"], sdil["final_gates"], atol=1e-8, rtol=0.0)
    print({
        "constant_neutral_rmse_v_per_s": constant_rmse,
        "quadratic_neutral_rmse_v_per_s": quadratic_rmse,
        "neutral_observations_each": constant_count,
        "sdil_matches_neutral_oracle": True,
        "autodiff_used": False,
    })


if __name__ == "__main__":
    main()