test_neuroname.py

"""
NeuroName Test Script - Validates all components work correctly.

Run: python test_neuroname.py

This script tests:
1. Model instantiation and forward pass
2. Phonotactic scoring
3. Sound symbolism engine
4. Morphological composition
5. Latent space operations
6. Full generation pipeline
7. Training step
"""

import sys
import torch
import numpy as np

def test_char_vocab():
    """Test character vocabulary encoding/decoding."""
    from neuroname.model import CharVocab
    
    vocab = CharVocab()
    assert len(vocab) == 44, f"Expected 44 tokens, got {len(vocab)}"
    
    # Test encode/decode roundtrip
    text = "spotify"
    encoded = vocab.encode(text)
    decoded = vocab.decode(encoded)
    assert decoded == text, f"Roundtrip failed: '{text}' → {encoded} → '{decoded}'"
    
    # Test batch operations
    texts = ["google", "tesla", "nexus"]
    batch = vocab.batch_encode(texts)
    assert batch.shape == (3, 32), f"Expected (3, 32), got {batch.shape}"
    
    decoded_batch = vocab.batch_decode(batch)
    assert decoded_batch == texts, f"Batch roundtrip failed"
    
    print("  ✓ CharVocab: encode/decode works correctly")


def test_model_forward():
    """Test model forward pass with random data."""
    from neuroname.model import NeuroNameModel
    
    model = NeuroNameModel()
    model.eval()
    
    batch_size = 4
    max_len = 32
    max_hints = 8
    
    # Create dummy batch
    char_ids = torch.randint(0, 44, (batch_size, max_len))
    hint_ids = torch.randint(0, 100, (batch_size, max_hints))
    target_length = torch.rand(batch_size, 1)
    style = torch.randint(0, 8, (batch_size,))
    language_feel = torch.randint(0, 8, (batch_size,))
    energy = torch.randint(0, 3, (batch_size,))
    char_padding = torch.zeros(batch_size, max_len, dtype=torch.bool)
    hint_padding = torch.zeros(batch_size, max_hints, dtype=torch.bool)
    
    # Forward pass
    outputs = model(
        char_ids=char_ids,
        hint_ids=hint_ids,
        target_length=target_length,
        style=style,
        language_feel=language_feel,
        energy=energy,
        char_padding_mask=char_padding,
        hint_padding_mask=hint_padding,
    )
    
    # Check outputs
    assert "recon_loss" in outputs
    assert "kl_loss" in outputs
    assert "style_loss" in outputs
    assert "lang_loss" in outputs
    assert "logits" in outputs
    assert "z" in outputs
    assert outputs["z"].shape == (batch_size, 128), f"z shape: {outputs['z'].shape}"
    assert outputs["logits"].shape == (batch_size, max_len - 1, 44)
    
    # Check loss is finite
    total_loss = model.total_loss(outputs)
    assert torch.isfinite(total_loss), f"Loss is not finite: {total_loss}"
    
    # Check backward pass
    total_loss.backward()
    
    print(f"  ✓ Model forward pass: loss={total_loss.item():.4f}")
    print(f"    Parameters: {sum(p.numel() for p in model.parameters()):,}")


def test_generation():
    """Test autoregressive generation."""
    from neuroname.model import NeuroNameModel
    
    model = NeuroNameModel()
    model.eval()
    
    # Prepare inputs
    hint_ids = torch.randint(0, 100, (1, 8))
    hint_mask = torch.zeros(1, 8, dtype=torch.bool)
    target_length = torch.tensor([[0.3]])
    style = torch.tensor([3])  # techy
    lang = torch.tensor([1])   # latin
    energy = torch.tensor([2]) # energetic
    
    # Generate
    with torch.no_grad():
        generated = model.generate_from_prior(
            hint_ids=hint_ids,
            target_length=target_length,
            style=style,
            language_feel=lang,
            energy=energy,
            hint_padding_mask=hint_mask,
            temperature=0.9,
            num_samples=3,
        )
    
    assert generated.shape[0] == 3, f"Expected 3 samples, got {generated.shape[0]}"
    
    # Decode
    names = model.char_vocab.batch_decode(generated)
    assert len(names) == 3
    
    print(f"  ✓ Generation works: {names}")


