"""Direct Preference Optimization (DPO) for Bee using TRL."""

import argparse
import logging
import sys
from pathlib import Path

from datasets import load_dataset
from transformers import AutoTokenizer, TrainingArguments, set_seed
from trl import DPOTrainer, DPOConfig

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from bee.register import register
from bee.modeling_bee import BeeForCausalLM

register()

logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s")
logger = logging.getLogger("bee.dpo")


def get_args():
    parser = argparse.ArgumentParser(description="DPO train Bee")
    parser.add_argument("--model_path", type=str, required=True, help="SFT checkpoint to align")
    parser.add_argument("--dataset", type=str, default="trl-lib/ultrafeedback_binarized", help="HF preference dataset")
    parser.add_argument("--output_dir", type=str, required=True)
    parser.add_argument("--max_length", type=int, default=2048)
    parser.add_argument("--batch_size", type=int, default=2)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=8)
    parser.add_argument("--learning_rate", type=float, default=5e-7)
    parser.add_argument("--num_train_epochs", type=int, default=1)
    parser.add_argument("--beta", type=float, default=0.1)
    parser.add_argument("--save_steps", type=int, default=500)
    parser.add_argument("--logging_steps", type=int, default=50)
    parser.add_argument("--bf16", action="store_true", default=True)
    parser.add_argument("--seed", type=int, default=42)
    return parser.parse_args()


def main():
    args = get_args()
    set_seed(args.seed)

    logger.info("Loading model from %s", args.model_path)
    model = BeeForCausalLM.from_pretrained(args.model_path)
    ref_model = BeeForCausalLM.from_pretrained(args.model_path)
    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    logger.info("Loading preference dataset: %s", args.dataset)
    ds = load_dataset(args.dataset, split="train")

    training_args = DPOConfig(
        output_dir=args.output_dir,
        num_train_epochs=args.num_train_epochs,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate,
        beta=args.beta,
        logging_steps=args.logging_steps,
        save_steps=args.save_steps,
        save_strategy="steps",
        bf16=args.bf16,
        max_length=args.max_length,
        report_to=["tensorboard"],
    )

    trainer = DPOTrainer(
        model=model,
        ref_model=ref_model,
        args=training_args,
        train_dataset=ds,
        tokenizer=tokenizer,
    )

    logger.info("Starting DPO training...")
    trainer.train()
    logger.info("DPO complete. Saving to %s", args.output_dir)
    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)


if __name__ == "__main__":
    main()