#!/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,
        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 distributed training support
    # Use FSDP for multi-GPU training
    from accelerate import Accelerator
    accelerator = Accelerator()
    
    # Load model on CPU first, then distribute
    model = AutoModelForCausalLM.from_pretrained(
        config["base_model"],
        torch_dtype=torch.bfloat16,
        device_map="cpu",  # Load on CPU first
    )
    
    # 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),
    )
    
    # Use FSDP for multi-GPU training
    from trl import SFTTrainer
    
    # Prepare model for FSDP
    model = accelerator.prepare(model)
    
    # 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"],
    )
    
    # Prepare trainer for distributed training
    trainer = accelerator.prepare(trainer)
    
    # 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()