def test_phonotactics():
    """Test phonotactic scoring."""
    from neuroname.phonotactics import PhonotacticScorer
    
    scorer = PhonotacticScorer()
    
    # Good names should score higher than bad names
    good_names = ["spotify", "nexura", "velocix", "tervon"]
    bad_names = ["xqzwp", "bkftl", "mmmmm", "crtstk"]
    
    good_scores = [scorer.score(n)["overall"] for n in good_names]
    bad_scores = [scorer.score(n)["overall"] for n in bad_names]
    
    avg_good = sum(good_scores) / len(good_scores)
    avg_bad = sum(bad_scores) / len(bad_scores)
    
    assert avg_good > avg_bad, f"Good ({avg_good:.3f}) should be > bad ({avg_bad:.3f})"
    
    print(f"  ✓ Phonotactics: good={avg_good:.3f} > bad={avg_bad:.3f}")


def test_sound_symbolism():
    """Test sound symbolism engine."""
    from neuroname.phonotactics import SoundSymbolismEngine
    
    engine = SoundSymbolismEngine()
    
    # Tesla should match "techy" better than "organic"
    tesla_techy = engine.style_match_score("tesla", "techy")
    tesla_organic = engine.style_match_score("tesla", "organic")
    
    # Bloom should match "organic" better than "techy"
    bloom_organic = engine.style_match_score("bloom", "organic")
    bloom_techy = engine.style_match_score("bloom", "techy")
    
    print(f"  ✓ Sound symbolism:")
    print(f"    Tesla→techy={tesla_techy:.3f}, Tesla→organic={tesla_organic:.3f}")
    print(f"    Bloom→organic={bloom_organic:.3f}, Bloom→techy={bloom_techy:.3f}")


def test_morphology():
    """Test morphological composition."""
    from neuroname.morphology import MorphologicalComposer
    
    composer = MorphologicalComposer(seed=42)
    
    results = composer.compose(
        source_words=["velocity", "nexus", "quantum"],
        style="techy",
        num_results=10,
        min_length=4,
        max_length=12,
    )
    
    assert len(results) > 0, "Should generate at least some results"
    
    for r in results:
        assert 4 <= len(r.name) <= 12, f"Name '{r.name}' outside length range"
        assert r.operation in ["blend", "clip_extend", "affix", "sound_shift", 
                               "back_form", "cross_blend"]
    
    print(f"  ✓ Morphology: generated {len(results)} names")
    for r in results[:5]:
        print(f"    • {r.name:12s} [{r.operation}]")


def test_latent_ops():
    """Test latent space operations."""
    from neuroname.latent_ops import LatentSpaceController, NoveltyFilter, LatentArithmetic
    from neuroname.model import NeuroNameModel
    
    model = NeuroNameModel()
    
    controller = LatentSpaceController(
        model.style_classifier,
        model.lang_classifier,
    )
    
    # Test energy computation
    z = torch.randn(4, 128)
    energy = controller.compute_energy(z, {"style": 3, "language": 1})
    assert energy.shape == (4,), f"Energy shape: {energy.shape}"
    assert torch.all(torch.isfinite(energy))
    
    # Test interpolation
    z1 = torch.randn(1, 128)
    z2 = torch.randn(1, 128)
    interp = controller.interpolate(z1, z2, num_steps=5, method="slerp")
    assert interp.shape == (1, 5, 128)
    
    # Test novelty filter
    novelty = NoveltyFilter()
    assert novelty.is_novel("velocix")[0] == True
    assert novelty.is_novel("google")[0] == False
    assert novelty.is_novel("googlx")[0] == False  # Too similar
    
    # Test latent arithmetic
    result = LatentArithmetic.concept_blend([z1, z2], weights=[0.5, 0.5])
    assert result.shape == (1, 128)
    
    print("  ✓ Latent ops: energy, interpolation, novelty, arithmetic all work")


