#!/usr/bin/env python3
"""Mechanics checks for the reconstructed Figure-5 physical grid."""

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_grid import (  # noqa: E402
    EdgePolynomialPredictor,
    GridCircuit,
    GridSquareLawImperfection,
    _residual_and_jacobian,
    edge_voltage_drops,
    fit_edge_predictor,
    output_difference,
    solve_grid_state,
)


def sample_local_states(
    circuit: GridCircuit, count: int, seed: int
) -> np.ndarray:
    rng = np.random.default_rng(seed)
    states = []
    for _ in range(count):
        gates = rng.normal(2.33, 0.08, circuit.edge_count)
        sources = circuit.source_values(
            rng.uniform(circuit.low_voltage, circuit.high_voltage),
            rng.uniform(circuit.low_voltage, circuit.high_voltage),
        )
        voltages = solve_grid_state(circuit, gates, sources)
        states.append(edge_voltage_drops(circuit, voltages))
    return np.asarray(states)


def main() -> None:
    circuit = GridCircuit()
    assert circuit.edge_count == 32
    assert circuit.source_nodes == (15, 13, 7, 5)
    assert circuit.target_nodes == (10, 8)
    gates = np.full(circuit.edge_count, 2.33)
    sources = circuit.source_values(0.2266, 0.2266)
    free = solve_grid_state(circuit, gates, sources)
    assert np.allclose(free[list(circuit.source_nodes)], sources)
    assert abs(output_difference(circuit, free)) < 1e-12
    residual, _ = _residual_and_jacobian(circuit, gates, free)
    unknown = [
        node for node in range(circuit.node_count)
        if node not in circuit.source_nodes
    ]
    assert np.linalg.norm(residual[unknown], ord=np.inf) < 1e-10

    target_difference = 0.018
    target_mean = float(np.mean(free[list(circuit.target_nodes)]))
    target_values = np.asarray((
        target_mean + 0.5 * target_difference,
        target_mean - 0.5 * target_difference,
    ))
    clamped = solve_grid_state(
        circuit, gates, sources, target_values=target_values,
        initial_state=free)
    assert np.allclose(clamped[list(circuit.target_nodes)], target_values)
    assert abs(output_difference(circuit, clamped) - target_difference) < 1e-12

    ideal = GridSquareLawImperfection.ideal(circuit.edge_count)
    free_drops = edge_voltage_drops(circuit, free)
    clamped_drops = edge_voltage_drops(circuit, clamped)
    assert np.array_equal(
        ideal.observed_rate(
            circuit.measured_learning_rate, free_drops, clamped_drops),
        ideal.ideal_rate(
            circuit.measured_learning_rate, free_drops, clamped_drops),
    )

    hardware = GridSquareLawImperfection.sample_appendix_c(
        circuit.edge_count, 20260829)
    calibration_states = sample_local_states(circuit, 48, 20260830)
    neutral = np.asarray([
        hardware.neutral_bias(circuit.measured_learning_rate, state)
        for state in calibration_states
    ])
    center = np.mean(calibration_states, axis=0)
    scale = np.maximum(np.ptp(calibration_states, axis=0), 1e-3)
    constant = EdgePolynomialPredictor.zeros(center, scale, degree=0)
    quadratic = EdgePolynomialPredictor.zeros(center, scale, degree=2)
    assert fit_edge_predictor(constant, calibration_states, neutral) == 48
    assert fit_edge_predictor(quadratic, calibration_states, neutral) == 48
    heldout_states = sample_local_states(circuit, 16, 20260831)
    heldout_neutral = np.asarray([
        hardware.neutral_bias(circuit.measured_learning_rate, state)
        for state in heldout_states
    ])
    constant_rmse = float(np.sqrt(np.mean([
        np.square(measurement - constant.predict(state))
        for state, measurement in zip(heldout_states, heldout_neutral)
    ])))
    quadratic_rmse = float(np.sqrt(np.mean([
        np.square(measurement - quadratic.predict(state))
        for state, measurement in zip(heldout_states, heldout_neutral)
    ])))
    assert quadratic_rmse < 1e-7
    assert constant_rmse > 1e-3
    print({
        "nodes": circuit.node_count,
        "edges": circuit.edge_count,
        "source_nodes": circuit.source_nodes,
        "target_nodes": circuit.target_nodes,
        "constant_neutral_rmse_v_per_s": constant_rmse,
        "quadratic_neutral_rmse_v_per_s": quadratic_rmse,
        "autodiff_used": False,
    })


if __name__ == "__main__":
    main()