"""
LoRA fine-tuning using DPO (Direct Preference Optimization).

Unlike SFT which trains on positive examples only, DPO trains on
preference pairs (chosen vs rejected) to directly optimize for preferences.
"""

from __future__ import annotations

import argparse
import gc
import json
import tempfile
from pathlib import Path
from typing import TYPE_CHECKING

import torch
from datasets import Dataset
from peft import LoraConfig, TaskType, get_peft_model
from trl import DPOTrainer, DPOConfig

from wisent.comparison.utils import (
    generate_contrastive_pairs,
    create_test_only_task,
    extract_accuracy,
    run_lm_eval_evaluation,
    run_ll_evaluation,
    load_model_and_tokenizer,
)
from wisent.core.utils.device import preferred_dtype

if TYPE_CHECKING:
    from wisent.core.models.wisent_model import WisentModel


def create_dpo_dataset(pairs: list[dict]) -> Dataset:
    """
    Convert contrastive pairs to DPO dataset format.

    DPO expects:
    - prompt: the input prompt
    - chosen: the preferred response
    - rejected: the non-preferred response
    """
    data = {
        "prompt": [],
        "chosen": [],
        "rejected": [],
    }

    for pair in pairs:
        prompt = pair["prompt"]
        chosen = pair["positive_response"]["model_response"]
        rejected = pair["negative_response"]["model_response"]

        data["prompt"].append(prompt)
        data["chosen"].append(chosen)
        data["rejected"].append(rejected)

    return Dataset.from_dict(data)


