"""CPU correctness checks using a tiny random Qwen and deterministic tokenizer."""

import unittest
from unittest.mock import patch

import torch
from transformers import Qwen2Config, Qwen2ForCausalLM

from decision_model import DecisionScorer


class CharacterTokenizer:
    pad_token_id = 0
    eos_token_id = 2

    def encode(self, text, add_special_tokens=False):
        assert add_special_tokens is False
        if text in (" yes", " no"):
            return [1 if text == " yes" else 2]
        return [ord(character) + 3 for character in text]


class ScorerTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        torch.set_num_threads(1)

    def setUp(self):
        torch.manual_seed(42)
        config = Qwen2Config(
            vocab_size=256, hidden_size=32, intermediate_size=48,
            num_hidden_layers=3, num_attention_heads=4, num_key_value_heads=2,
            max_position_embeddings=2048, attention_dropout=0.0,
        )
        config._attn_implementation = "sdpa"
        with (
            patch("decision_model.AutoTokenizer.from_pretrained", return_value=CharacterTokenizer()),
            patch("decision_model.AutoModelForCausalLM.from_pretrained", return_value=Qwen2ForCausalLM(config)),
        ):
            self.scorer = DecisionScorer("unused", train_layers=1, device="cpu").eval()
        self.state = "The customer needs a refund."
        self.questions = [
            {"question": "Which team?", "choices": ["sales", "billing support", "tech"]},
            {"question": "Is money involved?", "choices": ["yes", "no"]},
        ]
        self.examples = [dict(state=self.state, **question) for question in self.questions]

    def randomize_head(self):
        # A zero head would make cache/permutation comparisons vacuous.
        with torch.no_grad():
            self.scorer.head.weight.normal_(std=0.1)

    def test_zero_head_and_trainable_layers(self):
        self.assertEqual(torch.count_nonzero(self.scorer.head.weight).item(), 0)
        for index, layer in enumerate(self.scorer.lm.model.layers):
            self.assertTrue(all(p.requires_grad == (index == 2) for p in layer.parameters()))
        self.randomize_head()
        scores = self.scorer.score_examples(self.examples)
        sum(score.square().sum() for score in scores).backward()
        self.assertGreater(self.scorer.head.weight.grad.abs().sum().item(), 0)
        grads = [p.grad for p in self.scorer.lm.model.layers[-1].parameters() if p.grad is not None]
        self.assertGreater(sum(grad.abs().sum().item() for grad in grads), 0)
        self.assertIsNone(self.scorer.lm.model.embed_tokens.weight.grad)

    def test_shared_parity_chunking_permutation_and_isolation(self):
        self.randomize_head()
        with torch.inference_mode():
            full = self.scorer.score_examples(self.examples)
        for batch_size in (1, 2, 4, 16):
            shared = self.scorer.scores_shared(self.state, self.questions, batch_size)
            for expected, actual in zip(full, shared):
                torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
        self.assertEqual(self.scorer.last_shared_timings["branches"], 5)
        self.assertGreater(self.scorer.last_shared_timings["cache_bytes"], 0)
        permuted = dict(self.questions[0], choices=list(reversed(self.questions[0]["choices"])))
        actual = self.scorer.scores_shared(self.state, [permuted], 2)[0]
        torch.testing.assert_close(actual.flip(0), full[0], rtol=1e-5, atol=1e-5)
        altered = dict(self.questions[1], question="A completely different question?")
        actual = self.scorer.scores_shared(self.state, [self.questions[0], altered], 3)[0]
        torch.testing.assert_close(actual, full[0], rtol=1e-5, atol=1e-5)

    def test_token_baseline_matches_two_logits(self):
        with torch.inference_mode():
            hidden, counts = self.scorer._full_hidden(self.examples)
            logits = self.scorer.lm.lm_head(hidden)
            expected = (logits[:, 1] - logits[:, 2]).split(counts)
            actual = self.scorer.scores_token_baseline(self.examples)
        for want, got in zip(expected, actual):
            torch.testing.assert_close(got, want, rtol=1e-5, atol=1e-5)

    def test_invalid_input_and_length(self):
        with self.assertRaises(ValueError):
            self.scorer.score_examples([])
        with self.assertRaises(ValueError):
            self.scorer.scores_shared("", self.questions)
        with self.assertRaises(ValueError):
            self.scorer.scores_shared(self.state, [dict(question="Q?", choices=["a", " a "])])
        with self.assertRaises(ValueError):
            self.scorer.scores_shared(self.state, self.questions, 0)
        self.scorer.max_tokens = 10
        with self.assertRaisesRegex(ValueError, "never silently truncated"):
            self.scorer.score_examples(self.examples)
        with self.assertRaisesRegex(ValueError, "never silently truncated"):
            self.scorer.scores_shared(self.state, self.questions)


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