#!/usr/bin/env python3
"""DPO training skeleton using TRL's DPOTrainer.

This is a COMPLETE runnable script - generates synthetic preference pairs
dynamically so there are no external data dependencies.
Requires: pip install transformers datasets trl torch accelerate
"""

from __future__ import annotations

import sys

try:
    import torch
    from datasets import Dataset
    from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
    from trl import DPOTrainer
except ImportError:
    print("Install: pip install transformers datasets trl torch accelerate", file=sys.stderr)
    sys.exit(1)

# ── Synthetic preference data (stand-in for your real pairs) ──────────────
# Each entry: (prompt, chosen, rejected) — chosen is demonstrably better

SYNTHETIC_PAIRS = [
    ("Summarize the benefits of exercise.",
     "Exercise improves cardiovascular health, strengthens muscles, and boosts mood through endorphin release. Regular activity reduces the risk of chronic diseases.",
     "Exercise is good for you because it's healthy and people should do it more often."),
    ("Explain what a database index is.",
     "A database index is a data structure that improves the speed of data retrieval operations on a table. It works like a book's index: instead of scanning every page (full table scan), the database looks up the index to find the exact location of the data. The trade-off is slower writes, because the index must be updated on each insert/update/delete.",
     "An index makes queries faster. You create it with CREATE INDEX."),
    ("What is Python type hinting?",
     "Python type hints are optional annotations that specify the expected types of variables, function parameters, and return values. They are not enforced at runtime but enable static type checkers like mypy to catch type errors before execution. Syntax: `def greet(name: str) -> str:`. Introduced in PEP 484, they improve code readability, IDE autocompletion, and catch bugs early in development.",
     "Type hints tell you what type things are. Like `x: int = 5`. They're useful."),
    ("Write a function to compute factorial.",
     "```python\ndef factorial(n: int) -> int:\n    \"\"\"Return n! recursively.\"\"\"\n    if n < 0:\n        raise ValueError('n must be non-negative')\n    if n <= 1:\n        return 1\n    return n * factorial(n - 1)\n```",
     "```python\ndef fact(n):\n    x = 1\n    for i in range(1, n+1):\n        x *= i\n    return x\n```"),
    ("Is climate change real?",
     "Multiple independent lines of evidence confirm climate change is real and primarily driven by human greenhouse gas emissions: rising global temperatures (1.1°C above pre-industrial), accelerating ice sheet mass loss, sea level rise (20cm since 1900), and atmospheric CO2 concentrations exceeding 420 ppm — the highest in at least 800,000 years per ice core data.",
     "Climate change is debated. Some say it's real, others say it's natural cycles. Both sides have points."),
]


def create_dataset(model_name: str = "gpt2") -> tuple[Dataset, AutoTokenizer]:
    """Build a preference dataset from synthetic pairs."""
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    data = []
    for prompt, chosen, rejected in SYNTHETIC_PAIRS:
        # Format as conversation (share the prompt, differ the assistant reply)
        chosen_msg = [{"role": "user", "content": prompt}, {"role": "assistant", "content": chosen}]
        rejected_msg = [{"role": "user", "content": prompt}, {"role": "assistant", "content": rejected}]

        data.append({
            "prompt": tokenizer.apply_chat_template(
                [{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True
            ),
            "chosen": tokenizer.apply_chat_template(chosen_msg, tokenize=False),
            "rejected": tokenizer.apply_chat_template(rejected_msg, tokenize=False),
        })

    return Dataset.from_list(data), tokenizer


def main() -> None:
    model_name = "gpt2"  # Small model for demo speed; swap for your base/fine-tune

    print(f"Loading model: {model_name}")
    model = AutoModelForCausalLM.from_pretrained(model_name)
    ref_model = AutoModelForCausalLM.from_pretrained(model_name)  # Frozen reference
    dataset, tokenizer = create_dataset(model_name)

    print(f"Dataset: {len(dataset)} preference pairs")

    training_args = TrainingArguments(
        output_dir="./dpo_output",
        per_device_train_batch_size=2,
        num_train_epochs=1,            # DPO overfits fast — keep at ~1 epoch
        learning_rate=5e-6,             # Smaller than SFT (typ ~5e-6 to 5e-5)
        logging_steps=1,
        save_strategy="no",
        remove_unused_columns=False,
        report_to="none",               # Disable wandb/tensorboard for demo
    )

    trainer = DPOTrainer(
        model=model,
        ref_model=ref_model,
        args=training_args,
        train_dataset=dataset,
        tokenizer=tokenizer,
        beta=0.1,                       # Temperature of implicit reward (0.05-0.5)
        max_prompt_length=256,
        max_length=512,
    )

    print("Training DPO (this may take a minute on CPU)...")
    trainer.train()

    print("\nDone. Key levers:")
    print("  beta=0.1    — Lower -> pushes harder toward chosen (risk: mode collapse)")
    print("  epochs=1    — More -> overfits to style / reward hacks")
    print("  lr=5e-6     — DPO uses smaller LR than SFT")
    print("  Evaluate: check preference accuracy + capability benchmarks")

    # Optional: save
    # trainer.save_model("./dpo_final")
    # tokenizer.save_pretrained("./dpo_final")


if __name__ == "__main__":
    main()