"""Candidate-route workflow tested with a real deterministic decoder."""
import math
from pathlib import Path
import sys
import unittest

import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from demos.path_routes import aggregate_pair_effects, measure_receiver_paths, select_candidates
from test_path_patching import CircuitModel
from vlm_probing import CapabilityError, Prober


class CharacterTokenizer:
    """Exact token boundaries for the real-decoder integration test."""
    chat_template = None

    def __call__(self, text, **kwargs):
        assert kwargs == {"return_tensors": "pt", "add_special_tokens": False}
        ids = torch.tensor([[ord(character) - ord("a") + 1 for character in text]])
        return {"input_ids": ids, "attention_mask": torch.ones_like(ids)}


class ReceiverPathTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.threads = torch.get_num_threads()
        torch.set_num_threads(1)

    @classmethod
    def tearDownClass(cls):
        torch.set_num_threads(cls.threads)

    def test_candidate_selection_respects_decoder_order_and_physical_head_counts(self):
        counts = {layer: 2 if layer % 2 else 3 for layer in range(20)}
        ranking = [{"layer": 17, "head": 1, "mean_effect": 0.9},
                   {"layer": 19, "head": 1, "mean_effect": -0.8},
                   {"layer": 18, "head": 2, "mean_effect": 0.7},
                   {"layer": 2, "head": 2, "mean_effect": -0.6},
                   {"layer": 3, "head": 0, "mean_effect": 0.6}]
        senders, receivers = select_candidates(list(reversed(ranking)), counts, 3, 2)
        self.assertEqual(receivers, [(19, 1, "q"), (18, 2, "q")])
        self.assertEqual(senders, [(17, 1), (2, 2), (3, 0)])
        for sender in senders:
            self.assertLess(sender[0], min(receiver[0] for receiver in receivers))
        with self.assertRaisesRegex(ValueError, "unavailable"):
            select_candidates([*ranking, {"layer": 19, "head": 2, "mean_effect": 5}], counts)
        with self.assertRaisesRegex(ValueError, "duplicate"):
            select_candidates([*ranking, ranking[0]], counts)
        with self.assertRaisesRegex(ValueError, "not enough ranked sender"):
            select_candidates(ranking, counts, 4, 2)

    def test_aggregation_averages_per_pair_ratios_and_keeps_zero(self):
        records = [{"raw_delta": 2, "normalizer": 4, "normalized_effect": 0.5},
                   {"raw_delta": 0, "normalizer": 1, "normalized_effect": 0.0},
                   {"raw_delta": 1, "normalizer": 0, "normalized_effect": None}]
        self.assertEqual(aggregate_pair_effects(records), 0.25)
        self.assertNotEqual(aggregate_pair_effects(records), (2 + 0) / (4 + 1))
        self.assertIsNone(aggregate_pair_effects([records[-1]]))
        with self.assertRaisesRegex(ValueError, "finite"):
            aggregate_pair_effects([{"normalized_effect": math.inf}])

    def test_workflow_records_actual_query_paths_and_zero_controls(self):
        with torch.random.fork_rng(devices=[]):
            torch.manual_seed(194)
            model = CircuitModel()
        probe = Prober(model)
        pairs = [{"id": identifier, "clean_prompt": clean, "donor_prompt": donor,
                  "clean_answer": "c", "donor_answer": "k"}
                 for identifier, clean, donor in (("one", "abcd", "efgh"),
                                                   ("two", "dcba", "hgfe"))]
        artifact = measure_receiver_paths(probe, pairs, CharacterTokenizer(), "cpu",
                                          [(0, 0), (1, 0)], [(2, 0, "q")])
        self.assertEqual(len(artifact["edges"]), 2)
        self.assertEqual(artifact["actual_forward_passes"], 28)
        self.assertEqual(len(model.records), 28)
        self.assertEqual(len(artifact["controls"]), 4)
        self.assertTrue(all(control["passed"] for control in artifact["controls"]))
        self.assertTrue(all(abs(control["raw_delta"]) <= 1e-8 for control in artifact["controls"]))
        for edge in artifact["edges"]:
            self.assertEqual(edge["receiver"], [2, 0, "q"])
            self.assertEqual([pair["pair_id"] for pair in edge["per_pair"]], ["one", "two"])
            for pair in edge["per_pair"]:
                self.assertAlmostEqual(pair["raw_delta"], pair["patched_margin"] - pair["clean_margin"])
                self.assertAlmostEqual(pair["normalizer"], pair["donor_margin"] - pair["clean_margin"])
                self.assertAlmostEqual(pair["normalized_effect"], pair["raw_delta"] / pair["normalizer"], places=6)
            self.assertEqual(edge["mean_normalized_effect"],
                             aggregate_pair_effects(edge["per_pair"]))
        self.assertFalse(probe.adapter._active)
        for module in model.modules():
            self.assertFalse(module._forward_hooks)
            self.assertFalse(module._forward_pre_hooks)

    def test_invalid_query_capability_or_direction_fails_before_any_forward(self):
        model = CircuitModel()
        probe = Prober(model)
        pair = {"id": "one", "clean_prompt": "abcd", "donor_prompt": "efgh",
                "clean_answer": "c", "donor_answer": "k"}
        with self.assertRaisesRegex(ValueError, "precede"):
            measure_receiver_paths(probe, [pair], CharacterTokenizer(), "cpu", [(2, 0)], [(1, 0, "q")])
        probe.spec.path_qkv.pop(2)
        with self.assertRaises(CapabilityError):
            measure_receiver_paths(probe, [pair], CharacterTokenizer(), "cpu", [(0, 0)], [(2, 0, "q")])
        self.assertFalse(model.records)


if __name__ == "__main__":
    unittest.main()