#!/usr/bin/env python3
"""Endpoint-free mechanics checks for the independent oral-B-v2 branch."""
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from sdil.bci import generate_trajectories
from sdil.bci_v2 import BCIV2, BCIV2Config, run_day_v2


def check_role_estimator():
    cfg = BCIV2Config(
        days=2, episodes_per_day=8192, steps_per_episode=2,
        feedback="performance_velocity",
    )
    model = BCIV2(cfg, model_seed=1)
    generator = torch.Generator().manual_seed(2)
    soma = 0.2 * torch.randn(
        cfg.episodes_per_day, cfg.n_neurons, generator=generator
    )
    perturbations = torch.empty_like(soma).bernoulli_(
        0.5, generator=generator
    ).mul_(2).sub_(1)
    sigma = cfg.perturb_sigma
    cursor_plus = (soma + sigma * perturbations) @ model.role
    cursor_minus = (soma - sigma * perturbations) @ model.role
    saved_role = model.role.clone()
    target = model.causal_role_targets(
        cursor_plus, cursor_minus, perturbations
    )
    model.role.zero_()
    assert torch.equal(
        target,
        model.causal_role_targets(
            cursor_plus, cursor_minus, perturbations
        ),
    )
    model.role.copy_(saved_role)
    estimate = target.mean(0)
    cosine = torch.nn.functional.cosine_similarity(
        estimate, model.role, dim=0
    ).item()
    assert cosine > 0.99
    active = torch.ones(cfg.episodes_per_day, dtype=torch.bool)
    model.A.zero_()
    model.update_role(target, active)
    learned_cosine = torch.nn.functional.cosine_similarity(
        model.A, model.role, dim=0
    ).item()
    assert learned_cosine > 0.99
    print(
        "v2 antithetic causal-role estimator: "
        f"target_cos={cosine:.6f}, update_cos={learned_cosine:.6f}"
    )


def check_predictor():
    cfg = BCIV2Config(
        days=2, episodes_per_day=256, steps_per_episode=2,
        feedback="performance_velocity",
    )
    model = BCIV2(cfg, model_seed=3)
    generator = torch.Generator().manual_seed(4)
    active = torch.ones(cfg.episodes_per_day, dtype=torch.bool)
    for _ in range(100):
        soma = torch.randn(
            cfg.episodes_per_day,
            cfg.n_neurons,
            generator=generator,
        ).tanh()
        model.update_predictor(soma, model.coupling * soma, active)
    assert (model.P - model.coupling).abs().max() < 1e-4
    assert model.P_bias.abs().max() < 1e-5
    print("v2 neutral predictor identifies soma-predictable traffic: exact")


def check_td_and_eligibility_updates():
    cfg = BCIV2Config(
        days=2, episodes_per_day=4, steps_per_episode=3,
        target=10.0, process_noise=0.0, inertia=0.0,
        predictor_eta=0.2, vectorizer_eta=0.03,
        forward_eta=0.1, perturb_every=99,
        gamma=0.8, critic_eta=0.05, eligibility_decay=0.6,
        velocity_reward_scale=0.25, terminal_reward=1.0,
        feedback="performance_velocity",
    )
    initial = BCIV2(cfg, model_seed=5)
    initial.P.copy_(initial.coupling)
    initial.A.copy_(initial.role)
    initial.critic.copy_(torch.linspace(
        -0.1, 0.1, cfg.n_neurons + 1
    ))
    model = initial.clone()
    state = model.initial_episode_state(cfg.episodes_per_day)
    context = torch.arange(
        cfg.episodes_per_day * cfg.context_dim, dtype=torch.float32
    ).reshape(cfg.episodes_per_day, cfg.context_dim) / 100.0
    noise = torch.zeros(cfg.episodes_per_day, cfg.n_neurons)
    perturbations = torch.ones_like(noise)

    previous_critic = model.critic.clone()
    previous_w = model.W.clone()
    previous_features = state["value_features"].clone()
    previous_value = state["value"].clone()
    next_state, record = model.step(
        context,
        noise,
        perturbations,
        state,
        global_step=1,
        final_step=False,
        learn_role=False,
        learn_predictor=False,
    )
    expected_local = (
        (1.0 - record["soma"].square()).unsqueeze(2)
        * context.unsqueeze(1)
    )
    assert torch.allclose(record["eligibility"], expected_local)
    expected_delta = (
        record["reward"]
        + cfg.gamma * record["next_value_prediction"]
        - previous_value
    )
    assert torch.allclose(record["td_innovation"], expected_delta)
    expected_critic = previous_critic + cfg.critic_eta * (
        expected_delta.unsqueeze(1) * previous_features
    ).mean(0)
    assert torch.allclose(model.critic, expected_critic)
    expected_w = previous_w + cfg.forward_eta * (
        record["innovation"][:, :, None] * expected_local
    ).mean(0)
    assert torch.allclose(model.W, expected_w)
    assert torch.allclose(
        next_state["value"], model.value(next_state["value_features"])
    )

    second_state, second = model.step(
        context,
        noise,
        perturbations,
        next_state,
        global_step=2,
        final_step=False,
        plasticity_gain=0.0,
        learn_role=False,
        learn_predictor=False,
        learn_critic=False,
    )
    second_local = (
        (1.0 - second["soma"].square()).unsqueeze(2)
        * context.unsqueeze(1)
    )
    assert torch.allclose(
        second["eligibility"],
        cfg.eligibility_decay * expected_local + second_local,
    )
    assert torch.equal(second_state["active"], torch.ones(4, dtype=torch.bool))
    print("v2 TD(0), local eligibility, critic, and forward updates: exact")


