"""
FastTRL PPOTrainer - Proximal Policy Optimization
==================================================
RLHF için tam PPO implementasyonu. Value head, GAE, KL kontrolü dahil.
"""

from __future__ import annotations

import os
import math
import warnings
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 PPOConfig
from fasttrl.models import AutoModelForCausalLMWithValueHead, create_reference_model
from fasttrl.utils import (
    get_logger,
    set_seed,
    logprobs_from_logits,
    entropy_from_logits,
    get_kl_penalty,
    compute_gae,
    masked_mean,
    masked_whiten,
    clip_by_value,
    normalize_reward,
    scale_reward,
    clip_reward,
    print_trainable_parameters,
    get_dtype,
    flatten_dict,
    stack_dicts,
)

logger = get_logger(__name__)


class AdaptiveKLController:
    """
    KL ıraksama için adaptif kontrol.
    Hedef KL'ye göre KL katsayısını otomatik ayarlar.
    """

    def __init__(self, init_kl_coef: float, target: float, horizon: int):
        self.value = init_kl_coef
        self.target = target
        self.horizon = horizon

    def update(self, current_kl: float, n_steps: int) -> None:
        """KL katsayısını günceller."""
        proportional_error = (current_kl / self.target) - 1.0
        mult = 1.0 + proportional_error * n_steps / self.horizon
        self.value *= mult
        self.value = max(1e-6, self.value)  # Negatife düşme koruması


class FixedKLController:
    """Sabit KL katsayısı."""

    def __init__(self, kl_coef: float):
        self.value = kl_coef

    def update(self, *args, **kwargs) -> None:
        pass


