#!/usr/bin/env python3
"""
RLHF Training Script — Qwen3.8-27B DPO (Direct Preference Optimization)
──────────────────────────────────────────────────────────────────────────────
Trains a reward model and applies DPO to align the SFT model with human
preference data (win/loss pairs from SFT training).

Usage:
  python3 train_rlhf.py

Dependencies:
  pip install trl datasets torch accelerate
"""

import os
import json
import random
import math
from dataclasses import dataclass, field
from typing import List, Optional
from pathlib import Path

import torch
import torch.nn as nn
from datasets import Dataset, DatasetDict
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
)
from peft import PeftConfig, PeftModel, LoraConfig, TaskType, get_peft_model
from trl import DPOConfig, DPOTrainer, SFTTrainer

# ─── Configuration ──────────────────────────────────────────────────────────
@dataclass
class DPOConfig:
    model_name: str = "unsloth/Qwen3.8-27B-GGUF"
    sft_checkpoint: str = "./models/qwen3.8-27b-sft-full"
    output_dir: str = "./models/qwen3.8-27b-dpo"
    adapter_dir: str = "./adapters/qwen3.8-27b-dpo-lora"
    beta: float = 0.1
    learning_rate: float = 1e-7
    lr_scheduler: str = "cosine"
    lr_warmup_ratio: float = 0.1
    batch_size: int = 1
    gradient_accumulation_steps: int = 4
    num_train_epochs: int = 1
    save_steps: int = 50
    eval_steps: int = 25
    log_steps: int = 10
    seed: int = 42
    lora_rank: int = 64
    lora_alpha: int = 128
    lora_dropout: float = 0.05
    target_modules: List[str] = field(default_factory=lambda: [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ])
    use_8bit_adam: bool = True
    fp16: bool = True
    bf16: bool = True
    use_llama_flash_attn2: bool = True
    use_dora: bool = False
    use_rope_scaling: bool = False
    max_length: int = 4096
    max_prompt_length: int = 512
    max_target_length: int = 2048
    pad_to_multiple_of: int = 8
    pretrain_path: Optional[str] = None

# ─── Data Loading ───────────────────────────────────────────────────────────
def load_preference_dataset(data_path: str = "./data/dpo_pairs.json") -> DatasetDict:
    """Load preference dataset with (chosen, rejected) pairs."""
    if not os.path.exists(data_path):
        print(f"[INFO] Data not found at {data_path}, generating synthetic pairs...")
        return generate_preference_data()

    with open(data_path, "r") as f:
        data = json.load(f)

    if isinstance(data, dict) and "chosen" in data[0]:
        return DatasetDict({
            "chosen": Dataset.from_list(data),
            "rejected": Dataset.from_list(data),
        })
    else:
        raise ValueError(f"Unexpected data format: {type(data[0])}")


def generate_preference_data(num_samples: int = 2000) -> DatasetDict:
    """Generate synthetic preference pairs for DPO training."""

    def generate_response(prompt: str, difficulty: int = 0, make_wrong: bool = False) -> str:
        """Generate a response, optionally making it worse."""
        if make_wrong:
            return f"### Instruction:\n{prompt}\n\n### Response:\nI don't know. This is a bad response."
        else:
            good_responses = {
                "Explain quantum entanglement": "Quantum entanglement is a phenomenon where two or more particles become correlated in such a way that the quantum state of each particle cannot be described independently. Measuring one particle instantly determines the state of the other, regardless of distance.",
                "Write a Python function to reverse a string": "def reverse_string(s):\n    return s[::-1]\n\nprint(reverse_string('hello'))  # 'olleh'",
                "Translate 'La vida es bella'": "The Spanish phrase 'La vida es bella' translates to 'Life is beautiful' in English.",
                "Summarize: Artificial intelligence is transforming how we live and work. Machine learning models can now generate text, create images, and even play games. However, AI also raises concerns about job displacement, bias, and misinformation.": "AI is transforming society, enabling text generation, image creation, and gaming. However, it also raises concerns about job displacement, bias, and misinformation.",
                "Solve: What is 123456789 * 987654321?": "123456789 * 987654321 = 121,932,631,112,635,269",
                "Write a haiku about mountains": "Towering peaks pierce sky —\nSnow-capped giants stand in silence —\nNature's ancient might.",
                "What is the capital of France?": "The capital of France is Paris.",
                "Explain how photosynthesis works": "Photosynthesis is the process by which plants convert light energy into chemical energy. Chlorophyll in chloroplasts captures sunlight, splitting water molecules and releasing oxygen. The energy is stored in glucose (C6H12O6).",
                "Debug: def fib(n):\n    if n <= 1:\n        return n\n    return fib(n-1) + fib(n-2)": "The issue is that the base case should handle both n=0 and n=1 together. Also, add type hints and docstring.\n\nCorrected:\n```python\ndef fib(n: int) -> int:\n    \"\"\"Compute the nth Fibonacci number.\"\"\"\n    if n <= 1:\n        return n\n    return fib(n-1) + fib(n-2)\n```",
            }
            return good_responses.get(prompt, f"This is a good response to: {prompt}")

    examples = []
    for _ in range(num_samples):
        prompt = random.choice(list(generate_response.keys()))
        chosen = generate_response(prompt, difficulty=random.randint(0, 3))
        rejected = generate_response(prompt, difficulty=random.randint(0, 3), make_wrong=True)
        examples.append({
            "prompt": prompt,
            "chosen": chosen,
            "rejected": rejected,
        })

    return DatasetDict({
        "chosen": Dataset.from_list(examples),
        "rejected": Dataset.from_list(examples),
    })


