"""Small CPU checks for shared-prefix inference; no checkpoint download needed."""
import unittest
from unittest.mock import patch

import numpy as np
import torch
from transformers import (Qwen3_5Config, Qwen3_5ForSequenceClassification,
                          Qwen3_5TextConfig, Qwen3_5VisionConfig)

from modeling_openjev import OpenJevCrossEncoder


class CharacterTokenizer:
    pad_token_id = 0

    def __call__(self, texts, truncation=True, max_length=4096, padding=False, return_tensors=None):
        if isinstance(texts, str):
            return {"input_ids": [ord(c) + 1 for c in texts][:max_length]}
        rows = [[ord(c) + 1 for c in text][:max_length] for text in texts]
        width = max(map(len, rows))
        ids = [row + [0] * (width - len(row)) for row in rows]
        mask = [[1] * len(row) + [0] * (width - len(row)) for row in rows]
        return {"input_ids": torch.tensor(ids), "attention_mask": torch.tensor(mask)}


class SharedPrefixTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        torch.manual_seed(7)
        config = Qwen3_5TextConfig(
            vocab_size=256, hidden_size=32, intermediate_size=64,
            num_hidden_layers=4, num_attention_heads=4, num_key_value_heads=2,
            head_dim=8, linear_num_key_heads=2, linear_num_value_heads=4,
            linear_key_head_dim=8, linear_value_head_dim=8,
            layer_types=["linear_attention", "full_attention"] * 2,
            rope_parameters={"rope_type": "default", "rope_theta": 10000,
                             "partial_rotary_factor": 0.25, "mrope_section": [1, 0, 0]},
            pad_token_id=0, num_labels=3, use_cache=False,
        )
        vision = Qwen3_5VisionConfig(depth=1, hidden_size=32, intermediate_size=64,
                                    num_heads=4, out_hidden_size=32)
        model = Qwen3_5ForSequenceClassification(
            Qwen3_5Config(text_config=config, vision_config=vision, num_labels=3,
                           pad_token_id=0)).eval()
        cls.jev = OpenJevCrossEncoder.__new__(OpenJevCrossEncoder)
        cls.jev.model = model
        cls.jev.backbone = model.model
        cls.jev.tok = CharacterTokenizer()
        cls.jev.template = "Premise: {premise}\nHypothesis: {hypothesis}"
        cls.jev.device = "cpu"
        cls.jev.bs = 32
        cls.jev.max_len = 128

    def test_shared_matches_independent_and_isolates_branches(self):
        premise = "A cat sleeps."
        hypotheses = ["A cat rests.", "The dog runs.", "A cat."]
        expected = self.jev.latents([(premise, h) for h in hypotheses])
        actual = self.jev.latents_hypotheses(premise, hypotheses)
        np.testing.assert_allclose(actual, expected, atol=2e-4, rtol=2e-4)
        expected_probs = self.jev.predict([(premise, h) for h in hypotheses])
        actual_probs = self.jev.predict_hypotheses(premise, hypotheses)
        np.testing.assert_allclose(actual_probs, expected_probs, atol=2e-4, rtol=2e-4)

        changed = hypotheses.copy()
        changed[1] = "Birds fly around today."
        other = self.jev.latents_hypotheses(premise, changed)
        np.testing.assert_allclose(other[[0, 2]], actual[[0, 2]], atol=2e-4, rtol=2e-4)

    def test_small_option_sets_use_pairwise_batch(self):
        for hypotheses in (["B"], ["B", "C"]):
            with self.subTest(count=len(hypotheses)):
                with patch.object(self.jev, "_pooled", wraps=self.jev._pooled) as pooled:
                    actual = self.jev.predict_hypotheses("A", hypotheses)
                    self.assertEqual(actual.shape, (len(hypotheses), 3))
                    pooled.assert_called_once()

    def test_empty(self):
        with self.assertRaises(ValueError):
            self.jev.predict_hypotheses("A", [])


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