#!/usr/bin/env python3
"""
Train LoRA adapter on Cyron summary dataset.

Uses Hugging Face TRL for SFT training.
"""

import argparse
import os
import yaml
from pathlib import Path

import torch


def train(config_path):
    """Train LoRA adapter using TRL."""
    
    with open(config_path) as f:
        config = yaml.safe_load(f)
    
    from transformers import (
        AutoModelForCausalLM,
        AutoTokenizer,
        TrainingArguments,
    )
    from peft import (
        LoraConfig,
        get_peft_model,
    )
    from trl import SFTTrainer
    
    print(f"Loading model: {config['base_model']}")
    
    # Load model - let the model's own quantization config handle it
    # (Ornith uses CompressedTensors, not BitsAndBytes)
    # Load on CPU first, then DeepSpeed will distribute
    model = AutoModelForCausalLM.from_pretrained(
        config["base_model"],
        dtype=torch.bfloat16,
        device_map="cpu",  # Load on CPU, DeepSpeed distributes
    )
    
    # Add LoRA
    lora_config = LoraConfig(
        r=config["lora_r"],
        lora_alpha=config["lora_alpha"],
        lora_dropout=config["lora_dropout"],
        target_modules=config["target_modules"],
        task_type=config.get("lora_task_type", "CAUSAL_LM"),
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(config["base_model"])
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"
    
    # Load dataset
    from datasets import load_dataset
    import os
    
    # Get dataset paths - dataset is in training/data/ relative to repo root
    repo_root = Path(__file__).parent.parent.parent
    train_path = str(repo_root / "training" / "data" / "train.jsonl")
    test_path = str(repo_root / "training" / "data" / "test.jsonl")
    
    dataset = load_dataset(
        "json",
        data_files={
            "train": train_path,
            "test": test_path,
        },
    )
    
    # Training arguments (max_seq_length removed - passed to SFTTrainer instead)
    training_args = TrainingArguments(
        output_dir=config["train_params"]["output_dir"],
        num_train_epochs=config["train_params"]["num_train_epochs"],
        per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"],
        gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"],
        learning_rate=config["train_params"]["learning_rate"],
        lr_scheduler_type=config["train_params"]["lr_scheduler_type"],
        weight_decay=config["train_params"]["weight_decay"],
        warmup_ratio=config["train_params"]["warmup_ratio"],
        logging_steps=config["train_params"]["logging_steps"],
        save_steps=config["train_params"]["save_steps"],
        save_total_limit=config["train_params"]["save_total_limit"],
        eval_strategy=config.get("eval_strategy", "steps"),
        eval_steps=config.get("eval_steps", 100),
        bf16=config.get("mixed_precision", "bf16") == "bf16",
        fp16=config.get("mixed_precision", "bf16") == "fp16",
        gradient_checkpointing=config.get("gradient_checkpointing", True),
    )
    
    # SFT Trainer (DeepSpeed handles distributed training via config)
    from trl import SFTTrainer
    
    trainer = SFTTrainer(
        model=model,
        processing_class=tokenizer,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        args=training_args,
        dataset_text_field="text",
        max_seq_length=config["train_params"]["max_seq_length"],
    )
    
    # Train
    print("Starting training...")
    trainer.train()
    
    # Save
    trainer.save_model(config["train_params"]["output_dir"])
    tokenizer.save_pretrained(config["train_params"]["output_dir"])
    
    print(f"Training complete! Model saved to {config['train_params']['output_dir']}")


def main():
    parser = argparse.ArgumentParser(description="Train LoRA adapter")
    parser.add_argument("--config", type=str, default="configs/llama2-7b-lora.yaml", 
                       help="Training configuration file")
    args = parser.parse_args()
    
    train(args.config)


if __name__ == "__main__":
    main()