def check_terminal_outcomes_and_lesions():
    cfg = BCIV2Config(
        days=2, episodes_per_day=2, steps_per_episode=2,
        target=0.2, inertia=0.0, process_noise=0.0,
        forward_eta=0.1, gamma=0.9, critic_eta=0.03,
        velocity_reward_scale=0.0, terminal_reward=1.0,
        feedback="performance_velocity",
    )
    initial = BCIV2(cfg, model_seed=6)
    initial.P.copy_(initial.coupling)
    initial.A.copy_(initial.role)
    initial.critic.zero_()
    initial.critic[0] = 0.6
    context = torch.zeros(2, cfg.context_dim)
    context[:, 0] = 1.0
    noise = torch.zeros(2, cfg.n_neurons)
    noise[0, :cfg.n_plus] = 1.0
    perturbations = torch.ones_like(noise)
    state = initial.initial_episode_state(2)

    intact = initial.clone()
    next_state, record = intact.step(
        context,
        noise,
        perturbations,
        state,
        global_step=1,
        final_step=True,
        plasticity_gain=0.0,
        learn_role=False,
        learn_predictor=False,
        learn_critic=False,
    )
    assert record["terminal"].all()
    assert record["outcome"].tolist() == [True, False]
    assert torch.allclose(
        record["td_innovation"], torch.tensor([0.4, -0.6])
    )
    assert not bool(next_state["active"].any())

    no_critic = initial.clone()
    _, critic_lesion = no_critic.step(
        context,
        noise,
        perturbations,
        state,
        global_step=1,
        final_step=True,
        plasticity_gain=0.0,
        learn_role=False,
        learn_predictor=False,
        learn_critic=False,
        critic_enabled=False,
    )
    assert torch.allclose(
        critic_lesion["td_innovation"], torch.tensor([1.0, 0.0])
    )

    no_outcome = initial.clone()
    _, outcome_lesion = no_outcome.step(
        context,
        noise,
        perturbations,
        state,
        global_step=1,
        final_step=True,
        plasticity_gain=0.0,
        learn_role=False,
        learn_predictor=False,
        learn_critic=False,
        terminal_outcome_enabled=False,
    )
    assert torch.allclose(
        outcome_lesion["td_innovation"], torch.tensor([-0.6, -0.6])
    )
    print("v2 terminal reward, timeout surprise, and acute lesions: exact")


def check_horizon_and_pairing():
    cfg = BCIV2Config(
        days=2, episodes_per_day=16, steps_per_episode=8,
        target=10.0, feedback="performance_velocity",
    )
    trajectories = generate_trajectories(cfg, 7)
    first = BCIV2(cfg, model_seed=8)
    second = first.clone()
    report = run_day_v2(
        first,
        trajectories,
        0,
        collect=True,
        horizon=4,
        plasticity_gain=0.0,
        learn_role=False,
        learn_predictor=False,
        learn_critic=False,
    )
    repeated = run_day_v2(
        second,
        trajectories,
        0,
        collect=True,
        horizon=4,
        plasticity_gain=0.0,
        learn_role=False,
        learn_predictor=False,
        learn_critic=False,
    )
    assert report["global_step"] == 4
    assert len(report["events"]) == 4
    terminal_count = torch.stack([
        event["terminal"] for event in report["events"]
    ]).sum(0)
    assert torch.equal(terminal_count, torch.ones(16, dtype=torch.long))
    for left, right in zip(report["events"], repeated["events"]):
        assert all(
            torch.equal(left[key], right[key])
            for key in left
            if isinstance(left[key], torch.Tensor)
        )
    print("v2 fixed-horizon terminal accounting and paired replay: exact")


if __name__ == "__main__":
    check_role_estimator()
    check_predictor()
    check_td_and_eligibility_updates()
    check_terminal_outcomes_and_lesions()
    check_horizon_and_pairing()
    print("ALL ORAL-B-V2 MECHANICS CHECKS PASSED")