def test_data_pipeline():
    """Test data loading pipeline."""
    from neuroname.model import CharVocab
    from neuroname.data import (
        SemanticVocab, NameDataset, get_curated_brand_names,
        get_synthetic_training_data, create_dataloader,
    )
    
    char_vocab = CharVocab()
    semantic_vocab = SemanticVocab()
    
    # Test curated data
    curated = get_curated_brand_names()
    assert len(curated) > 50, f"Expected >50 curated names, got {len(curated)}"
    
    # Test synthetic generation
    synthetic = get_synthetic_training_data(num_samples=100, seed=42)
    assert len(synthetic) == 100
    
    # Test dataloader
    loader = create_dataloader(synthetic, char_vocab, semantic_vocab, batch_size=16)
    batch = next(iter(loader))
    
    assert batch["char_ids"].shape[0] == 16
    assert batch["char_ids"].shape[1] == 32
    assert batch["hint_ids"].shape == (16, 8)
    assert batch["style"].shape == (16,)
    
    print(f"  ✓ Data pipeline: {len(curated)} curated + {len(synthetic)} synthetic samples")


def test_full_generator():
    """Test the full generator pipeline."""
    from neuroname.generator import NeuroNameGenerator
    
    generator = NeuroNameGenerator(device="cpu")
    
    # Generate with morphological only (works without training)
    results = generator.generate(
        semantic_hints=["speed", "technology"],
        style="techy",
        language_feel="latin",
        energy="energetic",
        length_range=(4, 10),
        num_names=5,
        use_guided_sampling=False,
    )
    
    assert len(results) > 0, "Should generate at least some names"
    
    for r in results:
        assert 4 <= len(r.name) <= 10
        assert 0 <= r.phonotactic_score <= 1
        assert 0 <= r.style_match_score <= 1
    
    print(f"  ✓ Full generator: produced {len(results)} names")
    for r in results:
        print(f"    • {r.name:12s} (phon={r.phonotactic_score:.2f}, style={r.style_match_score:.2f})")


def test_training_step():
    """Test one training step to ensure gradients flow."""
    from neuroname.model import NeuroNameModel, CharVocab
    from neuroname.data import SemanticVocab, get_synthetic_training_data, create_dataloader
    
    model = NeuroNameModel()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    
    char_vocab = CharVocab()
    semantic_vocab = SemanticVocab()
    data = get_synthetic_training_data(num_samples=32, seed=42)
    loader = create_dataloader(data, char_vocab, semantic_vocab, batch_size=16)
    
    model.train()
    batch = next(iter(loader))
    
    outputs = model(
        char_ids=batch["char_ids"],
        hint_ids=batch["hint_ids"],
        target_length=batch["target_length"],
        style=batch["style"],
        language_feel=batch["language_feel"],
        energy=batch["energy"],
        char_padding_mask=batch["char_padding_mask"],
        hint_padding_mask=batch["hint_padding_mask"],
    )
    
    loss = model.total_loss(outputs)
    loss.backward()
    
    # Check gradients exist
    grad_norms = []
    for name, param in model.named_parameters():
        if param.grad is not None:
            grad_norms.append(param.grad.norm().item())
    
    assert len(grad_norms) > 0, "No gradients computed!"
    assert sum(grad_norms) > 0, "All gradients are zero!"
    
    optimizer.step()
    
    print(f"  ✓ Training step: loss={loss.item():.4f}, "
          f"grad_norm_avg={sum(grad_norms)/len(grad_norms):.6f}")


def main():
    print("=" * 60)
    print("NeuroName Test Suite")
    print("=" * 60)
    print()
    
    tests = [
        ("Character Vocabulary", test_char_vocab),
        ("Model Forward Pass", test_model_forward),
        ("Autoregressive Generation", test_generation),
        ("Phonotactic Scoring", test_phonotactics),
        ("Sound Symbolism", test_sound_symbolism),
        ("Morphological Composition", test_morphology),
        ("Latent Space Operations", test_latent_ops),
        ("Data Pipeline", test_data_pipeline),
        ("Full Generator", test_full_generator),
        ("Training Step", test_training_step),
    ]
    
    passed = 0
    failed = 0
    
    for name, test_fn in tests:
        print(f"\n[{name}]")
        try:
            test_fn()
            passed += 1
        except Exception as e:
            print(f"  ✗ FAILED: {e}")
            import traceback
            traceback.print_exc()
            failed += 1
    
    print("\n" + "=" * 60)
    print(f"Results: {passed} passed, {failed} failed out of {len(tests)} tests")
    print("=" * 60)
    
    if failed > 0:
        sys.exit(1)
    else:
        print("\n🎉 All tests passed! NeuroName is working correctly.")
        sys.exit(0)


if __name__ == "__main__":
    main()