import unittest
from types import SimpleNamespace

import torch
from torch import nn

from roboimi.vla.models.smolvla import NativeSmolVLAConfig
from roboimi.vla.models.smolvla.modeling import (
    VLAFlowMatching,
    make_att_2d_masks,
    pad_vector,
    resize_with_pad,
)


class FakeTokenizer:
    fake_image_token_id = 32000
    global_image_token_id = 32001


class FakeProcessor:
    tokenizer = FakeTokenizer()


class FakeVLMWithExpert(nn.Module):
    def __init__(self, vlm_hidden_size=8, expert_hidden_size=6, image_tokens=3, vocab_size=64):
        super().__init__()
        self.config = SimpleNamespace(text_config=SimpleNamespace(hidden_size=vlm_hidden_size))
        self.expert_hidden_size = expert_hidden_size
        self.image_tokens = image_tokens
        self.processor = FakeProcessor()
        self.vlm = SimpleNamespace(device=torch.device("cpu"))
        self.image_proj = nn.Linear(3, vlm_hidden_size)
        self.token_emb = nn.Embedding(vocab_size, vlm_hidden_size)
        self.suffix_proj = nn.Linear(expert_hidden_size, expert_hidden_size)

    def embed_image(self, image):
        # Deterministic lightweight image embedding: pool pixels, then repeat.
        pooled = image.mean(dim=(-1, -2)).to(dtype=torch.float32)
        return self.image_proj(pooled).unsqueeze(1).expand(-1, self.image_tokens, -1)

    def embed_language_tokens(self, tokens):
        return self.token_emb(tokens)

    def forward(self, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, fill_kv_cache):
        prefix_embs, suffix_embs = inputs_embeds
        prefix_out = prefix_embs if prefix_embs is not None else None
        suffix_out = self.suffix_proj(suffix_embs) if suffix_embs is not None else None
        cache = ("fake-cache",) if fill_kv_cache else past_key_values
        return (prefix_out, suffix_out), cache


class NativeSmolVLAModelingTest(unittest.TestCase):
    def test_config_rejects_action_steps_greater_than_chunk_size(self):
        with self.assertRaisesRegex(ValueError, "n_action_steps"):
            NativeSmolVLAConfig(chunk_size=2, n_action_steps=3)

    def test_pad_vector_pads_returns_original_for_equal_and_rejects_truncation(self):
        vector = torch.tensor([[1.0, 2.0, 3.0]])
        padded = pad_vector(vector, 5)
        self.assertEqual(tuple(padded.shape), (1, 5))
        torch.testing.assert_close(padded, torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0]]))

        same = pad_vector(vector, 3)
        self.assertIs(same, vector)

        with self.assertRaisesRegex(ValueError, "target dimension"):
            pad_vector(vector, 2)

    def test_resize_with_pad_returns_requested_spatial_size(self):
        img = torch.arange(2 * 3 * 4 * 8, dtype=torch.float32).reshape(2, 3, 4, 8)
        resized = resize_with_pad(img, width=10, height=10, pad_value=-1)
        self.assertEqual(tuple(resized.shape), (2, 3, 10, 10))

    def test_make_att_2d_masks_implements_prefix_lm_semantics(self):
        pad_masks = torch.tensor([[True, True, True, True, False]])
        # First two tokens are bidirectional prefix, later valid tokens are causal.
        att_masks = torch.tensor([[False, False, True, True, True]])
        mask = make_att_2d_masks(pad_masks, att_masks)
        expected = torch.tensor(
            [[
                [True, True, False, False, False],
                [True, True, False, False, False],
                [True, True, True, False, False],
                [True, True, True, True, False],
                [False, False, False, False, False],
            ]]
        )
        torch.testing.assert_close(mask, expected)

    def test_vla_flow_matching_forward_and_sample_actions_with_fake_vlm(self):
        torch.manual_seed(0)
        config = NativeSmolVLAConfig(
            chunk_size=4,
            n_action_steps=4,
            max_state_dim=5,
            max_action_dim=3,
            num_steps=2,
            prefix_length=8,
            add_image_special_tokens=False,
        )
        fake_vlm = FakeVLMWithExpert(vlm_hidden_size=8, expert_hidden_size=6)
        model = VLAFlowMatching(config, vlm_with_expert=fake_vlm)

        bsize = 2
        images = [torch.randn(bsize, 3, 6, 6)]
        img_masks = [torch.tensor([True, False])]
        lang_tokens = torch.tensor([[1, 2, 3], [4, 5, 0]])
        lang_masks = torch.tensor([[True, True, True], [True, True, False]])
        state = torch.randn(bsize, config.max_state_dim)
        actions = torch.randn(bsize, config.chunk_size, config.max_action_dim)

        losses = model(images, img_masks, lang_tokens, lang_masks, state, actions)
        self.assertEqual(tuple(losses.shape), (bsize, config.chunk_size, config.max_action_dim))

        sampled = model.sample_actions(images, img_masks, lang_tokens, lang_masks, state)
        self.assertEqual(tuple(sampled.shape), (bsize, config.chunk_size, config.max_action_dim))


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