#!/usr/bin/env python3
"""
OSRS Agent Training - CPU Only (No GPU Required)
Uses tiny models that train reasonably fast on CPU.
"""
import json
import os
import sys
from pathlib import Path
import torch
from transformers import (
    AutoModelForCausalLM, 
    AutoTokenizer, 
    TrainingArguments,
    DataCollatorForSeq2Seq,
    Trainer
)
from peft import LoraConfig, get_peft_model, TaskType
from datasets import Dataset

# Configuration - Optimized for CPU
HOME = Path("/home/donn")
TRAIN_DATA = HOME / "organized" / "02-ai-ml" / "agents" / "osrs-agent" / "training_data" / "osrs_complete_combined.jsonl"
OUTPUT_DIR = HOME / "organized" / "02-ai-ml" / "agents" / "osrs-agent" / "trained_model_cpu"

# Tiny models for fast CPU training
# MODEL_NAME = "HuggingFaceTB/SmolLM-135M-Instruct"  # Ultra-tiny, very fast
MODEL_NAME = "HuggingFaceTB/SmolLM-360M-Instruct"     # Better quality, still fast

# Training hyperparameters
LORA_R = 8
LORA_ALPHA = 16
LORA_DROPOUT = 0.05
LEARNING_RATE = 2e-4
EPOCHS = 3
BATCH_SIZE = 1  # Small batch for CPU
GRAD_ACCUM = 8  # Larger grad accum to compensate
MAX_LENGTH = 128  # Shorter sequences for speed

def load_data(path):
    """Load training data from JSONL"""
    examples = []
    with open(path) as f:
        for line in f:
            if line.strip():
                examples.append(json.loads(line))
    return examples

def format_example(ex):
    """Format training example - compact for CPU"""
    state = ex['game_state']
    action = ex['agent_action']
    
    # Super compact format
    user = f"HP:{state['health']}/{state['max_health']} Pray:{state['prayer']} Combat:{state['combat_level']} Task:{state.get('task', 'gameplay')}"
    assistant = f"Action:{action['action_type']} {action.get('target', '')}"
    
    return {"prompt": user, "response": assistant}

def main():
    print("="*60)
    print("OSRS Agent Training - CPU Only")
    print("="*60)
    print(f"Model: {MODEL_NAME}")
    print(f"Device: CPU (no GPU)")
    print(f"Warning: This will be slow but free!")
    print()
    
    # Force CPU
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
    
    # Check device
    device = "cpu"
    print(f"[Device] {device.upper()}")
    print(f"[Warning] Training on CPU is ~20x slower than GPU")
    print()
    
    # Load data
    print(f"[1] Loading training data...")
    examples = load_data(TRAIN_DATA)
    print(f"    {len(examples)} examples loaded")
    
    # Format
    print(f"[2] Formatting examples...")
    formatted = [format_example(ex) for ex in examples]
    ds = Dataset.from_list(formatted)
    
    # Tokenizer
    print(f"[3] Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token
    
    def tokenize(sample):
        text = f"User: {sample['prompt']}\nAssistant: {sample['response']}"
        return tokenizer(text, truncation=True, max_length=MAX_LENGTH, padding='max_length')
    
    print(f"[4] Tokenizing...")
    ds = ds.map(tokenize, batched=False, remove_columns=ds.column_names)
    
    # Load model on CPU
    print(f"[5] Loading model on CPU...")
    print(f"    This may take a few minutes...")
    
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        device_map="cpu",  # Force CPU
        torch_dtype=torch.float32,  # Use float32 for CPU
        trust_remote_code=True,
    )
    model.config.use_cache = False
    
    # LoRA
    print(f"[6] Setting up LoRA...")
    lora_config = LoraConfig(
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        lora_dropout=LORA_DROPOUT,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    
    # Training
    print(f"[7] Configuring training...")
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    
    training_args = TrainingArguments(
        output_dir=str(OUTPUT_DIR),
        num_train_epochs=EPOCHS,
        per_device_train_batch_size=BATCH_SIZE,
        gradient_accumulation_steps=GRAD_ACCUM,
        learning_rate=LEARNING_RATE,
        logging_steps=50,  # Log less frequently
        save_steps=1000,
        save_total_limit=1,
        fp16=False,  # No FP16 on CPU
        bf16=False,
        report_to="none",
        gradient_checkpointing=False,  # Save memory
        optim="adamw_torch",
        lr_scheduler_type="cosine",
        warmup_ratio=0.1,
        remove_unused_columns=False,
        dataloader_num_workers=0,  # Single process for CPU
    )
    
    data_collator = DataCollatorForSeq2Seq(tokenizer, padding=True, return_tensors="pt")
    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=ds,
        data_collator=data_collator,
    )
    
    # Estimate time
    est_time_hours = (len(examples) * EPOCHS) / 60  # Rough estimate
    
    # Train
    print(f"[8] Starting training...")
    print(f"    Epochs: {EPOCHS}")
    print(f"    Estimated time: {est_time_hours:.1f} hours")
    print(f"    (You can stop anytime with Ctrl+C)")
    print()
    
    try:
        trainer.train()
    except KeyboardInterrupt:
        print("\n\nTraining interrupted by user.")
    
    # Save
    print(f"\n[9] Saving model to {OUTPUT_DIR}...")
    model.save_pretrained(str(OUTPUT_DIR / "lora_adapter"))
    tokenizer.save_pretrained(str(OUTPUT_DIR))
    
    # Config
    config = {
        "model_name": MODEL_NAME,
        "training_examples": len(examples),
        "lora_r": LORA_R,
        "lora_alpha": LORA_ALPHA,
        "epochs": EPOCHS,
        "device": "cpu",
    }
    with open(OUTPUT_DIR / "config.json", "w") as f:
        json.dump(config, f, indent=2)
    
    print("\n" + "="*60)
    print("TRAINING COMPLETE!")
    print("="*60)
    print(f"Model saved: {OUTPUT_DIR}")
    print(f"\nThis model was trained on CPU and is smaller")
    print(f"but can still make OSRS gameplay predictions!")

if __name__ == "__main__":
    main()