#!/usr/bin/env python3
"""Train from COMPLETE scratch - no existing checkpoint."""

from pathlib import Path
from types import SimpleNamespace
import sys

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

from salience_os_seed.training.run_corpus import train_corpus

def main():
    corpus_path = Path("data/local_benchmarks/synthetic_baseline_corpus.txt")
    checkpoint_path = Path("storage/proto_lm/fresh_synthetic.pt")
    
    # Ensure corpus exists
    if not corpus_path.exists():
        raise FileNotFoundError(f"Corpus not found at {corpus_path}")
    
    # DELETE old checkpoint if it exists
    if checkpoint_path.exists():
        print(f"Deleting old checkpoint: {checkpoint_path}")
        checkpoint_path.unlink()
    
    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
    
    print("="*80)
    print("TRAINING FROM COMPLETE SCRATCH")
    print("="*80)
    print(f"Corpus: {corpus_path}")
    print(f"Checkpoint: {checkpoint_path}")
    print(f"Resume: False (fresh start)")
    print()
    
    args = SimpleNamespace(
        corpus=corpus_path,
        epochs=20,  # More epochs
        chunk_size=2048,
        shuffle_buffer=64,
        seed=13,
        patience=5,  # More patience
        min_delta=0.05,
        log_every=100,  # Log more frequently
        checkpoint_path=checkpoint_path,
        checkpoint_interval=5000,
        resume=False,  # DON'T resume - start fresh!
        salience_filter=False,  # Disable filter initially
        min_uncertainty=0.0,
        min_novelty=0.0,
        max_drag=1.0,
    )
    
    train_corpus(args)
    
    print()
    print("="*80)
    print("Training complete! Testing generation...")
    print("="*80)
    
    # Test immediately
    from salience_os_seed.proto_lm.trainer import ProtoLanguageModel, TrainingConfig
    
    cfg = TrainingConfig()
    cfg.checkpoint_path = str(checkpoint_path)
    model = ProtoLanguageModel(cfg, learning_enabled=False)
    model.load_checkpoint()
    
    print(f"\nFinal state:")
    print(f"  Step: {model.step}")
    print(f"  Vocab: {model.vocab.size()}")
    print()
    
    test_prompts = ["Hello", "What", "Hi"]
    for prompt in test_prompts:
        output = model.sample(prompt, max_tokens=15, temperature=0.7, repetition_penalty=1.8)
        print(f"{prompt} → {output}")

if __name__ == "__main__":
    main()