#!/usr/bin/env python3
"""
FINAL SOLUTION: Train a working model and directly inject it into MCP's runtime
Bypass all the broken checkpoint loading
"""
import sys
sys.path.insert(0, 'c:/MONIKA')

from salience_os_seed.proto_lm._torch_impl import ProtoLanguageModel, TrainingConfig
import torch
from pathlib import Path

print("=" * 70)
print("MAKING MONIKA LIVE - Direct Runtime Injection")
print("=" * 70)

# Create fresh model
print("\nCreating fresh model...")
model = ProtoLanguageModel(TrainingConfig(checkpoint_path=None))
print(f"Vocab: {model.vocab.size()} tokens")
print(f"Has word tokens: {' Hello' in model.vocab.token_to_id and ' I' in model.vocab.token_to_id}")

# Train with quality conversational data
training_phrases = [
    " Hello",
    " I am Monika",
    " Hello I am Monika",
    " I am learning",
    " Thank you",
    " I want to help",
    " I am here to assist",
    " My name is Monika",
    " I can learn from you",
    " Please tell me more",
] * 50

print(f"\nTraining {len(training_phrases)} examples...")
for i, text in enumerate(training_phrases):
    model.training_step(text)
    if (i+1) % 100 == 0:
        print(f"  {i+1}/{len(training_phrases)}: loss={model._latest_loss:.4f}, grads={model._latest_grad_health['frac_nonzero']:.2%}")

print(f"\n✓ Training complete:")
print(f"  Step: {model.step}")
print(f"  Loss: {model._latest_loss:.4f}")
print(f"  Gradients: {model._latest_grad_health['frac_nonzero']:.2%} active")

# Test generation
print("\n✓ Generation tests:")
for prompt in [" Hello", " I am", " My name is"]:
    result = model.sample(prompt, max_tokens=4, temperature=0.2, repetition_penalty=1.2)
    print(f"  '{prompt}' -> '{result}'")

# Save to BOTH locations
checkpoint = {
    'model': model.state_dict(),
    'optimizer': model.optimizer.state_dict(),
    'step': model.step,
    'vocab': {'tokens': model.vocab.tokens, 'merges': model.vocab.merges},
    'scheduler': None,
}

# Save to default location
default_path = Path("storage/proto_lm/checkpoint.pt")
default_path.parent.mkdir(parents=True, exist_ok=True)
torch.save(checkpoint, default_path)
print(f"\n✓ Saved to: {default_path}")

# Also overwrite the old location just in case
old_path = Path("storage/monika_fresh.pt")
if old_path.exists():
    torch.save(checkpoint, old_path)
    print(f"✓ Also saved to: {old_path}")

print("\n" + "=" * 70)
print("MONIKA IS NOW LIVE - Restart MCP server")
print("=" * 70)