#!/usr/bin/env python3
"""Strict-locality smoke test for the shared two-state SDIL filter."""

from __future__ import annotations

from pathlib import Path
import sys

import torch

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

from sdil.two_state_debias import (  # noqa: E402
    BatchedLocalAffineDebiaser,
    LocalAffineDebiaser,
    two_state_difference,
)


def main() -> None:
    torch.manual_seed(20260806)
    templates = [torch.zeros(7), torch.zeros(3, 4)]
    centers = [0.2, -0.1]
    scales = [1.3, 0.8]
    affine = LocalAffineDebiaser(
        templates,
        feature_centers=centers,
        feature_scales=scales,
        affine=True,
    )
    constant = LocalAffineDebiaser(
        templates,
        feature_centers=centers,
        feature_scales=scales,
        affine=False,
    )
    intercepts = [torch.linspace(-0.4, 0.5, 7), torch.randn(3, 4) * 0.2]
    slopes = [torch.linspace(0.3, 1.0, 7), torch.randn(3, 4) * 0.4]

    train_features = []
    for value in torch.linspace(-1.5, 1.1, 50):
        train_features.append([
            torch.full_like(templates[0], value),
            torch.full_like(templates[1], -0.6 * value + 0.2),
        ])
    for _ in range(20):
        for features in train_features:
            measurements = [
                intercept + slope * ((feature - center) / scale)
                for intercept, slope, feature, center, scale in zip(
                    intercepts, slopes, features, centers, scales)
            ]
            affine.update_neutral(features, measurements, 0.2)
            constant.update_neutral(features, measurements, 0.2)

    held_features = [
        torch.full_like(templates[0], 1.35),
        torch.full_like(templates[1], -0.75),
    ]
    held_bias = [
        intercept + slope * ((feature - center) / scale)
        for intercept, slope, feature, center, scale in zip(
            intercepts, slopes, held_features, centers, scales)
    ]
    affine_error = torch.mean(torch.stack([
        (prediction - target).square().mean()
        for prediction, target in zip(affine.predict(held_features), held_bias)
    ])).sqrt()
    constant_error = torch.mean(torch.stack([
        (prediction - target).square().mean()
        for prediction, target in zip(constant.predict(held_features), held_bias)
    ])).sqrt()
    assert affine_error < 1e-5
    assert constant_error > 0.1

    first = [torch.randn_like(template) for template in templates]
    common = [torch.randn_like(template) for template in templates]
    second = [value + 0.03 * torch.ones_like(value) for value in first]
    difference_a = two_state_difference(first, second, 0.2)
    difference_b = two_state_difference(
        [value + offset for value, offset in zip(first, common)],
        [value + offset for value, offset in zip(second, common)],
        0.2,
    )
    common_mode_error = max(
        float((a - b).abs().max()) for a, b in zip(difference_a, difference_b)
    )
    assert common_mode_error < 2e-6

    teaching = [torch.randn_like(template) for template in templates]
    eligibility = [torch.randn_like(template) for template in templates]
    update_a = affine.replay_updates(
        held_features, teaching, eligibility, 0.07)
    downstream = torch.randn(1024, 1024)
    downstream.normal_()
    update_b = affine.replay_updates(
        held_features, teaching, eligibility, 0.07)
    assert all(torch.equal(a, b) for a, b in zip(update_a, update_b))
    assert all(not update.requires_grad for update in update_a)

    try:
        affine.predict([held_features[0].requires_grad_(), held_features[1]])
    except ValueError:
        pass
    else:
        raise AssertionError("requires-grad input was not rejected")

    feature = torch.linspace(-1.0, 1.0, 64).reshape(64, 1, 1)
    feature = torch.cat((feature, feature.square()), dim=2)
    target = 0.2 + torch.tensor([[[0.7, -0.4]]]) * feature
    batched_affine = BatchedLocalAffineDebiaser(
        [feature], feature_centers=[0.0], feature_scales=[1.0], affine=True)
    batched_constant = BatchedLocalAffineDebiaser(
        [feature], feature_centers=[0.0], feature_scales=[1.0], affine=False)
    for _ in range(20):
        batched_affine.update_neutral([feature], [target], 0.2)
        batched_constant.update_neutral([feature], [target], 0.2)
    held_feature = torch.tensor([[[-1.3, 1.4]], [[1.3, 1.4]]])
    held_target = 0.2 + torch.tensor([[[0.7, -0.4]]]) * held_feature
    held_affine = batched_affine.residual([held_feature], [held_target])[0]
    held_constant = batched_constant.residual([held_feature], [held_target])[0]
    batched_affine_rmse = float(held_affine.square().mean().sqrt())
    batched_constant_rmse = float(held_constant.square().mean().sqrt())
    assert batched_affine_rmse < 0.05 * batched_constant_rmse

    print({
        "affine_heldout_rmse": float(affine_error),
        "constant_heldout_rmse": float(constant_error),
        "neutral_observations_each": affine.neutral_observations,
        "common_mode_max_float_error": common_mode_error,
        "downstream_independence_exact": True,
        "requires_grad_rejected": True,
        "batched_affine_heldout_rmse": batched_affine_rmse,
        "batched_constant_heldout_rmse": batched_constant_rmse,
        "batched_neutral_observations": batched_affine.neutral_observations,
    })


if __name__ == "__main__":
    main()