# ─── SFT (optional pre-training step) ───────────────────────────────────────
@dataclass
class SFTConfig:
    model_name: str = "unsloth/Qwen3.8-27B-GGUF"
    output_dir: str = "./models/qwen3.8-27b-sft-full"
    adapter_dir: str = "./adapters/qwen3.8-27b-sft-lora"
    batch_size: int = 1
    gradient_accumulation_steps: int = 4
    learning_rate: float = 1e-4
    lr_scheduler: str = "cosine"
    num_train_epochs: int = 3
    save_steps: int = 100
    eval_steps: int = 50
    log_steps: int = 10
    weight_decay: float = 0.05
    seed: int = 42
    lora_rank: int = 64
    lora_alpha: int = 128
    lora_dropout: float = 0.05
    target_modules: List[str] = field(default_factory=lambda: [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ])
    use_8bit_adam: bool = True
    bf16: bool = True
    use_llama_flash_attn2: bool = True
    use_dora: bool = False
    use_rope_scaling: bool = False


def train_sft(config: SFTConfig) -> str:
    """Train SFT model and return the checkpoint path."""
    from transformers import BitsAndBytesConfig, DataCollatorForLanguageModeling
    from peft import LoraConfig, get_peft_model, TaskType
    from datasets import Dataset
    from transformers import TrainingArguments, Trainer, AutoTokenizer

    cfg = config
    print(f"[SFT] Model: {cfg.model_name}")
    print(f"[SFT] Output dir: {cfg.output_dir}")

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(cfg.model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # Quantization
    quant_config = BitsAndBytesConfig(
        load_in_8bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
        bnb_4bit_use_double_quant=True,
        llm_int8_enable_fp32_cpu_offload=False,
        llm_int8_threshold=6.0,
    )

    # Load model
    model = AutoModelForCausalLM.from_pretrained(
        cfg.model_name,
        quantization_config=quant_config,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
        device_map="auto",
    )

    # LoRA
    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        inference_mode=False,
        r=cfg.lora_rank,
        lora_alpha=cfg.lora_alpha,
        lora_dropout=cfg.lora_dropout,
        target_modules=cfg.target_modules,
    )
    model = get_peft_model(model, peft_config)
    model.print_trainable_parameters()

    # Synthetic SFT data
    prompts = [
        "Explain quantum entanglement in simple terms.",
        "Write a Python function to reverse a string.",
        "Translate 'La vida es bella' to English.",
        "Summarize this paragraph in one sentence: AI is transforming society.",
        "Solve: What is 123456789 * 987654321?",
        "Write a haiku about space.",
        "What is the capital of Japan?",
        "Explain how photosynthesis works.",
        "Debug this code: def fib(n): if n <= 1: return n; return fib(n-1) + fib(n-2)",
    ]
    responses = [
        "Quantum entanglement is a phenomenon where particles become correlated so that measuring one instantly determines the state of the other, regardless of distance. This defies classical intuition.",
        "def reverse_string(s):\n    return s[::-1]\n\nprint(reverse_string('hello'))  # 'olleh'",
        "'La vida es bella' translates to 'Life is beautiful' in English.",
        "AI is transforming society by enabling text generation, image creation, and gaming. However, it raises concerns about job displacement, bias, and misinformation.",
        "123456789 * 987654321 = 121,932,631,112,635,269",
        "Starry skies above,\nNebulas dance in cosmic light,\nUniverse expands.",
        "The capital of Japan is Tokyo.",
        "Photosynthesis converts light energy into chemical energy. Chlorophyll captures sunlight, splitting water and releasing oxygen. Energy is stored in glucose.",
        "The base case handles n=0 and n=1 together, which is correct. The recursive step adds the two previous Fibonacci numbers. This is the standard recursive Fibonacci definition.",
    ]

    data = [{"instruction": p, "output": r} for p, r in zip(prompts, responses)]
    dataset = Dataset.from_list(data)
    dataset = dataset.train_test_split(test_size=0.1, seed=cfg.seed)

    collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
    training_args = TrainingArguments(
        output_dir=cfg.output_dir,
        per_device_train_batch_size=cfg.batch_size,
        gradient_accumulation_steps=cfg.gradient_accumulation_steps,
        learning_rate=cfg.learning_rate,
        fp16=cfg.fp16,
        bf16=cfg.bf16,
        logging_steps=cfg.log_steps,
        save_steps=cfg.save_steps,
        save_total_limit=1,
        evaluation_strategy="steps",
        eval_steps=cfg.eval_steps,
        per_device_eval_batch_size=cfg.batch_size,
        num_train_epochs=cfg.num_train_epochs,
        weight_decay=cfg.weight_decay,
        lr_scheduler_type=cfg.lr_scheduler,
        optim="paged_adamw_8bit" if cfg.use_8bit_adam else "adamw_torch",
        logging_strategy="steps",
        logging_first_step=True,
        report_to="none",
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        tokenizer=tokenizer,
        data_collator=collator,
    )
    trainer.train()

    model.save_pretrained(cfg.output_dir)
    tokenizer.save_pretrained(cfg.output_dir)

    return cfg.output_dir


