import unittest
from types import SimpleNamespace
from unittest.mock import patch

import torch

from scoring.functions.peptiverse_binding import (
    PeptiVerseBindingAffinity,
    PeptiVersePooledAffinityModel,
)


class DummyTokenizer:
    pad_token_id = 0
    cls_token_id = 1
    eos_token_id = 2
    sep_token_id = None
    bos_token_id = None
    mask_token_id = None

    def __call__(self, texts, **kwargs):
        del texts, kwargs
        return {
            "input_ids": torch.tensor([[1, 3, 4, 2, 0]]),
            "attention_mask": torch.tensor([[1, 1, 1, 1, 0]]),
        }


class DummyEncoder(torch.nn.Module):
    def forward(self, input_ids, attention_mask):
        del attention_mask
        hidden = input_ids.float().unsqueeze(-1).repeat(1, 1, 2)
        return SimpleNamespace(last_hidden_state=hidden)


class DummyAffinityHead(torch.nn.Module):
    def forward(self, target, binder):
        del target
        return binder.sum(dim=-1), torch.zeros(len(binder), 3)


class PeptiVerseBindingTests(unittest.TestCase):
    def test_affinity_head_shapes(self):
        model = PeptiVersePooledAffinityModel(
            target_dim=8,
            binder_dim=6,
            hidden_dim=12,
            n_heads=3,
            n_layers=2,
            dropout=0.0,
        ).eval()
        affinity, classes = model(torch.randn(4, 8), torch.randn(4, 6))
        self.assertEqual(tuple(affinity.shape), (4,))
        self.assertEqual(tuple(classes.shape), (4, 3))

    def test_pool_excludes_special_tokens(self):
        predictor = object.__new__(PeptiVerseBindingAffinity)
        predictor.device = torch.device("cpu")
        pooled = predictor._pool(
            ["unused"], DummyTokenizer(), DummyEncoder(), max_length=8
        )
        expected = torch.tensor([[3.5, 3.5]])
        self.assertTrue(torch.equal(pooled, expected))

    def test_factory_keeps_original_as_default(self):
        from scoring.functions import binding

        original = object()
        with patch.object(
            binding, "MultiTargetBindingAffinity", return_value=original
        ) as constructor:
            result = binding.create_multi_target_affinity_predictor(
                tokenizer=object(), base_path="/tmp", device="cpu"
            )
        self.assertIs(result, original)
        constructor.assert_called_once()

    def test_factory_selects_peptiverse(self):
        from scoring.functions import binding

        peptiverse = object()
        with patch.object(
            binding, "PeptiVerseBindingAffinity", return_value=peptiverse
        ) as constructor:
            result = binding.create_multi_target_affinity_predictor(
                backend="peptiverse",
                device="cpu",
                peptiverse_checkpoint="model.pt",
            )
        self.assertIs(result, peptiverse)
        constructor.assert_called_once_with(
            device="cpu",
            checkpoint_path="model.pt",
            repo_id="ChatterjeeLab/PeptiVerse",
            revision=None,
            cache_dir=None,
            local_files_only=False,
            batch_size=32,
        )

    def test_forward_batches_binder_smiles(self):
        predictor = object.__new__(PeptiVerseBindingAffinity)
        predictor.batch_size = 2
        predictor.binder_tokenizer = object()
        predictor.binder_encoder = object()
        predictor.max_smiles_length = 16
        predictor.model = DummyAffinityHead()
        predictor.get_protein_embedding = lambda _: torch.zeros(1, 3)

        batch_sizes = []

        def fake_pool(texts, tokenizer, encoder, max_length):
            del tokenizer, encoder, max_length
            batch_sizes.append(len(texts))
            return torch.ones(len(texts), 2)

        predictor._pool = fake_pool
        scores = predictor.forward(["a", "b", "c", "d", "e"], "TARGET")
        self.assertEqual(batch_sizes, [2, 2, 1])
        self.assertEqual(scores, [2.0] * 5)


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