#!/usr/bin/env python3
"""Train the zero-start model with the synthetic corpus."""

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/zero_start.pt")
    
    if not corpus_path.exists():
        raise FileNotFoundError(f"Corpus not found at {corpus_path}")
    
    if not checkpoint_path.exists():
        print(f"ERROR: Run train_truly_fresh.py first to create the zero-start checkpoint!")
        return
    
    print("="*80)
    print("CONTINUING ZERO-START MODEL WITH SYNTHETIC CORPUS")
    print("="*80)
    print()
    
    args = SimpleNamespace(
        corpus=corpus_path,
        epochs=30,  # More epochs since starting fresh
        chunk_size=2048,
        shuffle_buffer=64,
        seed=13,
        patience=6,  # More patience
        min_delta=0.05,
        log_every=50,  # Log more frequently
        checkpoint_path=checkpoint_path,
        checkpoint_interval=500,  # Save more often
        resume=True,  # Resume from our fresh checkpoint
        salience_filter=False,
        min_uncertainty=0.0,
        min_novelty=0.0,
        max_drag=1.0,
    )
    
    print(f"Resuming from: {checkpoint_path}")
    print(f"Corpus: {corpus_path}")
    print()
    
    train_corpus(args)

if __name__ == "__main__":
    main()