# ─── DPO ────────────────────────────────────────────────────────────────────
def train_dpo(config: DPOConfig) -> None:
    """Train DPO model using TRL."""
    cfg = config

    print(f"[DPO] Model: {cfg.model_name}")
    print(f"[DPO] SFT checkpoint: {cfg.sft_checkpoint}")
    print(f"[DPO] Output dir: {cfg.output_dir}")
    print(f"[DPO] Beta: {cfg.beta}")
    print(f"[DPO] Learning rate: {cfg.learning_rate}")
    print(f"[DPO] Epochs: {cfg.num_train_epochs}")

    # Load SFT checkpoint (or full model if no SFT exists)
    if os.path.exists(cfg.sft_checkpoint):
        print(f"[DPO] Loading SFT checkpoint from {cfg.sft_checkpoint}")
        sft_config = PeftConfig.from_pretrained(cfg.sft_checkpoint)
        model = PeftModel.from_pretrained(
            AutoModelForCausalLM.from_pretrained(
                sft_config.base_model_name_or_path,
                torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
                device_map="auto",
            ),
            cfg.sft_checkpoint,
        )
    else:
        print(f"[DPO] No SFT checkpoint found, training from scratch...")
        sft_path = train_sft(SFTConfig(
            model_name=cfg.model_name,
            output_dir=cfg.sft_checkpoint,
            adapter_dir=cfg.adapter_dir + "-sft",
            num_train_epochs=1,
        ))
        model = PeftModel.from_pretrained(
            AutoModelForCausalLM.from_pretrained(
                sft_config.base_model_name_or_path if os.path.exists(cfg.sft_checkpoint) else cfg.model_name,
                torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
                device_map="auto",
            ),
            cfg.sft_checkpoint,
        )

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(cfg.model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # LoRA for DPO
    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        inference_mode=False,
        r=cfg.lora_rank,
        lora_alpha=cfg.lora_alpha,
        lora_dropout=cfg.lora_dropout,
        target_modules=cfg.target_modules,
    )
    model = get_peft_model(model, peft_config)
    model.print_trainable_parameters()

    # Load preference data
    dataset_dict = load_preference_dataset()
    print(f"[DPO] Dataset: {len(dataset_dict['chosen'])} chosen, {len(dataset_dict['rejected'])} rejected")

    # DPO config
    dpo_config = DPOConfig(
        model_name=cfg.model_name,
        beta=cfg.beta,
        learning_rate=cfg.learning_rate,
        lr_scheduler=cfg.lr_scheduler,
        lr_warmup_ratio=cfg.lr_warmup_ratio,
        batch_size=cfg.batch_size,
        gradient_accumulation_steps=cfg.gradient_accumulation_steps,
        num_train_epochs=cfg.num_train_epochs,
        save_steps=cfg.save_steps,
        eval_steps=cfg.eval_steps,
        log_steps=cfg.log_steps,
        seed=cfg.seed,
        lora_rank=cfg.lora_rank,
        lora_alpha=cfg.lora_alpha,
        lora_dropout=cfg.lora_dropout,
        target_modules=cfg.target_modules,
        use_8bit_adam=cfg.use_8bit_adam,
        fp16=cfg.fp16,
        bf16=cfg.bf16,
        use_llama_flash_attn2=cfg.use_llama_flash_attn2,
        use_dora=cfg.use_dora,
        use_rope_scaling=cfg.use_rope_scaling,
        max_length=cfg.max_length,
        max_prompt_length=cfg.max_prompt_length,
        max_target_length=cfg.max_target_length,
        pad_to_multiple_of=cfg.pad_to_multiple_of,
    )

    # DPO trainer
    trainer = DPOTrainer(
        model=model,
        ref_model=None,  # Can load a reference model or use the same model
        reward_model=None,  # DPO directly optimizes the policy
        policy_tokenizer=tokenizer,
        beta=cfg.beta,
        loss_kwargs={},
        args=TrainingArguments(
            output_dir=cfg.output_dir,
            per_device_train_batch_size=cfg.batch_size,
            gradient_accumulation_steps=cfg.gradient_accumulation_steps,
            learning_rate=cfg.learning_rate,
            lr_scheduler_type=cfg.lr_scheduler,
            lr_warmup_ratio=cfg.lr_warmup_ratio,
            logging_steps=cfg.log_steps,
            save_steps=cfg.save_steps,
            save_total_limit=1,
            evaluation_strategy="steps",
            eval_steps=cfg.eval_steps,
            per_device_eval_batch_size=cfg.batch_size,
            num_train_epochs=cfg.num_train_epochs,
            report_to="none",
        ),
        train_dataset=dataset_dict["chosen"],
        eval_dataset=dataset_dict["chosen"],  # Use same for eval
        tokenizer=tokenizer,
        max_prompt_length=cfg.max_prompt_length,
        max_length=cfg.max_length,
        padding_side="right",
    )

    # Train
    print("[DPO] Starting DPO training...")
    trainer.train()

    # Save
    trainer.save_model(cfg.output_dir)
    tokenizer.save_pretrained(cfg.output_dir)

    print(f"[DPO] Training complete! Model saved to: {cfg.output_dir}")


