#!/usr/bin/env python3
"""
Test script to see how the text tokenizer encodes various sequence formats.

This tests encoding of:
- Simple arrays like [1,2,3]
- Value sequences like [v1,v2,...,v100] with different lengths
- Mixed negative values like [-1,-1,100,60,51,41]
- Character mixing like [15,12,10,x,1,3,4]
"""

import sys
import os

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

from text_utils import tokenize_batch

# Test cases
test_sequences = []

# 1. Simple short sequences
test_sequences.append(("[1,2,3]", "Simple short sequence"))
test_sequences.append(("[10,20,30]", "Simple sequence with larger numbers"))

# 2. Value sequences with 10 different lengths (5 to 50 values)
for num_values in [5, 10, 15, 20, 25, 30, 35, 40, 45, 50]:
    values = [str(i % 101) for i in range(num_values)]  # Values 0-100
    seq_str = f"[{','.join(values)}]"
    test_sequences.append((seq_str, f"Sequence with {num_values} values (0-100)"))

# 3. Mixed negative values
test_sequences.append(("[-1,-1,100,60,51,41]", "Sequence with negative values"))
test_sequences.append(("[-1,50,75,-1,25,-1]", "Alternating negatives and positives"))

# 4. Character mixing
test_sequences.append(("[15,12,10,x,1,3,4]", "Sequence with character 'x'"))
test_sequences.append(("[10,abc,20,30]", "Sequence with string 'abc'"))
test_sequences.append(("[5,10,?,20,25]", "Sequence with '?' character"))

# 5. Full range test: -1 to 100
full_range_values = ["-1"] + [str(i) for i in range(101)]  # -1, 0, 1, 2, ..., 100
full_range_str = f"[{','.join(full_range_values)}]"
test_sequences.append((full_range_str, "Full range from -1 to 100 (102 values)"))

# 6. Double-digit tokenization tests
test_sequences.append(("[61]", "Single value 61 in array"))
test_sequences.append(("[71]", "Single value 71 in array"))
test_sequences.append(("[61,71]", "Two values 61 and 71"))
test_sequences.append(("[6]", "Single digit 6 in array"))
test_sequences.append(("[1]", "Single digit 1 in array"))
test_sequences.append(("[7]", "Single digit 7 in array"))

# Tokenize and display results
print("=" * 100)
print("Text Tokenizer Encoding Tests")
print("=" * 100)
print()

tokenizer_fp = "/app2/suno/data/auk_v0/tokenizer_60k.json"

for seq_str, description in test_sequences:
    tokens = tokenize_batch([seq_str], tokenizer_fp=tokenizer_fp, pad_token_id=60001)[0]

    print(f"Description: {description}")
    print(f"Input:       {seq_str}")
    print(f"Tokens:      {tokens.tolist()}")
    print(f"Num tokens:  {len(tokens)}")
    print(f"Tokens/char: {len(tokens) / len(seq_str):.2f}")
    print()

print("=" * 100)
print("Summary:")
print("- Short sequences like [1,2,3] are compact")
print("- Longer sequences scale linearly with number of values")
print("- Negative values and special characters affect tokenization")
print("- Token count correlates with string length, not value count")
print("=" * 100)
print()

# Detailed assertion test for full range -1 to 100
print("=" * 100)
print("Detailed Token Analysis: Full Range [-1, 0, 1, ..., 100]")
print("=" * 100)
print()

# Tokenize the full range sequence
full_range_tokens = tokenize_batch([full_range_str], tokenizer_fp=tokenizer_fp, pad_token_id=60001)[
    0
].tolist()

print(f"Input: {full_range_str[:100]}... (truncated)")
print(f"Total tokens: {len(full_range_tokens)}")
print()

# Expected structure:
# [ (1 token) + -1 (2 tokens) + , (1 token) + 0 (1 token) + , (1 token) + ... + 100 (1 token) + ] (1 token)
# = 1 + 2 + (101 commas * 1) + (101 values: 0-100 each 1 token) + 1
# = 1 + 2 + 101 + 101 + 1 = 206 tokens expected

# Test individual value tokenization
print("Testing individual value tokenization:")
print()

# Test -1
minus_one_tokens = tokenize_batch(["-1"], tokenizer_fp=tokenizer_fp, pad_token_id=60001)[0].tolist()
print(f"  -1 encodes as: {minus_one_tokens} ({len(minus_one_tokens)} tokens)")
assert len(minus_one_tokens) == 2, f"-1 should use 2 tokens, got {len(minus_one_tokens)}"

# Test 0-100 (should each be single token)
failed_values = []
for i in range(101):
    value_tokens = tokenize_batch([str(i)], tokenizer_fp=tokenizer_fp, pad_token_id=60001)[0].tolist()
    if len(value_tokens) != 1:
        failed_values.append((i, len(value_tokens), value_tokens))

if failed_values:
    print()
    print(f"  ❌ FAILED: {len(failed_values)} values did not encode as single token:")
    for val, num_tokens, tokens in failed_values[:10]:  # Show first 10 failures
        print(f"     Value {val}: {num_tokens} tokens {tokens}")
    if len(failed_values) > 10:
        print(f"     ... and {len(failed_values) - 10} more")
else:
    print(f"  ✓ All values 0-100 encode as single token")

# Test comma
comma_tokens = tokenize_batch([","], tokenizer_fp=tokenizer_fp, pad_token_id=60001)[0].tolist()
print(f"  Comma ',' encodes as: {comma_tokens} ({len(comma_tokens)} token)")
assert len(comma_tokens) == 1, f"Comma should use 1 token, got {len(comma_tokens)}"

print()
print("Expected token count calculation:")
print(f"  [ bracket: 1 token")
print(f"  -1 value: 2 tokens")
print(f"  101 commas: 101 tokens")
print(f"  101 values (0-100): 101 tokens")
print(f"  ] bracket: 1 token")
print(f"  TOTAL EXPECTED: 206 tokens")
print()
print(f"Actual token count: {len(full_range_tokens)}")

if len(full_range_tokens) == 206:
    print("✓ Token count matches expected!")
else:
    print(f"❌ Token count mismatch! Expected 206, got {len(full_range_tokens)}")
    print(f"   Difference: {len(full_range_tokens) - 206}")

print("=" * 100)