#!/usr/bin/env python3
"""Strict-locality smoke test for the Rain neuron-state adapter."""

from __future__ import annotations

import argparse
from pathlib import Path
import sys

import torch

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

from sdil.rain_ep_adapter import (  # noqa: E402
    RainLayerStateCorrector,
    attach_layer_to_rain_estimator,
)


def build_estimator(author_root: Path):
    sys.path.insert(0, str(author_root))
    from model.function.cost import SquaredError
    from model.function.network import Network
    from model.hopfield.minimizer import FixedPointMinimizer
    from model.hopfield.network import DeepHopfieldEnergy
    from training.sgd import AugmentedFunction, EquilibriumProp

    energy = DeepHopfieldEnergy([(4,), (7,), (3,)], [0.5, 0.5])
    energy.set_device("cpu")
    network = Network(energy)
    cost = SquaredError(energy.layers()[-1])
    augmented = AugmentedFunction(energy, cost)
    minimizer = FixedPointMinimizer(augmented, network.free_layers())
    minimizer.mode = "asynchronous"
    minimizer.num_iterations = 12
    estimator = EquilibriumProp(
        energy.params(), energy.layers(), augmented, cost, minimizer)
    estimator.variant = "positive"
    estimator.nudging = 0.25
    return network, cost, augmented, minimizer, estimator


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--author-root", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    torch.manual_seed(20260806)
    network, cost, augmented, minimizer, estimator = build_estimator(
        args.author_root)
    x = torch.randn(64, 4)
    labels = torch.arange(64) % 3
    network.set_input(x, reset=True)
    cost.set_target(labels)
    augmented.nudging = 0.0
    minimizer.compute_equilibrium()
    free = [layer.state.clone() for layer in minimizer._layers]
    clean = [value.clone() for value in estimator.compute_gradient()]
    for layer, state in zip(minimizer._layers, free):
        layer.state = state.clone()
    oracle_corrector = RainLayerStateCorrector(
        mode="oracle", bias_ratio=4.0, seed=19)
    attach_layer_to_rain_estimator(estimator, oracle_corrector)
    oracle = estimator.compute_gradient()
    oracle_relative_error = max(
        float((actual - target).norm() / target.norm().clamp_min(1e-30))
        for actual, target in zip(oracle, clean)
    )
    assert oracle_relative_error < 2e-5, oracle_relative_error

    first = {
        "input": torch.randn(64, 4),
        "hidden": torch.randn(64, 7),
        "output": torch.randn(64, 3),
    }
    clean_difference = {
        name: 0.01 * torch.randn_like(value)
        for name, value in first.items()
    }
    second = {
        name: value + clean_difference[name]
        for name, value in first.items()
    }
    innovation = RainLayerStateCorrector(
        mode="innovation", bias_ratio=4.0, predictor_rate=0.2,
        calibration_steps=1, seed=31)
    constant = RainLayerStateCorrector(
        mode="constant", bias_ratio=4.0, predictor_rate=0.2,
        calibration_steps=1, seed=31)
    layer_names = ["hidden", "output"]
    innovation.apply(first, second, layer_names)
    constant.apply(first, second, layer_names)
    held_first = {
        name: 1.25 * value + 0.1 for name, value in first.items()
    }
    held_clean = {
        name: 0.01 * torch.randn_like(value)
        for name, value in first.items()
    }
    held_second = {
        name: value + held_clean[name]
        for name, value in held_first.items()
    }
    innovation_used = innovation.apply(
        held_first, held_second, layer_names)
    constant_used = constant.apply(held_first, held_second, layer_names)

    def residual_rms(used):
        errors = [
            used[name] - held_second[name] for name in layer_names
        ]
        return (
            sum(float(error.square().sum()) for error in errors)
            / sum(error.numel() for error in errors)
        ) ** 0.5

    innovation_error = residual_rms(innovation_used)
    constant_error = residual_rms(constant_used)
    assert innovation_error < 0.1 * constant_error, (
        innovation_error, constant_error)
    assert innovation.debiaser.neutral_observations == 64
    assert constant.debiaser.neutral_observations == 64
    assert all(not value.requires_grad for value in innovation_used.values())
    print({
        "oracle_parameter_gradient_relative_error": oracle_relative_error,
        "innovation_heldout_state_residual_rms": innovation_error,
        "constant_heldout_state_residual_rms": constant_error,
        "matched_neutral_observations": 64,
        "extra_equilibrium_phases": 0,
        "requires_grad": False,
    })


if __name__ == "__main__":
    main()