"""
FastTRL GRPOTrainer - Group Relative Policy Optimization
=========================================================
DeepSeek-R1'de kullanılan GRPO implementasyonu.
Value head gerektirmez — grup içi normalize ödüllerle çalışır.
"""

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 GRPOConfig
from fasttrl.models import create_reference_model
from fasttrl.utils import (
    get_logger,
    set_seed,
    logprobs_from_logits,
    get_kl_penalty,
    masked_mean,
    print_trainable_parameters,
    get_dtype,
)

logger = get_logger(__name__)


class GRPOTrainer:
    """
    Group Relative Policy Optimization Trainer.

    Her prompt için G yanıt üretilir, grup içi normalize ödüller
    avantaj olarak kullanılır (value head gerekmez).

    Kayıp = -E[min(r·A, clip(r, 1-ε, 1+ε)·A)] + β·KL

    Kullanım::

        from fasttrl import GRPOTrainer, GRPOConfig
        from transformers import AutoModelForCausalLM, AutoTokenizer
        from datasets import load_dataset

        model = AutoModelForCausalLM.from_pretrained("gpt2")
        tokenizer = AutoTokenizer.from_pretrained("gpt2")

        def reward_fn(completions, **kwargs):
            # Basit uzunluk ödülü
            return [len(c) / 100 for c in completions]

        config = GRPOConfig(
            output_dir="./grpo_output",
            num_generations=4,
            max_new_tokens=128,
        )

        trainer = GRPOTrainer(
            model=model,
            tokenizer=tokenizer,
            args=config,
            reward_funcs=[reward_fn],
            train_dataset=dataset,
        )
        trainer.train()
    """

    def __init__(
        self,
        model: Union[nn.Module, str],
        tokenizer,
        args: GRPOConfig,
        reward_funcs: List[Callable],
        train_dataset,
        eval_dataset=None,
        ref_model: Optional[nn.Module] = None,
        peft_config=None,
    ) -> None:
        self.args = args
        self.tokenizer = tokenizer
        self.reward_funcs = reward_funcs if isinstance(reward_funcs, list) else [reward_funcs]

        set_seed(args.seed)

        # Model yükle
        self.model = self._load_model(model, peft_config)
        self.device = next(self.model.parameters()).device

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

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

        # Veri seti
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset

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

        print_trainable_parameters(self.model)
        logger.info(
            f"GRPOTrainer hazır. "
            f"G={args.num_generations}, β={args.beta}, ε={args.epsilon}"
        )

    def _load_model(self, model, 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:
            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)

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

        return model

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

    @torch.no_grad()
    def _generate_completions(
        self, prompts: List[str]
    ) -> Tuple[List[str], List[Tensor]]:
        """Her prompt için G tane tamamlama üretir."""
        args = self.args
        G = args.num_generations

        all_texts = []
        all_tensors = []

        self.model.eval()

        for prompt in prompts:
            prompt_ids = self.tokenizer(
                prompt,
                return_tensors="pt",
                truncation=True,
                max_length=args.max_prompt_length,
            ).input_ids.to(self.device)

            # G tane yanıt üret
            outputs = self.model.generate(
                prompt_ids.repeat(G, 1),
                max_new_tokens=args.max_new_tokens,
                temperature=args.temperature,
                top_p=args.top_p,
                top_k=args.top_k,
                do_sample=True,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
            )

            for out in outputs:
                completion_ids = out[prompt_ids.size(1):]
                text = self.tokenizer.decode(
                    completion_ids, skip_special_tokens=True
                )
                all_texts.append(text)
                all_tensors.append(completion_ids)

        return all_texts, all_tensors

    # ------------------------------------------------------------------
    # Ödül hesaplama
    # ------------------------------------------------------------------

    def _compute_rewards(
        self,
        prompts: List[str],
        completions: List[str],
    ) -> Tensor:
        """Tüm reward fonksiyonlarının toplamını hesaplar."""
        G = self.args.num_generations
        n_prompts = len(prompts)

        # Her prompt için G tane tamamlama
        expanded_prompts = [p for p in prompts for _ in range(G)]

        total_rewards = torch.zeros(n_prompts * G)

        for reward_fn in self.reward_funcs:
            try:
                rewards = reward_fn(
                    completions,
                    prompts=expanded_prompts,
                )
            except TypeError:
                rewards = reward_fn(completions)

            if isinstance(rewards, list):
                rewards = torch.tensor(rewards, dtype=torch.float)
            elif isinstance(rewards, Tensor):
                rewards = rewards.float().cpu()

            if rewards.shape[0] != n_prompts * G:
                raise ValueError(
                    f"Reward fonksiyonu {n_prompts * G} değer döndürmeli "
                    f"({n_prompts} prompt × {G} yanıt), "
                    f"ancak {rewards.shape[0]} döndürdü."
                )

            total_rewards += rewards

        return total_rewards

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

    def _get_logprobs(
        self,
        model: nn.Module,
        input_ids: Tensor,
        attention_mask: Tensor,
        labels: Tensor,
    ) -> Tensor:
        """Sequence log-prob toplamı döner."""
        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]
        labels = labels[:, 1:]
        per_token = logprobs_from_logits(logits, labels)
        mask = (labels != self.tokenizer.pad_token_id).float()
        return (per_token * mask).sum(-1)

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

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

        # Veri setini hazırla
        if hasattr(self.train_dataset, "column_names"):
            prompt_field = next(
                (c for c in ("prompt", "question", "input", "text")
                 if c in self.train_dataset.column_names),
                self.train_dataset.column_names[0],
            )
            prompts_all = self.train_dataset[prompt_field]
        else:
            prompts_all = list(self.train_dataset)

        # Per-prompt batch boyutu
        per_prompt_batch = max(1, args.per_device_train_batch_size)
        total_prompts = len(prompts_all)

        # Optimizer
        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,
        )

        steps_per_epoch = math.ceil(
            (total_prompts / per_prompt_batch) / args.gradient_accumulation_steps
        )
        total_steps = steps_per_epoch * args.num_train_epochs

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

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

        logger.info(
            f"GRPO eğitimi başlıyor: {args.num_train_epochs} epoch, "
            f"{total_steps} adım"
        )

        total_loss = 0.0
        n_updates = 0

        for epoch in range(args.num_train_epochs):
            perm = torch.randperm(total_prompts).tolist()

            for batch_start in range(0, total_prompts, per_prompt_batch):
                batch_idx = perm[batch_start : batch_start + per_prompt_batch]
                prompts = [prompts_all[i] for i in batch_idx]

                # ---- Generasyon (no_grad) ----
                completions, completion_tensors = self._generate_completions(prompts)

                # ---- Ödüller ----
                rewards = self._compute_rewards(prompts, completions)  # [B*G]

                # Normalize rewards
                if args.normalize_reward:
                    B = len(prompts)
                    G = args.num_generations
                    rewards_grouped = rewards.view(B, G)
                    mean = rewards_grouped.mean(dim=1, keepdim=True)
                    std = rewards_grouped.std(dim=1, keepdim=True).clamp(min=1e-8)
                    rewards = ((rewards_grouped - mean) / std).view(-1)

                if args.reward_clip is not None:
                    rewards = rewards.clamp(-args.reward_clip, args.reward_clip)

                # ---- Model girişlerini hazırla ----
                G = args.num_generations
                expanded_prompts = [p for p in prompts for _ in range(G)]
                all_texts = [
                    expanded_prompts[i] + completions[i]
                    for i in range(len(completions))
                ]

                encoding = self.tokenizer(
                    all_texts,
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                    max_length=args.max_prompt_length + args.max_new_tokens,
                ).to(self.device)

                input_ids = encoding["input_ids"]
                attention_mask = encoding["attention_mask"]
                rewards = rewards.to(self.device)

                # ---- Eski log-prob (referans) ----
                with torch.no_grad():
                    old_logprobs = self._get_logprobs(
                        self.model, input_ids, attention_mask, input_ids
                    )
                    ref_ids = input_ids.to(ref_device)
                    ref_mask = attention_mask.to(ref_device)
                    ref_logprobs = self._get_logprobs(
                        self.ref_model, ref_ids, ref_mask, ref_ids
                    ).to(self.device)

                # ---- Policy güncelleme ----
                self.model.train()
                new_logprobs = self._get_logprobs(
                    self.model, input_ids, attention_mask, input_ids
                )

                log_ratio = new_logprobs - old_logprobs
                ratio = torch.exp(log_ratio)

                # Asimetrik clipping
                eps_lo = args.epsilon
                eps_hi = args.epsilon_high if args.epsilon_high is not None else args.epsilon

                clipped_ratio = torch.clamp(ratio, 1 - eps_lo, 1 + eps_hi)
                advantages = rewards  # GAE yok, doğrudan normalize ödüller

                pg_loss = -torch.min(
                    ratio * advantages,
                    clipped_ratio * advantages,
                ).mean()

                # KL ceza
                kl = get_kl_penalty(new_logprobs, ref_logprobs, penalty_type="kl")
                kl_loss = args.beta * kl.mean()

                loss = (pg_loss + kl_loss) / args.gradient_accumulation_steps

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

                total_loss += loss.item() * args.gradient_accumulation_steps
                n_updates += 1

                is_update = (
                    n_updates % args.gradient_accumulation_steps == 0
                )

                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 args.ref_model_sync_steps > 0
                        and self._global_step % args.ref_model_sync_steps == 0
                    ):
                        import copy
                        self.ref_model = copy.deepcopy(self.model)
                        self.ref_model.eval()
                        for p in self.ref_model.parameters():
                            p.requires_grad_(False)
                        logger.info("Referans model senkronize edildi.")

                    if self._global_step % args.logging_steps == 0:
                        avg_loss = total_loss / n_updates
                        avg_reward = rewards.mean().item()
                        logger.info(
                            f"Adım {self._global_step} | "
                            f"Loss: {avg_loss:.4f} | "
                            f"KL: {kl.mean().item():.4f} | "
                            f"Ortalama Ödül: {avg_reward:.4f}"
                        )

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

            logger.info(f"Epoch {epoch+1}/{args.num_train_epochs} tamamlandı.")

        self._save_checkpoint("final_model")

        result = {
            "train_loss": total_loss / max(n_updates, 1),
        }
        logger.info(f"GRPO eğitimi tamamlandı: {result}")
        return result

    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}")
