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

Uses Hugging Face TRL for SFT training with QLoRA.
"""

import argparse
import yaml
from pathlib import Path


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,
        BitsAndBytesConfig,
        TrainingArguments,
    )
    from peft import (
        prepare_model_for_kbit_training,
        LoraConfig,
        get_peft_model,
    )
    from trl import SFTTrainer, SFTConfig
    
    print(f"Loading model: {config['base_model']}")
    
    # Load model with quantization
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=config.get("load_in_4bit", True),
        bnb_4bit_compute_dtype=config.get("bnb_4bit_compute_dtype", "bfloat16"),
        bnb_4bit_quant_type=config.get("bnb_4bit_quant_type", "nf4"),
        use_nested_quant=config.get("use_nested_quant", False),
    )
    
    # Use all available GPUs
    device_map = "auto" if torch.cuda.device_count() == 1 else "balanced"
    
    model = AutoModelForCausalLM.from_pretrained(
        config["base_model"],
        quantization_config=bnb_config,
        device_map=device_map,
    )
    model = prepare_model_for_kbit_training(model)
    
    # 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
    
    dataset = load_dataset(
        "json",
        data_files={
            "train": config["dataset"][0]["path"].replace("../", ""),
            "test": config["dataset"][0]["path"].replace("../", "").replace("combined_20k.jsonl", "test.jsonl"),
        },
    )
    
    # Training arguments
    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"],
        max_seq_length=config["train_params"]["max_seq_length"],
        logging_steps=config["train_params"]["logging_steps"],
        save_steps=config["train_params"]["save_steps"],
        save_total_limit=config["train_params"]["save_total_limit"],
        evaluation_strategy=config.get("eval_strategy", "steps"),
        eval_steps=config.get("eval_steps", 100),
        mixed_precision=config.get("mixed_precision", "bf16"),
        gradient_checkpointing=config.get("gradient_checkpointing", True),
    )
    
    # SFT Trainer
    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        args=training_args,
        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()