#!/usr/bin/env python3
"""
Test script to verify dual-mode implementation for text conditioning pairs.
Tests that both original and new modes work correctly with random selection.

Created: 2025-09-29
"""

import unittest
import os
import random

import torch
import numpy as np

import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from data_utils import (
    make_text_conditioning_pair,
    make_conditioning_block,
    INSTRUCTION_TEMPLATES,
)
from modules.gpt import GPTConfig
from data_types import SamplingParams


class TestDualMode(unittest.TestCase):
    """Test dual-mode implementation for text conditioning pairs."""

    @classmethod
    def setUpClass(cls):
        """Set up test fixtures."""
        cls.cfg = GPTConfig(
            semantic_codebook_size=2048,
            semantic_vocab_size=2100,
        )
        cls.tokenizer_path = "/app2/suno/data/auk_v0/tokenizer_60k.json"
        cls.has_tokenizer = os.path.exists(cls.tokenizer_path)

    def test_gpt_config_has_text_desc_token(self):
        """Test that GPTConfig has semantic_cond_audio_token."""
        self.assertIsNotNone(self.cfg.semantic_cond_audio_token)
        self.assertEqual(self.cfg.semantic_cond_audio_token, self.cfg.semantic_codebook_size + 13)
        self.assertLess(self.cfg.semantic_cond_audio_token, self.cfg.semantic_vocab_size)

    def test_instruction_template_coverage(self):
        """Test that all expected block types have instruction templates."""
        expected_types = [
            "cover",
            "artist",
            "playlist",
            "overpaint",
            "underpaint",
            "vox",
            "stem_add",
            "stem_extract",
            "stem_remove",
            "sample",
            "prefix",
            "suffix",
        ]

        for block_type in expected_types:
            with self.subTest(block_type=block_type):
                self.assertIn(block_type, INSTRUCTION_TEMPLATES)
                self.assertGreaterEqual(len(INSTRUCTION_TEMPLATES[block_type]), 5)

    @unittest.skipIf(
        not os.path.exists("/app2/suno/data/auk_v0/tokenizer_60k.json"),
        "Tokenizer not available",
    )
    def test_new_mode_text_conditioning_pairs(self):
        """Test new mode - text conditioning pairs."""
        # Create mock audio data with semantic_cond_audio_token
        mock_audio = torch.tensor(
            [self.cfg.semantic_cond_audio_token] + [100, 101, 102, 103, 104],
            dtype=torch.int64,
        )

        test_cases = [
            ("cover", None),
            ("artist", None),
            ("stem", "extract Bass"),
            ("vox", None),
            ("sample", None),
        ]

        for block_type, stem_type in test_cases:
            with self.subTest(block_type=block_type):
                pair = make_text_conditioning_pair(
                    data=mock_audio,
                    block_type=block_type,
                    is_diffusion=False,
                    tokenizer_fp=self.tokenizer_path,
                    stem_type=stem_type,
                )

                # Verify it returns a tuple
                self.assertIsInstance(pair, tuple)
                self.assertEqual(len(pair), 2)

                text_block, content_block = pair

                # Verify text block has debug text
                self.assertIsNotNone(text_block.debug_text)
                self.assertIn("[", text_block.debug_text)

                # Verify content block has data
                self.assertGreater(len(content_block), 0)

    def test_original_mode_simple_conditioning_blocks(self):
        """Test original mode - simple conditioning blocks."""
        test_cases = [
            ("cover", self.cfg.semantic_cover_token),
            ("artist", self.cfg.semantic_artist_token),
            ("stem", self.cfg.semantic_stem_token),
            ("vox", self.cfg.semantic_vox_token),
        ]

        for block_type, token_name in test_cases:
            with self.subTest(block_type=block_type):
                mock_audio = torch.tensor([token_name] + [100, 101, 102, 103, 104], dtype=torch.int64)
                block = make_conditioning_block(mock_audio, block_type, is_diffusion=False)

                # Verify block structure
                self.assertIsNotNone(block)
                self.assertEqual(block.spec.name, block_type)
                self.assertGreater(len(block), 0)

    def test_sampling_params_has_prob_text_conditioning_pairs(self):
        """Test that SamplingParams has prob_text_conditioning_pairs."""
        # Default probability
        params_default = SamplingParams()
        self.assertEqual(params_default.prob_text_conditioning_pairs, 0.5)

        # Custom probability
        params_always = SamplingParams(prob_text_conditioning_pairs=1.0)
        self.assertEqual(params_always.prob_text_conditioning_pairs, 1.0)

        params_never = SamplingParams(prob_text_conditioning_pairs=0.0)
        self.assertEqual(params_never.prob_text_conditioning_pairs, 0.0)

    def test_random_selection_distribution(self):
        """Test random selection distribution."""
        random.seed(42)

        prob = 0.5
        num_samples = 1000
        new_mode_count = sum(1 for _ in range(num_samples) if random.random() < prob)

        # Should be approximately 50% with some tolerance
        expected = num_samples * prob
        tolerance = num_samples * 0.1  # 10% tolerance
        self.assertAlmostEqual(new_mode_count, expected, delta=tolerance)


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