# ─── Main ───────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="RLHF Training: SFT + DPO")
    parser.add_argument("--mode", choices=["sft", "dpo", "full"], default="sft",
                        help="Training mode: sft, dpo, or full (sft+dpo)")
    parser.add_argument("--model_name", type=str, default="unsloth/Qwen3.8-27B-GGUF")
    parser.add_argument("--sft_output", type=str, default="./models/qwen3.8-27b-sft-full")
    parser.add_argument("--dpo_output", type=str, default="./models/qwen3.8-27b-dpo")
    parser.add_argument("--beta", type=float, default=0.1)
    parser.add_argument("--lr", type=float, default=1e-7)
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--data_path", type=str, default="./data/dpo_pairs.json")
    parser.add_argument("--num_samples", type=int, default=2000)

    args = parser.parse_args()

    if args.mode == "sft":
        config = SFTConfig(
            model_name=args.model_name,
            output_dir=args.sft_output,
            num_train_epochs=args.epochs,
            batch_size=args.batch_size,
        )
        result = train_sft(config)
        print(f"\n[COMPLETE] SFT model saved to: {result}")

    elif args.mode == "dpo":
        config = DPOConfig(
            model_name=args.model_name,
            sft_checkpoint=args.sft_output,
            output_dir=args.dpo_output,
            beta=args.beta,
            learning_rate=args.lr,
            num_train_epochs=args.epochs,
            batch_size=args.batch_size,
        )
        train_dpo(config)

    elif args.mode == "full":
        # Train SFT first, then DPO
        sft_config = SFTConfig(
            model_name=args.model_name,
            output_dir=args.sft_output,
            num_train_epochs=3,
            batch_size=args.batch_size,
        )
        sft_path = train_sft(sft_config)

        dpo_config = DPOConfig(
            model_name=args.model_name,
            sft_checkpoint=sft_path,
            output_dir=args.dpo_output,
            beta=args.beta,
            learning_rate=args.lr,
            num_train_epochs=1,
            batch_size=args.batch_size,
        )
        train_dpo(dpo_config)

        print(f"\n[COMPLETE] Full pipeline done!")
        print(f"  SFT model: {sft_path}")
        print(f"  DPO model: {args.dpo_output}")