def train_lora_dpo(
    task: str,
    model_name: str,
    output_path: str | Path,
    num_pairs: int = 50,
    device: str = "cuda:0",
    keep_intermediate: bool = False,
    lora_r: int = 16,
    lora_alpha: int = 32,
    lora_dropout: float = 0.05,
    learning_rate: float = 5e-5,
    num_epochs: int = 1,
    batch_size: int = 1,
    max_length: int = 512,
    max_prompt_length: int = 256,
    beta: float = 0.1,
) -> Path:
    """
    Train a LoRA adapter using DPO on contrastive pairs from an lm-eval task.

    Args:
        task: lm-eval task name (e.g., 'boolq', 'cb')
        model_name: HuggingFace model name
        output_path: Where to save the trained adapter
        num_pairs: Number of preference pairs to use
        device: Device to run on
        keep_intermediate: Whether to keep intermediate files
        lora_r: LoRA rank
        lora_alpha: LoRA alpha
        lora_dropout: LoRA dropout
        learning_rate: Learning rate
        num_epochs: Number of training epochs
        batch_size: Training batch size
        max_length: Max total sequence length
        max_prompt_length: Max prompt length
        beta: DPO beta parameter (controls deviation from reference model)

    Returns:
        Path to saved adapter
    """
    output_path = Path(output_path)

    # Step 1: Generate contrastive pairs
    print(f"\n{'='*60}")
    print(f"Step 1: Generating {num_pairs} preference pairs from {task}")
    print(f"{'='*60}")

    pairs, pairs_file = generate_contrastive_pairs(task, num_pairs)
    print(f"Generated {len(pairs)} preference pairs")

    # Step 2: Create DPO dataset
    print(f"\n{'='*60}")
    print(f"Step 2: Creating DPO dataset")
    print(f"{'='*60}")

    dataset = create_dpo_dataset(pairs)
    print(f"Dataset size: {len(dataset)}")

    # Step 3: Load model
    print(f"\n{'='*60}")
    print(f"Step 3: Loading model {model_name}")
    print(f"{'='*60}")

    model, tokenizer = load_model_and_tokenizer(model_name, device, eval_mode=False)

    # Ensure tokenizer has padding
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"  # DPO typically uses left padding

    # Step 4: Configure LoRA
    print(f"\n{'='*60}")
    print(f"Step 4: Configuring LoRA (r={lora_r}, alpha={lora_alpha})")
    print(f"{'='*60}")

    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=lora_r,
        lora_alpha=lora_alpha,
        lora_dropout=lora_dropout,
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
        bias="none",
    )

    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # Step 5: Configure DPO training
    print(f"\n{'='*60}")
    print(f"Step 5: Configuring DPO training")
    print(f"{'='*60}")

    training_output_dir = tempfile.mkdtemp(prefix="lora_dpo_training_")

    # Determine dtype
    dtype = preferred_dtype(device)

    training_args = DPOConfig(
        output_dir=training_output_dir,
        num_train_epochs=num_epochs,
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=1,
        learning_rate=learning_rate,
        weight_decay=0.01,
        warmup_ratio=0.1,
        logging_steps=10,
        save_strategy="no",
        bf16=(dtype == torch.bfloat16),
        fp16=(dtype == torch.float16),
        report_to="none",
        max_length=max_length,
        max_prompt_length=max_prompt_length,
        beta=beta,
        loss_type="sigmoid",  # Standard DPO loss
    )

    print(f"Beta: {beta}")
    print(f"Max length: {max_length}")
    print(f"Max prompt length: {max_prompt_length}")
    print(f"Learning rate: {learning_rate}")
    print(f"Epochs: {num_epochs}")
    print(f"Batch size: {batch_size}")

    # Step 6: Train with DPO
    print(f"\n{'='*60}")
    print(f"Step 6: Training with DPO")
    print(f"{'='*60}")

    trainer = DPOTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        processing_class=tokenizer,
    )

    trainer.train()

    # Step 7: Save adapter
    print(f"\n{'='*60}")
    print(f"Step 7: Saving LoRA adapter")
    print(f"{'='*60}")

    output_path.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(output_path)
    tokenizer.save_pretrained(output_path)

    # Save metadata
    metadata = {
        "task": task,
        "model": model_name,
        "training_method": "dpo",
        "num_pairs": len(pairs),
        "lora_r": lora_r,
        "lora_alpha": lora_alpha,
        "lora_dropout": lora_dropout,
        "learning_rate": learning_rate,
        "num_epochs": num_epochs,
        "batch_size": batch_size,
        "max_length": max_length,
        "max_prompt_length": max_prompt_length,
        "beta": beta,
    }
    with open(output_path / "metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)

    # Cleanup
    del model, trainer
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    if not keep_intermediate:
        import os
        import shutil
        os.unlink(pairs_file)
        shutil.rmtree(training_output_dir, ignore_errors=True)

    print(f"\nDPO LoRA adapter saved to {output_path}")
    return output_path


def evaluate_lora_dpo(
    model_name: str,
    lora_path: str | Path,
    task: str,
    train_ratio: float = 0.8,
    device: str = "cuda:0",
    batch_size: int = 1,
    max_batch_size: int = 8,
    limit: int | None = None,
    output_dir: str | Path = None,
) -> dict:
    """
    Evaluate a trained DPO LoRA adapter.

    Compares base model vs DPO-LoRA model accuracy.
    """
    from wisent.core.models.wisent_model import WisentModel
    from wisent.comparison.lora import apply_lora_to_model, remove_lora

    lora_path = Path(lora_path)

    # Create test task
    print(f"\n{'='*60}")
    print(f"Creating test task for: {task}")
    print(f"{'='*60}")

    task_dict = create_test_only_task(task, train_ratio=train_ratio)

    # Load model
    print(f"\n{'='*60}")
    print(f"Loading model: {model_name}")
    print(f"{'='*60}")
    wisent_model = WisentModel(model_name=model_name, device=device)

    # Base evaluation
    print(f"\n{'='*60}")
    print(f"Running BASE evaluation")
    print(f"{'='*60}")

    base_results = run_lm_eval_evaluation(wisent_model, task_dict, task, batch_size, max_batch_size, limit)
    base_acc_lm_eval = extract_accuracy(base_results, task)
    print(f"Base accuracy (lm-eval): {base_acc_lm_eval:.4f}")

    base_acc_ll = run_ll_evaluation(wisent_model, task_dict, task, limit)
    print(f"Base accuracy (LL): {base_acc_ll:.4f}")

    # Apply DPO LoRA
    print(f"\n{'='*60}")
    print(f"Applying DPO LoRA adapter from: {lora_path}")
    print(f"{'='*60}")
    apply_lora_to_model(wisent_model, lora_path)

    # LoRA evaluation
    print(f"\n{'='*60}")
    print(f"Running DPO-LORA evaluation")
    print(f"{'='*60}")

    lora_results = run_lm_eval_evaluation(wisent_model, task_dict, task, batch_size, max_batch_size, limit)
    lora_acc_lm_eval = extract_accuracy(lora_results, task)
    print(f"DPO-LoRA accuracy (lm-eval): {lora_acc_lm_eval:.4f}")

    lora_acc_ll = run_ll_evaluation(wisent_model, task_dict, task, limit)
    print(f"DPO-LoRA accuracy (LL): {lora_acc_ll:.4f}")

    # Cleanup
    remove_lora(wisent_model)
    del wisent_model
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # Results
    results = {
        "task": task,
        "model": model_name,
        "training_method": "dpo",
        "lora_path": str(lora_path),
        "train_ratio": train_ratio,
        "base_accuracy_lm_eval": base_acc_lm_eval,
        "base_accuracy_ll": base_acc_ll,
        "lora_accuracy_lm_eval": lora_acc_lm_eval,
        "lora_accuracy_ll": lora_acc_ll,
        "diff_lm_eval": lora_acc_lm_eval - base_acc_lm_eval,
        "diff_ll": lora_acc_ll - base_acc_ll,
    }

    # Print summary
    print(f"\n{'='*60}")
    print(f"RESULTS SUMMARY")
    print(f"{'='*60}")
    print(f"Task: {task}")
    print(f"Model: {model_name}")
    print(f"Training: DPO")
    print(f"{'-'*60}")
    print(f"{'Metric':<20} {'Base':<12} {'DPO-LoRA':<12} {'Diff':<12}")
    print(f"{'-'*60}")
    print(f"{'lm-eval accuracy':<20} {base_acc_lm_eval:<12.4f} {lora_acc_lm_eval:<12.4f} {lora_acc_lm_eval - base_acc_lm_eval:+.4f}")
    print(f"{'LL accuracy':<20} {base_acc_ll:<12.4f} {lora_acc_ll:<12.4f} {lora_acc_ll - base_acc_ll:+.4f}")
    print(f"{'='*60}")

    # Save results
    if output_dir:
        output_dir = Path(output_dir)
        model_dir_name = model_name.replace("/", "_")
        output_dir = output_dir / model_dir_name
        output_dir.mkdir(parents=True, exist_ok=True)
        results_file = output_dir / f"{task}_lora_dpo_eval_results.json"
        with open(results_file, "w") as f:
            json.dump(results, f, indent=2)
        print(f"\nResults saved to: {results_file}")

    return results


def main():
    parser = argparse.ArgumentParser(description="Train and evaluate LoRA adapter using DPO")
    parser.add_argument("--model", required=True, help="HuggingFace model name")
    parser.add_argument("--task", default="boolq", help="lm-eval task name")
    parser.add_argument("--output-dir", default="/home/ubuntu/output", help="Output directory")
    parser.add_argument("--num-pairs", type=int, default=50, help="Number of preference pairs")
    parser.add_argument("--device", default="cuda:0", help="Device")
    parser.add_argument("--lora-r", type=int, default=16, help="LoRA rank")
    parser.add_argument("--lora-alpha", type=int, default=32, help="LoRA alpha")
    parser.add_argument("--lora-dropout", type=float, default=0.05, help="LoRA dropout")
    parser.add_argument("--learning-rate", type=float, default=5e-5, help="Learning rate")
    parser.add_argument("--num-epochs", type=int, default=1, help="Number of epochs")
    parser.add_argument("--batch-size", type=int, default=1, help="Training batch size")
    parser.add_argument("--max-length", type=int, default=512, help="Max total sequence length")
    parser.add_argument("--max-prompt-length", type=int, default=256, help="Max prompt length")
    parser.add_argument("--beta", type=float, default=0.1, help="DPO beta (controls KL penalty)")
    parser.add_argument("--keep-intermediate", action="store_true", help="Keep intermediate files")
    # Eval args
    parser.add_argument("--train-ratio", type=float, default=0.8, help="Train/test split ratio")
    parser.add_argument("--eval-batch-size", default="auto", help="Eval batch size")
    parser.add_argument("--eval-max-batch-size", type=int, default=64, help="Max eval batch size")
    parser.add_argument("--eval-limit", type=int, default=None, help="Limit eval examples")
    parser.add_argument("--skip-eval", action="store_true", help="Skip evaluation after training")

    args = parser.parse_args()

    output_path = Path(args.output_dir) / f"{args.task}_lora_dpo_adapter"

    # Train
    train_lora_dpo(
        task=args.task,
        model_name=args.model,
        output_path=output_path,
        num_pairs=args.num_pairs,
        device=args.device,
        keep_intermediate=args.keep_intermediate,
        lora_r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        learning_rate=args.learning_rate,
        num_epochs=args.num_epochs,
        batch_size=args.batch_size,
        max_length=args.max_length,
        max_prompt_length=args.max_prompt_length,
        beta=args.beta,
    )

    # Evaluate
    if not args.skip_eval:
        eval_batch_size = args.eval_batch_size
        if eval_batch_size != "auto":
            eval_batch_size = int(eval_batch_size)

        evaluate_lora_dpo(
            model_name=args.model,
            lora_path=output_path,
            task=args.task,
            train_ratio=args.train_ratio,
            device=args.device,
            batch_size=eval_batch_size,
            max_batch_size=args.eval_max_batch_size,
            limit=args.eval_limit,
            output_dir=args.output_dir,
        )


if __name__ == "__main__":
    main()
