#!/usr/bin/env python3
"""FINAL fresh start with proper vocab and training."""
import sys
import torch
sys.path.insert(0, 'c:/MONIKA')
from salience_os_seed.proto_lm._torch_impl import ProtoLanguageModel, TrainingConfig

print("Creating TRULY fresh model...")
m = ProtoLanguageModel(TrainingConfig(checkpoint_path=None))

print(f"Vocab size: {m.vocab.size()}")
print(f"Has 'G': {'G' in m.vocab.token_to_id}")
print(f"Has ' Hello': {' Hello' in m.vocab.token_to_id}")
print(f"Has ' I': {' I' in m.vocab.token_to_id}")

# Train with SHORT, simple phrases
training = [
    "Hello",
    " I am Monika",
    " Hello I am Monika",
    " Thank you",
    " I am learning",
] * 100

print(f"\nTraining {len(training)} examples...")
for i, text in enumerate(training):
    m.training_step(text)
    if (i+1) % 100 == 0:
        print(f"  Step {i+1}: loss={m._latest_loss:.4f}")

print(f"\nFinal: step={m.step}, loss={m._latest_loss:.4f}, grad_health={m._latest_grad_health}")

# Test
print("\nGeneration:")
result = m.sample(" Hello I", max_tokens=3, temperature=0.1, repetition_penalty=1.0)
print(f"  ' Hello I' -> '{result}'")

# Save
torch.save({
    'model': m.state_dict(),
    'optimizer': m.optimizer.state_dict(),
    'step': m.step,
    'vocab': {'tokens': m.vocab.tokens, 'merges': m.vocab.merges},
    'scheduler': None
}, 'storage/monika_FINAL_FRESH.pt')
print("\n✓ Saved to storage/monika_FINAL_FRESH.pt")