#!/usr/bin/env python3
"""Final test - completely fresh, no merges."""
import sys
import torch
sys.path.insert(0, 'c:/MONIKA')
from salience_os_seed.proto_lm._torch_impl import ProtoLanguageModel, TrainingConfig
from salience_os_seed.proto_lm.vocab import Vocabulary

# Create TRULY fresh vocab with no merges
fresh_vocab = Vocabulary()
assert len(fresh_vocab.merges) == 0, "Vocab should have no merges"

m = ProtoLanguageModel(TrainingConfig(checkpoint_path=None))
m.vocab = fresh_vocab
m.vocab._refresh_index()
m._ensure_capacity(m.vocab.size())

print(f"Fresh vocab size: {m.vocab.size()}, merges: {len(m.vocab.merges)}")

# Train
phrase = "Hello I am Monika"
print(f"Training 100x on: '{phrase}'")
for i in range(100):
    m.training_step(phrase)
    if i == 99:
        print(f"Final loss: {m._latest_loss:.8f}")

# Test character-level encoding
print(f"\nCharacter tokens for 'Hello I am M':")
ids = m.encode("Hello I am M", mutate=False)
print(f"IDs: {ids}")
print(f"Tokens: {[repr(m.vocab.tokens[i]) for i in ids]}")

# Generate
print(f"\nGeneration test (temp=0.01, no rep penalty):")
result = m.sample("Hello I am M", max_tokens=6, temperature=0.01, repetition_penalty=1.0)
print(f"Output: '{result}'")

# Check if it predicts the right next chars
with torch.no_grad():
    m.eval()
    context = torch.tensor(ids, device=m.device).unsqueeze(0)
    logits = m._forward_logits(context)
    
    for i in range(6):
        next_logits = logits[0, -1, :]
        next_id = next_logits.argmax().item()
        next_char = m.vocab.tokens[next_id]
        print(f"  Step {i+1}: Predicted '{next_char}' (ID {next_id})")
        
        # Add to context
        next_tensor = torch.tensor([[next_id]], device=m.device)
        context = torch.cat([context, next_tensor], dim=1)
        logits = m._forward_logits(context)