"""
FastTRL DPOTrainer - Direct Preference Optimization
====================================================
Tüm DPO varyantlarını destekleyen, hızlı ve güvenilir implementasyon.
"""

from __future__ import annotations

import os
import math
import warnings
from contextlib import contextmanager
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.utils.data import DataLoader

from fasttrl.config import DPOConfig
from fasttrl.models import create_reference_model
from fasttrl.utils import (
    get_logger,
    set_seed,
    print_trainable_parameters,
    logprobs_from_logits,
    get_dtype,
    compute_perplexity,
)

logger = get_logger(__name__)


class DPOTrainer:
    """
    Direct Preference Optimization Trainer.

    Desteklenen kayıp fonksiyonları:
    - sigmoid  : Orijinal DPO (Rafailov et al., 2023)
    - ipo      : Identity Preference Optimization
    - hinge    : Hinge DPO
    - robust   : Robust DPO
    - kto_pair : KTO çifti
    - bco_pair : BCO çifti
    - apo_zero : APO zero
    - apo_down : APO down

    Kullanım::

        from fasttrl import DPOTrainer, DPOConfig
        from transformers import AutoModelForCausalLM, AutoTokenizer
        from datasets import load_dataset

        model = AutoModelForCausalLM.from_pretrained("gpt2")
        tokenizer = AutoTokenizer.from_pretrained("gpt2")
        dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")

        config = DPOConfig(
            output_dir="./dpo_output",
            beta=0.1,
            loss_type="sigmoid",
        )

        trainer = DPOTrainer(
            model=model,
            tokenizer=tokenizer,
            args=config,
            train_dataset=dataset,
        )
        trainer.train()
    """

    def __init__(
        self,
        model: Union[nn.Module, str],
        tokenizer,
        args: DPOConfig,
        train_dataset,
        eval_dataset=None,
        ref_model: Optional[nn.Module] = None,
        data_collator=None,
        compute_metrics: Optional[Callable] = None,
        peft_config=None,
    ) -> None:
        self.args = args
        self.tokenizer = tokenizer
        self.compute_metrics = compute_metrics
        self._ref_precomputed = {}

        set_seed(args.seed)

        # Model
        self.model = self._load_model(model, peft_config)

        # Tokenizer düzelt
        self._fix_tokenizer()

        # Referans model
        if args.reference_free:
            self.ref_model = None
            logger.info("Referans modelsiz DPO (reference_free=True)")
        elif ref_model is not None:
            self.ref_model = ref_model
        else:
            self.ref_model = create_reference_model(self.model)

        # Veri setlerini tokenize et
        self.train_dataset = self._tokenize_dataset(train_dataset)
        self.eval_dataset = (
            self._tokenize_dataset(eval_dataset) if eval_dataset is not None else None
        )

        # Referans log-prob önhesaplama
        if args.precompute_ref_log_probs and self.ref_model is not None:
            logger.info("Referans log-prob'lar önhesaplanıyor...")
            self.train_dataset = self._precompute_ref_logprobs(self.train_dataset)

        # Collator
        from fasttrl.data import DataCollatorForPreferenceOptimization
        self.data_collator = data_collator or DataCollatorForPreferenceOptimization(
            tokenizer=self.tokenizer,
            max_length=args.max_length,
            max_prompt_length=args.max_prompt_length,
            truncation_mode=args.truncation_mode,
        )

        self.optimizer = None
        self.lr_scheduler = None
        self._global_step = 0

        print_trainable_parameters(self.model)

    # ------------------------------------------------------------------
    # Hazırlık
    # ------------------------------------------------------------------

    def _load_model(self, model: Union[nn.Module, str], peft_config) -> nn.Module:
        if isinstance(model, str):
            from transformers import AutoModelForCausalLM

            dtype = get_dtype(self.args.fp16, self.args.bf16)
            model = AutoModelForCausalLM.from_pretrained(
                model,
                torch_dtype=dtype,
                device_map=self.args.device_map,
                trust_remote_code=True,
            )

        if peft_config is not None:
            try:
                from peft import get_peft_model, prepare_model_for_kbit_training
                if getattr(model, "is_loaded_in_4bit", False):
                    model = prepare_model_for_kbit_training(model)
                model = get_peft_model(model, peft_config)
            except ImportError:
                raise ImportError("pip install peft gerekli.")

        if self.args.gradient_checkpointing:
            if hasattr(model, "enable_input_require_grads"):
                model.enable_input_require_grads()
            model.gradient_checkpointing_enable()

        return model

    def _fix_tokenizer(self) -> None:
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.tokenizer.padding_side = "left"  # DPO için sol padding tercih edilir

    def _tokenize_dataset(self, dataset) -> Any:
        """Tercih veri setini (chosen/rejected) tokenize eder."""
        if dataset is None:
            return None

        required = {"prompt", "chosen", "rejected"}
        cols = set(dataset.column_names)

        if not required.issubset(cols):
            missing = required - cols
            raise ValueError(
                f"DPO veri seti şu sütunları gerektiriyor: {required}. "
                f"Eksik: {missing}. Mevcut: {cols}"
            )

        def tokenize_example(example):
            prompt = example["prompt"]
            chosen = example["chosen"]
            rejected = example["rejected"]

            # Prompt + chosen
            chosen_full = prompt + chosen
            rejected_full = prompt + rejected

            # Tokenize
            tok_chosen = self.tokenizer(
                chosen_full,
                truncation=True,
                max_length=self.args.max_length,
                add_special_tokens=True,
            )
            tok_rejected = self.tokenizer(
                rejected_full,
                truncation=True,
                max_length=self.args.max_length,
                add_special_tokens=True,
            )
            tok_prompt = self.tokenizer(
                prompt,
                truncation=True,
                max_length=self.args.max_prompt_length,
                add_special_tokens=True,
            )

            prompt_len = len(tok_prompt["input_ids"])

            # Labels: prompt kısmını -100 yap
            chosen_labels = [-100] * prompt_len + tok_chosen["input_ids"][prompt_len:]
            rejected_labels = [-100] * prompt_len + tok_rejected["input_ids"][prompt_len:]

            # Uzunluk kontrolü
            chosen_labels = chosen_labels[: self.args.max_length]
            rejected_labels = rejected_labels[: self.args.max_length]

            return {
                "chosen_input_ids": tok_chosen["input_ids"],
                "chosen_attention_mask": tok_chosen["attention_mask"],
                "chosen_labels": chosen_labels,
                "rejected_input_ids": tok_rejected["input_ids"],
                "rejected_attention_mask": tok_rejected["attention_mask"],
                "rejected_labels": rejected_labels,
            }

        dataset = dataset.map(
            tokenize_example,
            desc="DPO tokenizasyonu",
            remove_columns=[
                c for c in dataset.column_names
                if c not in (
                    "chosen_input_ids", "chosen_attention_mask", "chosen_labels",
                    "rejected_input_ids", "rejected_attention_mask", "rejected_labels",
                )
            ],
        )
        return dataset

    # ------------------------------------------------------------------
    # Log-prob hesaplama
    # ------------------------------------------------------------------

    def _get_logprobs(
        self,
        model: nn.Module,
        input_ids: Tensor,
        attention_mask: Tensor,
        labels: Tensor,
    ) -> Tuple[Tensor, Tensor]:
        """
        Modelin seçimli ve reddedilen yanıtlar için log-prob'larını döner.

        Returns:
            (per_token_logprobs, sequence_logprob)
            - per_token: [batch, seq_len]
            - sequence:  [batch] (labels != -100 maskesi ile toplanmış)
        """
        with torch.cuda.amp.autocast(
            enabled=self.args.fp16 or self.args.bf16
        ):
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                return_dict=True,
            )

        logits = outputs.logits[:, :-1]  # Son token'ı hariç tut
        labels_shifted = labels[:, 1:]   # İlk token'ı hariç tut

        # Her token'ın log-prob'u
        per_token_logprobs = logprobs_from_logits(logits, labels_shifted)

        # Mask: sadece gerçek label token'ları
        loss_mask = labels_shifted != -100

        # Sıfır bölünmesini önle
        seq_logprob = (per_token_logprobs * loss_mask).sum(-1)

        return per_token_logprobs, seq_logprob

    @contextmanager
    def _null_ref_context(self):
        """Referans model yoksa (reference_free) null context."""
        yield

    def _precompute_ref_logprobs(self, dataset) -> Any:
        """Referans model log-prob'larını önceden hesaplar."""
        if self.ref_model is None:
            return dataset

        device = next(self.ref_model.parameters()).device
        self.ref_model.eval()

        chosen_logprobs = []
        rejected_logprobs = []

        loader = DataLoader(
            dataset,
            batch_size=self.args.per_device_train_batch_size,
            collate_fn=self.data_collator,
        )

        with torch.no_grad():
            for batch in loader:
                batch = {k: v.to(device) for k, v in batch.items()}

                _, chosen_lp = self._get_logprobs(
                    self.ref_model,
                    batch["chosen_input_ids"],
                    batch["chosen_attention_mask"],
                    batch["chosen_labels"],
                )
                _, rejected_lp = self._get_logprobs(
                    self.ref_model,
                    batch["rejected_input_ids"],
                    batch["rejected_attention_mask"],
                    batch["rejected_labels"],
                )

                chosen_logprobs.extend(chosen_lp.cpu().tolist())
                rejected_logprobs.extend(rejected_lp.cpu().tolist())

        dataset = dataset.add_column("ref_chosen_logprobs", chosen_logprobs)
        dataset = dataset.add_column("ref_rejected_logprobs", rejected_logprobs)
        logger.info("Referans log-prob'lar önhesaplandı.")
        return dataset

    # ------------------------------------------------------------------
    # Kayıp fonksiyonları
    # ------------------------------------------------------------------

    def _compute_loss(
        self,
        policy_chosen_logprob: Tensor,
        policy_rejected_logprob: Tensor,
        ref_chosen_logprob: Tensor,
        ref_rejected_logprob: Tensor,
    ) -> Tuple[Tensor, Tensor, Tensor]:
        """
        Seçilen kayıp fonksiyonuna göre DPO kaybı hesaplar.

        Returns:
            (loss, chosen_reward, rejected_reward)
        """
        beta = self.args.beta
        smoothing = self.args.label_smoothing

        pi_log_ratio = policy_chosen_logprob - policy_rejected_logprob
        ref_log_ratio = ref_chosen_logprob - ref_rejected_logprob
        log_odds = pi_log_ratio - ref_log_ratio

        chosen_reward = beta * (policy_chosen_logprob - ref_chosen_logprob)
        rejected_reward = beta * (policy_rejected_logprob - ref_rejected_logprob)

        loss_type = self.args.loss_type

        if loss_type == "sigmoid":
            # Orijinal DPO
            loss = (
                -F.logsigmoid(beta * log_odds) * (1 - smoothing)
                - F.logsigmoid(-beta * log_odds) * smoothing
            )

        elif loss_type == "ipo":
            # IPO: (log_odds - 1/(2*beta))²
            loss = (log_odds - 1 / (2 * beta)) ** 2

        elif loss_type == "hinge":
            loss = torch.relu(1 - beta * log_odds)

        elif loss_type == "robust":
            # Robust DPO — c değeri label_smoothing kullanıyor
            c = smoothing
            loss = (
                -F.logsigmoid(beta * log_odds) * (1 - c)
                - F.logsigmoid(-beta * log_odds) * c
            ) / (1 - 2 * c)

        elif loss_type in ("kto_pair", "bco_pair"):
            # KTO pair
            chosen_kl = (policy_chosen_logprob - ref_chosen_logprob).mean().clamp(min=0)
            rejected_kl = (policy_rejected_logprob - ref_rejected_logprob).mean().clamp(min=0)
            chosen_loss = 1 - F.sigmoid(beta * (policy_chosen_logprob - ref_chosen_logprob) - beta * chosen_kl)
            rejected_loss = 1 - F.sigmoid(-(beta * (policy_rejected_logprob - ref_rejected_logprob) - beta * rejected_kl))
            loss = (chosen_loss + rejected_loss) / 2

        elif loss_type == "apo_zero":
            # APO zero: maximize chosen, minimize rejected separately
            loss = (
                -F.logsigmoid(beta * (policy_chosen_logprob - ref_chosen_logprob))
                + F.sigmoid(beta * (policy_rejected_logprob - ref_rejected_logprob))
            )

        elif loss_type == "apo_down":
            loss = (
                -F.logsigmoid(beta * (policy_chosen_logprob - ref_chosen_logprob))
                * F.sigmoid(beta * (policy_rejected_logprob - ref_rejected_logprob))
            )

        else:
            raise ValueError(
                f"Desteklenmeyen loss_type: '{loss_type}'. "
                f"Kullanılabilir: sigmoid, ipo, hinge, robust, "
                f"kto_pair, bco_pair, apo_zero, apo_down"
            )

        # RPO: ek SFT terimi
        if self.args.rpo_alpha is not None:
            sft_loss = -policy_chosen_logprob
            loss = loss + self.args.rpo_alpha * sft_loss

        return loss.mean(), chosen_reward.detach(), rejected_reward.detach()

    # ------------------------------------------------------------------
    # Eğitim döngüsü
    # ------------------------------------------------------------------

    def train(self) -> Dict[str, float]:
        """Ana eğitim döngüsü."""
        args = self.args

        train_loader = DataLoader(
            self.train_dataset,
            batch_size=args.per_device_train_batch_size,
            shuffle=True,
            collate_fn=self.data_collator,
            num_workers=args.dataloader_num_workers,
        )

        steps_per_epoch = math.ceil(len(train_loader) / args.gradient_accumulation_steps)
        total_steps = steps_per_epoch * args.num_train_epochs

        # Optimizer & Scheduler
        self.optimizer = torch.optim.AdamW(
            [p for p in self.model.parameters() if p.requires_grad],
            lr=args.learning_rate,
            weight_decay=args.weight_decay,
        )

        from transformers import get_scheduler
        self.lr_scheduler = get_scheduler(
            name=args.lr_scheduler_type,
            optimizer=self.optimizer,
            num_warmup_steps=int(total_steps * args.warmup_ratio),
            num_training_steps=total_steps,
        )

        scaler = torch.cuda.amp.GradScaler() if args.fp16 else None
        device = next(self.model.parameters()).device

        if self.ref_model is not None:
            ref_device = next(self.ref_model.parameters()).device
            self.ref_model.eval()

        self.model.train()

        logger.info(
            f"DPO eğitimi başlıyor: {args.num_train_epochs} epoch, "
            f"loss_type={args.loss_type}, beta={args.beta}"
        )

        total_loss = 0.0
        total_chosen_reward = 0.0
        total_rejected_reward = 0.0
        ref_sync_step = getattr(args, "ref_model_sync_steps", 0)

        for epoch in range(args.num_train_epochs):
            epoch_loss = 0.0

            for step, batch in enumerate(train_loader):
                batch = {k: v.to(device) for k, v in batch.items()}

                # Policy log-prob'lar
                _, policy_chosen_lp = self._get_logprobs(
                    self.model,
                    batch["chosen_input_ids"],
                    batch["chosen_attention_mask"],
                    batch["chosen_labels"],
                )
                _, policy_rejected_lp = self._get_logprobs(
                    self.model,
                    batch["rejected_input_ids"],
                    batch["rejected_attention_mask"],
                    batch["rejected_labels"],
                )

                # Referans log-prob'lar
                if "ref_chosen_logprobs" in batch:
                    ref_chosen_lp = batch["ref_chosen_logprobs"]
                    ref_rejected_lp = batch["ref_rejected_logprobs"]
                elif self.ref_model is not None:
                    with torch.no_grad():
                        _, ref_chosen_lp = self._get_logprobs(
                            self.ref_model,
                            batch["chosen_input_ids"].to(ref_device),
                            batch["chosen_attention_mask"].to(ref_device),
                            batch["chosen_labels"].to(ref_device),
                        )
                        _, ref_rejected_lp = self._get_logprobs(
                            self.ref_model,
                            batch["rejected_input_ids"].to(ref_device),
                            batch["rejected_attention_mask"].to(ref_device),
                            batch["rejected_labels"].to(ref_device),
                        )
                    ref_chosen_lp = ref_chosen_lp.to(device)
                    ref_rejected_lp = ref_rejected_lp.to(device)
                else:
                    # Reference-free: ref = 0
                    ref_chosen_lp = torch.zeros_like(policy_chosen_lp)
                    ref_rejected_lp = torch.zeros_like(policy_rejected_lp)

                loss, chosen_rew, rejected_rew = self._compute_loss(
                    policy_chosen_lp,
                    policy_rejected_lp,
                    ref_chosen_lp,
                    ref_rejected_lp,
                )

                loss = loss / args.gradient_accumulation_steps

                if scaler:
                    scaler.scale(loss).backward()
                else:
                    loss.backward()

                epoch_loss += loss.item() * args.gradient_accumulation_steps
                total_chosen_reward += chosen_rew.mean().item()
                total_rejected_reward += rejected_rew.mean().item()

                is_update = (
                    (step + 1) % args.gradient_accumulation_steps == 0
                    or (step + 1) == len(train_loader)
                )

                if is_update:
                    if args.max_grad_norm:
                        if scaler:
                            scaler.unscale_(self.optimizer)
                        torch.nn.utils.clip_grad_norm_(
                            self.model.parameters(), args.max_grad_norm
                        )

                    if scaler:
                        scaler.step(self.optimizer)
                        scaler.update()
                    else:
                        self.optimizer.step()

                    self.lr_scheduler.step()
                    self.optimizer.zero_grad()
                    self._global_step += 1

                    # Ref model senkronizasyonu
                    if (
                        args.sync_ref_model
                        and ref_sync_step > 0
                        and self._global_step % ref_sync_step == 0
                        and self.ref_model is not None
                    ):
                        self._sync_ref_model()

                    if self._global_step % args.logging_steps == 0:
                        n = self._global_step * args.per_device_train_batch_size
                        avg_margin = (
                            total_chosen_reward - total_rejected_reward
                        ) / self._global_step
                        logger.info(
                            f"Adım {self._global_step}/{total_steps} | "
                            f"Loss: {epoch_loss/(step+1):.4f} | "
                            f"Margin: {avg_margin:.4f}"
                        )

                    if args.save_steps > 0 and self._global_step % args.save_steps == 0:
                        self._save_checkpoint(f"checkpoint-{self._global_step}")

            total_loss += epoch_loss / len(train_loader)
            logger.info(f"Epoch {epoch+1} tamamlandı.")

        self._save_checkpoint("final_model")

        avg_loss = total_loss / args.num_train_epochs
        result = {
            "train_loss": avg_loss,
            "train_chosen_reward": total_chosen_reward / max(self._global_step, 1),
            "train_rejected_reward": total_rejected_reward / max(self._global_step, 1),
        }
        logger.info(f"DPO eğitimi tamamlandı: {result}")
        return result

    def _sync_ref_model(self) -> None:
        """Referans modeli policy ile karıştırır (Polyak average)."""
        alpha = getattr(self.args, "ref_model_mixup_alpha", 0.9)
        with torch.no_grad():
            for p, rp in zip(
                self.model.parameters(), self.ref_model.parameters()
            ):
                rp.data.mul_(alpha).add_(p.data, alpha=1 - alpha)
        logger.debug("Referans model senkronize edildi.")

    def evaluate(self) -> Dict[str, float]:
        """Değerlendirme döngüsü."""
        if self.eval_dataset is None:
            raise ValueError("eval_dataset belirtilmedi.")

        loader = DataLoader(
            self.eval_dataset,
            batch_size=self.args.per_device_eval_batch_size,
            shuffle=False,
            collate_fn=self.data_collator,
        )

        device = next(self.model.parameters()).device
        self.model.eval()
        ref_device = (
            next(self.ref_model.parameters()).device
            if self.ref_model is not None
            else device
        )

        total_loss = 0.0
        total_acc = 0.0
        n = 0

        with torch.no_grad():
            for batch in loader:
                batch = {k: v.to(device) for k, v in batch.items()}

                _, pol_chosen = self._get_logprobs(
                    self.model, batch["chosen_input_ids"],
                    batch["chosen_attention_mask"], batch["chosen_labels"]
                )
                _, pol_rejected = self._get_logprobs(
                    self.model, batch["rejected_input_ids"],
                    batch["rejected_attention_mask"], batch["rejected_labels"]
                )

                if self.ref_model is not None:
                    _, ref_chosen = self._get_logprobs(
                        self.ref_model,
                        batch["chosen_input_ids"].to(ref_device),
                        batch["chosen_attention_mask"].to(ref_device),
                        batch["chosen_labels"].to(ref_device),
                    )
                    _, ref_rejected = self._get_logprobs(
                        self.ref_model,
                        batch["rejected_input_ids"].to(ref_device),
                        batch["rejected_attention_mask"].to(ref_device),
                        batch["rejected_labels"].to(ref_device),
                    )
                    ref_chosen = ref_chosen.to(device)
                    ref_rejected = ref_rejected.to(device)
                else:
                    ref_chosen = torch.zeros_like(pol_chosen)
                    ref_rejected = torch.zeros_like(pol_rejected)

                loss, chosen_rew, rejected_rew = self._compute_loss(
                    pol_chosen, pol_rejected, ref_chosen, ref_rejected
                )
                acc = (chosen_rew > rejected_rew).float().mean()

                total_loss += loss.item()
                total_acc += acc.item()
                n += 1

        results = {
            "eval_loss": total_loss / n,
            "eval_accuracy": total_acc / n,
        }
        logger.info(f"Değerlendirme: {results}")
        self.model.train()
        return results

    def _save_checkpoint(self, name: str) -> None:
        path = os.path.join(self.args.output_dir, name)
        os.makedirs(path, exist_ok=True)
        if hasattr(self.model, "save_pretrained"):
            self.model.save_pretrained(path)
        else:
            torch.save(self.model.state_dict(), os.path.join(path, "model.pt"))
        self.tokenizer.save_pretrained(path)
        logger.info(f"Checkpoint kaydedildi: {path}")
