"""
FastTRL RewardTrainer - Ödül modeli eğitimi
============================================
"""

from __future__ import annotations

import os
import math
from typing import Any, Callable, Dict, List, Optional, 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 RewardConfig
from fasttrl.utils import (
    get_logger,
    set_seed,
    print_trainable_parameters,
    get_dtype,
    compute_perplexity,
)

logger = get_logger(__name__)


class RewardTrainer:
    """
    Ödül Modeli Trainer.

    Bradley-Terry modeli: P(chosen > rejected) = σ(r_chosen - r_rejected)

    Veri formatı::

        {
            "chosen":   "<tercih edilen yanıt>",
            "rejected": "<reddedilen yanıt>",
            "prompt":   "<istem (opsiyonel)>",
        }

    Kullanım::

        from fasttrl import RewardTrainer, RewardConfig
        from transformers import AutoModelForSequenceClassification, AutoTokenizer

        model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=1)
        tokenizer = AutoTokenizer.from_pretrained("gpt2")

        config = RewardConfig(output_dir="./reward_output")

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

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

        set_seed(args.seed)

        self.model = self._load_model(model, peft_config)
        self._fix_tokenizer()

        self.train_dataset = self._tokenize_dataset(train_dataset)
        self.eval_dataset = (
            self._tokenize_dataset(eval_dataset) if eval_dataset is not None else None
        )

        self.data_collator = data_collator or self._default_collator
        self.optimizer = None
        self.lr_scheduler = None
        self._global_step = 0

        print_trainable_parameters(self.model)

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

            dtype = get_dtype(self.args.fp16, self.args.bf16)
            model = AutoModelForSequenceClassification.from_pretrained(
                model,
                num_labels=1,
                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
            model = get_peft_model(model, peft_config)

        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 = "right"

    def _tokenize_dataset(self, dataset) -> Any:
        required = {"chosen", "rejected"}
        if not required.issubset(set(dataset.column_names)):
            # prompt+chosen / prompt+rejected formatına bak
            if "prompt" in dataset.column_names:
                pass  # ok
            else:
                raise ValueError(
                    f"Veri seti 'chosen' ve 'rejected' sütunlarını gerektiriyor. "
                    f"Mevcut: {dataset.column_names}"
                )

        def tokenize(example):
            prompt = example.get("prompt", "")
            chosen_text = prompt + example["chosen"]
            rejected_text = prompt + example["rejected"]

            tok_chosen = self.tokenizer(
                chosen_text,
                truncation=True,
                max_length=self.args.max_length,
                padding=False,
            )
            tok_rejected = self.tokenizer(
                rejected_text,
                truncation=True,
                max_length=self.args.max_length,
                padding=False,
            )

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

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

    def _default_collator(self, features: List[Dict]) -> Dict[str, Tensor]:
        """Padding ve batch oluşturma."""
        pad_id = self.tokenizer.pad_token_id or 0
        max_len = max(
            max(len(f["chosen_input_ids"]) for f in features),
            max(len(f["rejected_input_ids"]) for f in features),
        )

        def pad_seq(seq, max_l):
            return seq + [pad_id] * (max_l - len(seq))

        def pad_mask(mask, max_l):
            return mask + [0] * (max_l - len(mask))

        return {
            "chosen_input_ids": torch.tensor(
                [pad_seq(f["chosen_input_ids"], max_len) for f in features]
            ),
            "chosen_attention_mask": torch.tensor(
                [pad_mask(f["chosen_attention_mask"], max_len) for f in features]
            ),
            "rejected_input_ids": torch.tensor(
                [pad_seq(f["rejected_input_ids"], max_len) for f in features]
            ),
            "rejected_attention_mask": torch.tensor(
                [pad_mask(f["rejected_attention_mask"], max_len) for f in features]
            ),
        }

    def _compute_reward_loss(
        self,
        chosen_rewards: Tensor,
        rejected_rewards: Tensor,
    ) -> Tensor:
        """Bradley-Terry kayıp + opsiyonel margin."""
        margin = self.args.margin
        if margin > 0:
            loss = -F.logsigmoid(chosen_rewards - rejected_rewards - margin)
        else:
            loss = -F.logsigmoid(chosen_rewards - rejected_rewards)

        # Center rewards regularization
        if self.args.center_rewards_coefficient is not None:
            c = self.args.center_rewards_coefficient
            center_loss = c * ((chosen_rewards + rejected_rewards) ** 2).mean()
            loss = loss.mean() + center_loss
        else:
            loss = loss.mean()

        return loss

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

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

        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(
            args.lr_scheduler_type,
            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
        self.model.train()

        logger.info(f"Reward eğitimi başlıyor: {total_steps} adım")

        total_loss = 0.0
        total_acc = 0.0
        n_batches = 0

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

                with torch.cuda.amp.autocast(enabled=args.fp16 or args.bf16):
                    chosen_out = self.model(
                        input_ids=batch["chosen_input_ids"],
                        attention_mask=batch["chosen_attention_mask"],
                    )
                    rejected_out = self.model(
                        input_ids=batch["rejected_input_ids"],
                        attention_mask=batch["rejected_attention_mask"],
                    )

                chosen_rewards = chosen_out.logits.squeeze(-1)
                rejected_rewards = rejected_out.logits.squeeze(-1)

                loss = self._compute_reward_loss(chosen_rewards, rejected_rewards)
                loss = loss / args.gradient_accumulation_steps

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

                total_loss += loss.item() * args.gradient_accumulation_steps
                acc = (chosen_rewards > rejected_rewards).float().mean().item()
                total_acc += acc
                n_batches += 1

                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

                    if self._global_step % args.logging_steps == 0:
                        avg_loss = total_loss / n_batches
                        avg_acc = total_acc / n_batches
                        logger.info(
                            f"Adım {self._global_step} | "
                            f"Loss: {avg_loss:.4f} | "
                            f"Doğruluk: {avg_acc:.4f}"
                        )

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

        self._save_checkpoint("final_model")

        result = {
            "train_loss": total_loss / n_batches,
            "train_accuracy": total_acc / n_batches,
        }
        logger.info(f"Reward 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}")