class PPOTrainer:
    """
    PPO (RLHF) Trainer.

    Pipeline:
    1. Policy model prompt'ları alır → yanıt üretir
    2. Reward modeli yanıtları puanlar
    3. Value model avantaj tahmini yapar (GAE)
    4. PPO clipped surrogate objective ile policy güncellenir
    5. KL sapması kontrol edilir

    Kullanım::

        from fasttrl import PPOTrainer, PPOConfig
        from fasttrl.models import AutoModelForCausalLMWithValueHead
        from transformers import AutoTokenizer, pipeline

        model = AutoModelForCausalLMWithValueHead.from_pretrained("gpt2")
        tokenizer = AutoTokenizer.from_pretrained("gpt2")
        reward_fn = pipeline("sentiment-analysis", model="distilbert...")

        config = PPOConfig(
            batch_size=64,
            mini_batch_size=16,
            ppo_epochs=4,
        )

        trainer = PPOTrainer(
            model=model,
            tokenizer=tokenizer,
            args=config,
        )

        for batch in dataloader:
            queries = batch["input_ids"]
            responses = trainer.generate(queries)
            rewards = [reward_fn(r) for r in responses]
            stats = trainer.step(queries, responses, rewards)
    """

    def __init__(
        self,
        model: Union[AutoModelForCausalLMWithValueHead, str],
        tokenizer,
        args: PPOConfig,
        ref_model: Optional[nn.Module] = None,
        reward_model: Optional[Callable] = None,
        dataset=None,
        data_collator=None,
        peft_config=None,
    ) -> None:
        self.args = args
        self.tokenizer = tokenizer
        self.reward_model = reward_model
        self.dataset = dataset

        set_seed(args.seed)

        # Model
        if isinstance(model, str):
            model = AutoModelForCausalLMWithValueHead.from_pretrained(
                model,
                peft_config=peft_config,
                torch_dtype=get_dtype(args.fp16, args.bf16),
                device_map=args.device_map,
            )
        elif not isinstance(model, AutoModelForCausalLMWithValueHead):
            # Sade LM verilmişse value head ekle
            warnings.warn(
                "AutoModelForCausalLMWithValueHead verilmesi önerilir. "
                "Otomatik dönüştürülüyor."
            )
            model = AutoModelForCausalLMWithValueHead(model)

        self.model = model
        self.device = next(model.parameters()).device

        # Tokenizer
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        tokenizer.padding_side = "left"

        # Referans model
        self.ref_model = (
            ref_model
            if ref_model is not None
            else create_reference_model(self.model.pretrained_model)
        )
        self.ref_model.eval()

        # KL kontrolcü
        if args.adap_kl_ctrl:
            self.kl_ctl = AdaptiveKLController(
                args.init_kl_coef, args.target_kl, args.horizon
            )
        else:
            self.kl_ctl = FixedKLController(args.init_kl_coef)

        # Optimizer
        self.optimizer = self._create_optimizer()

        # İstatistikler
        self._stats_history: List[Dict] = []
        self._global_step = 0

        print_trainable_parameters(self.model)
        logger.info(
            f"PPOTrainer hazır. "
            f"KL kontrolü: {'adaptif' if args.adap_kl_ctrl else 'sabit'} "
            f"(başlangıç={args.init_kl_coef})"
        )

    # ------------------------------------------------------------------
    # Optimizer
    # ------------------------------------------------------------------

    def _create_optimizer(self) -> torch.optim.Optimizer:
        params = [p for p in self.model.parameters() if p.requires_grad]
        try:
            return torch.optim.AdamW(
                params, lr=self.args.learning_rate, fused=True
            )
        except TypeError:
            return torch.optim.AdamW(params, lr=self.args.learning_rate)

    # ------------------------------------------------------------------
    # Generasyon
    # ------------------------------------------------------------------

    @torch.no_grad()
    def generate(
        self,
        query_tensors: Union[Tensor, List[Tensor]],
        return_prompt: bool = False,
        **generation_kwargs,
    ) -> List[Tensor]:
        """
        Query tensörlerinden yanıt üretir.

        Args:
            query_tensors:     [batch, prompt_len] veya liste
            return_prompt:     True → prompt+yanıt döner, False → sadece yanıt
            **generation_kwargs: ek generation parametreleri

        Returns:
            Yanıt tensörleri listesi
        """
        args = self.args

        gen_kwargs = {
            "max_new_tokens": args.max_new_tokens,
            "temperature": args.temperature,
            "top_k": args.top_k,
            "top_p": args.top_p,
            "do_sample": args.do_sample,
            "pad_token_id": self.tokenizer.pad_token_id,
            "eos_token_id": self.tokenizer.eos_token_id,
        }
        gen_kwargs.update(generation_kwargs)

        # Listeyi tensor'a çevir
        if isinstance(query_tensors, list):
            query_tensors = self._pad_queries(query_tensors)

        query_tensors = query_tensors.to(self.device)

        self.model.eval()
        outputs = self.model.generate(
            input_ids=query_tensors,
            attention_mask=(query_tensors != self.tokenizer.pad_token_id).long(),
            **gen_kwargs,
        )

        responses = []
        for i, (query, output) in enumerate(zip(query_tensors, outputs)):
            query_len = (query != self.tokenizer.pad_token_id).sum()
            if return_prompt:
                responses.append(output)
            else:
                responses.append(output[query_len:])

        return responses

    def _pad_queries(self, query_list: List[Tensor]) -> Tensor:
        """Sorgu listesini sol-padding ile tensor'a dönüştürür."""
        max_len = max(q.size(0) for q in query_list)
        pad_id = self.tokenizer.pad_token_id

        padded = []
        for q in query_list:
            pad_len = max_len - q.size(0)
            padded.append(
                torch.cat([
                    torch.full((pad_len,), pad_id, dtype=torch.long),
                    q,
                ])
            )
        return torch.stack(padded)

    # ------------------------------------------------------------------
    # PPO adımı
    # ------------------------------------------------------------------

    def step(
        self,
        queries: List[Tensor],
        responses: List[Tensor],
        scores: List[Tensor],
        response_masks: Optional[List[Tensor]] = None,
    ) -> Dict[str, float]:
        """
        Tek PPO adımı. Queries, responses ve reward skorlarını alır,
        politikayı günceller.

        Args:
            queries:        Prompt tensörleri listesi
            responses:      Yanıt tensörleri listesi
            scores:         Skalar ödül tensörleri listesi
            response_masks: İsteğe bağlı yanıt token maskeleri

        Returns:
            İstatistik sözlüğü
        """
        bs = self.args.batch_size
        assert len(queries) == len(responses) == len(scores), (
            f"queries ({len(queries)}), responses ({len(responses)}), "
            f"scores ({len(scores)}) aynı uzunlukta olmalı."
        )

        self.model.train()

        # Ödülleri işle
        scores_tensor = torch.stack(scores).float().to(self.device)

        if self.args.use_score_scaling:
            scores_tensor = scale_reward(scores_tensor)
        if self.args.use_score_norm:
            scores_tensor = normalize_reward(scores_tensor)
        if self.args.score_clip is not None:
            scores_tensor = clip_reward(scores_tensor, self.args.score_clip)

        # Tüm girdileri birleştir
        model_inputs = self._prepare_model_inputs(queries, responses)
        model_inputs_list = [
            {k: v[i] for k, v in model_inputs.items()}
            for i in range(len(queries))
        ]

        # İlk değerler ve log-prob'lar (eski policy)
        all_logprobs, all_ref_logprobs, all_values, all_masks = (
            self._get_initial_logprobs(model_inputs_list)
        )

        # Token bazında ödülleri hesapla (sadece son token'a skalar ödül)
        rewards, non_score_rewards, kl_rewards = self._compute_rewards(
            scores_tensor, all_logprobs, all_ref_logprobs, all_masks
        )

        # GAE avantajları
        advantages, returns = self._compute_advantages(
            all_values, rewards, all_masks
        )

        # PPO güncelleme döngüsü (ppo_epochs kez)
        stats_list = []
        for _ in range(self.args.ppo_epochs):
            batch_stats = self._ppo_update(
                model_inputs_list,
                all_logprobs,
                all_values,
                advantages,
                returns,
                all_masks,
            )
            stats_list.append(batch_stats)

        # KL güncelle
        mean_kl = torch.cat(
            [r.sum(-1) for r in kl_rewards]
        ).mean().item()
        self.kl_ctl.update(mean_kl, bs)

        self._global_step += 1

        # İstatistikleri topla
        stats = self._aggregate_stats(
            stats_list, scores_tensor, kl_rewards, advantages
        )
        self._stats_history.append(stats)

        if self._global_step % self.args.logging_steps == 0:
            logger.info(
                f"Adım {self._global_step} | "
                f"Loss: {stats.get('ppo/loss/total', 0):.4f} | "
                f"KL: {stats.get('ppo/mean_non_score_reward', 0):.4f} | "
                f"Ödül: {stats.get('ppo/mean_scores', 0):.4f}"
            )

        return stats

    def _prepare_model_inputs(
        self,
        queries: List[Tensor],
        responses: List[Tensor],
    ) -> Dict[str, Tensor]:
        """Query + response birleşimi için model girişi hazırlar."""
        pad_id = self.tokenizer.pad_token_id

        combined = [
            torch.cat([q.to(self.device), r.to(self.device)])
            for q, r in zip(queries, responses)
        ]
        max_len = max(c.size(0) for c in combined)

        input_ids = torch.stack([
            F.pad(c, (0, max_len - c.size(0)), value=pad_id)
            for c in combined
        ])
        attention_mask = (input_ids != pad_id).long()

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "query_lengths": torch.tensor([q.size(0) for q in queries]),
            "response_lengths": torch.tensor([r.size(0) for r in responses]),
        }

    @torch.no_grad()
    def _get_initial_logprobs(
        self,
        model_inputs_list: List[Dict],
    ) -> Tuple[List[Tensor], List[Tensor], List[Tensor], List[Tensor]]:
        """Batch içindeki her örnek için başlangıç değerlerini hesaplar."""
        all_logprobs = []
        all_ref_logprobs = []
        all_values = []
        all_masks = []

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

        for inp in model_inputs_list:
            ids = inp["input_ids"].unsqueeze(0).to(self.device)
            mask = inp["attention_mask"].unsqueeze(0).to(self.device)
            q_len = inp["query_lengths"].item()
            r_len = inp["response_lengths"].item()

            # Policy + Value
            out = self.model(ids, attention_mask=mask)
            logits = out.logits[0, :-1]
            values = out.value[0, :-1]
            labels = ids[0, 1:]

            logprobs = logprobs_from_logits(
                logits.unsqueeze(0), labels.unsqueeze(0)
            ).squeeze(0)

            # Referans model
            ref_ids = ids.to(ref_device)
            ref_mask = mask.to(ref_device)
            ref_out = self.ref_model(
                ref_ids,
                attention_mask=ref_mask,
                output_hidden_states=False,
            )
            ref_logits = ref_out.logits[0, :-1].to(self.device)
            ref_logprobs = logprobs_from_logits(
                ref_logits.unsqueeze(0), labels.unsqueeze(0)
            ).squeeze(0)

            # Response maskesi (sadece response token'ları için)
            seq_len = ids.size(1) - 1
            start = max(0, q_len - 1)
            end = min(start + r_len, seq_len)
            response_mask = torch.zeros(seq_len, device=self.device)
            response_mask[start:end] = 1.0

            all_logprobs.append(logprobs)
            all_ref_logprobs.append(ref_logprobs)
            all_values.append(values)
            all_masks.append(response_mask)

        self.model.train()
        return all_logprobs, all_ref_logprobs, all_values, all_masks

    def _compute_rewards(
        self,
        scores: Tensor,
        logprobs: List[Tensor],
        ref_logprobs: List[Tensor],
        masks: List[Tensor],
    ) -> Tuple[List[Tensor], List[Tensor], List[Tensor]]:
        """Token bazında ödülleri hesaplar (KL cezası dahil)."""
        rewards = []
        non_score_rewards = []
        kl_rewards = []

        for i, (lp, ref_lp, mask) in enumerate(zip(logprobs, ref_logprobs, masks)):
            kl = get_kl_penalty(lp, ref_lp, self.args.kl_penalty) * mask
            kl_rewards.append(kl)
            non_score_reward = -self.kl_ctl.value * kl
            non_score_rewards.append(non_score_reward)

            # Skalar ödülü son response token'ına ekle
            reward = non_score_reward.clone()
            last_non_pad = (mask == 1).nonzero(as_tuple=True)[0]
            if len(last_non_pad) > 0:
                reward[last_non_pad[-1]] += scores[i]

            if self.args.whiten_rewards:
                reward = masked_whiten(reward, mask)

            rewards.append(reward)

        return rewards, non_score_rewards, kl_rewards

    def _compute_advantages(
        self,
        values: List[Tensor],
        rewards: List[Tensor],
        masks: List[Tensor],
    ) -> Tuple[Tensor, Tensor]:
        """GAE ile avantaj ve getiri hesaplar."""
        # Stack (padding gerekebilir)
        max_len = max(v.size(0) for v in values)

        def pad(t, length):
            return F.pad(t, (0, length - t.size(0)))

        values_t = torch.stack([pad(v, max_len) for v in values])
        rewards_t = torch.stack([pad(r, max_len) for r in rewards])
        masks_t = torch.stack([pad(m, max_len) for m in masks])

        advantages, returns = compute_gae(
            values_t, rewards_t, masks_t,
            gamma=self.args.gamma, lam=self.args.lam
        )
        return advantages, returns

    def _ppo_update(
        self,
        model_inputs_list: List[Dict],
        old_logprobs: List[Tensor],
        old_values: List[Tensor],
        advantages: Tensor,
        returns: Tensor,
        masks: List[Tensor],
    ) -> Dict[str, float]:
        """Mini-batch PPO güncellemesi."""
        args = self.args
        n = len(model_inputs_list)
        mini_bs = args.mini_batch_size

        all_pg_loss = []
        all_vf_loss = []
        all_entropy = []
        all_kl = []

        # Mini-batch'lere böl
        indices = torch.randperm(n)

        for start in range(0, n, mini_bs):
            batch_idx = indices[start : start + mini_bs]

            for i in batch_idx.tolist():
                inp = model_inputs_list[i]
                ids = inp["input_ids"].unsqueeze(0).to(self.device)
                mask_in = inp["attention_mask"].unsqueeze(0).to(self.device)

                with torch.cuda.amp.autocast(enabled=args.fp16 or args.bf16):
                    out = self.model(ids, attention_mask=mask_in)

                logits = out.logits[0, :-1]
                new_values = out.value[0, :-1]
                labels = ids[0, 1:]
                resp_mask = masks[i]

                new_logprobs = logprobs_from_logits(
                    logits.unsqueeze(0), labels.unsqueeze(0)
                ).squeeze(0)

                old_lp = old_logprobs[i]
                old_v = old_values[i]
                adv = advantages[i, : old_lp.size(0)]
                ret = returns[i, : old_lp.size(0)]

                # Normalize avantajlar
                adv = (adv - adv[resp_mask.bool()].mean()) / (
                    adv[resp_mask.bool()].std() + 1e-8
                ) if resp_mask.sum() > 1 else adv

                # Policy gradient loss (clipped)
                ratio = torch.exp(new_logprobs - old_lp)
                pg_loss1 = -adv * ratio
                pg_loss2 = -adv * ratio.clamp(
                    1 - args.cliprange, 1 + args.cliprange
                )
                pg_loss = masked_mean(
                    torch.max(pg_loss1, pg_loss2), resp_mask
                )

                # Value function loss (clipped)
                vf_loss1 = (new_values - ret) ** 2
                vf_clipped = clip_by_value(
                    new_values,
                    old_v - args.cliprange_value,
                    old_v + args.cliprange_value,
                )
                vf_loss2 = (vf_clipped - ret) ** 2
                vf_loss = 0.5 * masked_mean(
                    torch.max(vf_loss1, vf_loss2), resp_mask
                )

                # Entropy bonus
                entropy = masked_mean(
                    entropy_from_logits(logits), resp_mask
                )

                # Toplam kayıp
                loss = (
                    pg_loss
                    + args.vf_coef * vf_loss
                    - args.entropy_coef * entropy
                )

                loss.backward()

                all_pg_loss.append(pg_loss.item())
                all_vf_loss.append(vf_loss.item())
                all_entropy.append(entropy.item())
                all_kl.append(
                    masked_mean(old_lp - new_logprobs, resp_mask).item()
                )

            # Gradient güncellemesi
            if args.max_grad_norm:
                torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(), args.max_grad_norm
                )
            self.optimizer.step()
            self.optimizer.zero_grad()

        return {
            "ppo/loss/policy": sum(all_pg_loss) / len(all_pg_loss),
            "ppo/loss/value": sum(all_vf_loss) / len(all_vf_loss),
            "ppo/loss/total": sum(all_pg_loss) / len(all_pg_loss) + sum(all_vf_loss) / len(all_vf_loss),
            "ppo/policy/entropy": sum(all_entropy) / len(all_entropy),
            "ppo/policy/approxkl": sum(all_kl) / len(all_kl),
        }

    def _aggregate_stats(
        self,
        stats_list: List[Dict],
        scores: Tensor,
        kl_rewards: List[Tensor],
        advantages: Tensor,
    ) -> Dict[str, float]:
        merged = {}
        for k in stats_list[0]:
            merged[k] = sum(s[k] for s in stats_list) / len(stats_list)

        merged["ppo/mean_scores"] = scores.mean().item()
        merged["ppo/std_scores"] = scores.std().item()
        merged["ppo/kl_coef"] = self.kl_ctl.value
        merged["ppo/mean_non_score_reward"] = torch.cat(kl_rewards).mean().item()

        return merged

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