#!/usr/bin/env python3
"""Train from ABSOLUTE ZERO - explicitly prevent any checkpoint loading."""

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))

from salience_os_seed.conversation.session import ConversationSession, ConversationConfig, IngestionConfig
from salience_os_seed.proto_lm.trainer import ProtoLanguageModel, TrainingConfig

print("="*80)
print("TRAINING FROM ABSOLUTE ZERO")
print("="*80)
print()

# Create a NEW checkpoint path that doesn't exist
new_checkpoint = Path("storage/proto_lm/zero_start.pt")

# DELETE it if it somehow exists
if new_checkpoint.exists():
    print(f"Deleting existing checkpoint: {new_checkpoint}")
    new_checkpoint.unlink()

# Create TrainingConfig with NO default checkpoint
training_cfg = TrainingConfig()
training_cfg.checkpoint_path = None  # Don't load anything!
training_cfg.device = "cuda"

# Create model WITHOUT loading anything
print("Creating fresh model (no checkpoint loading)...")
model = ProtoLanguageModel(training_cfg, learning_enabled=True)
print(f"  Step: {model.step} (should be 0)")
print(f"  Vocab: {model.vocab.size()} (should be 256 for base ASCII)")
print()

# NOW set the checkpoint path for saving
model.config.checkpoint_path = str(new_checkpoint)

# Train on simple examples
print("Training on simple examples...")
simple_examples = [
    "Hello",
    "Hi",
    "Hello there",
    "Hi there",
    "Thank you",
    "Thanks",
    "Yes",
    "No",
    "I am MONIKA",
    "My name is MONIKA",
] * 10  # Repeat 10 times

for idx, text in enumerate(simple_examples, 1):
    loss = model.training_step(text)
    if idx % 10 == 0:
        print(f"  Step {model.step}: loss={loss:.4f}, vocab={model.vocab.size()}")

print()
print(f"Training complete!")
print(f"  Final step: {model.step}")
print(f"  Final vocab: {model.vocab.size()}")
print()

# Save checkpoint
saved_path = model.save_checkpoint(str(new_checkpoint), reason="fresh_training")
print(f"Saved to: {saved_path}")
print()

# Test generation
print("="*80)
print("TESTING GENERATION")
print("="*80)
print()

test_prompts = ["Hello", "Hi", "Thank", "I am"]

for prompt in test_prompts:
    output = model.sample(
        prompt,
        max_tokens=10,
        temperature=0.7,
        repetition_penalty=1.8,
    )
    print(f"'{prompt}' → '{output}'")

print()
print("="*80)