"""
FastTRL Config - Tüm eğitim konfigürasyonları
==============================================
Tüm parametreler tip-kontrollü, açıklamalı ve akıllı default değerli.
"""

from __future__ import annotations

import os
import json
import warnings
from dataclasses import dataclass, field, asdict
from typing import Optional, List, Union, Literal, Dict, Any
from pathlib import Path


def _validate_positive(value: float, name: str) -> None:
    if value <= 0:
        raise ValueError(f"`{name}` pozitif olmalı, verilen: {value}")


def _validate_range(value: float, lo: float, hi: float, name: str) -> None:
    if not (lo <= value <= hi):
        raise ValueError(f"`{name}` [{lo}, {hi}] aralığında olmalı, verilen: {value}")


@dataclass
class BaseConfig:
    """Tüm config'lerin temel sınıfı."""

    output_dir: str = "outputs"
    seed: int = 42
    logging_steps: int = 10
    save_steps: int = 500
    eval_steps: int = 500
    save_total_limit: int = 3
    report_to: List[str] = field(default_factory=lambda: ["tensorboard"])
    run_name: Optional[str] = None
    fp16: bool = False
    bf16: bool = False
    gradient_checkpointing: bool = False
    dataloader_num_workers: int = 4
    dataloader_pin_memory: bool = True
    remove_unused_columns: bool = False
    push_to_hub: bool = False
    hub_model_id: Optional[str] = None
    hub_token: Optional[str] = None
    # Otomatik cihaz tespiti
    device_map: Optional[Union[str, Dict]] = "auto"

    def __post_init__(self) -> None:
        os.makedirs(self.output_dir, exist_ok=True)
        if self.fp16 and self.bf16:
            raise ValueError("fp16 ve bf16 aynı anda etkinleştirilemez.")
        self._validate()

    def _validate(self) -> None:
        """Alt sınıflar override eder."""
        pass

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    def save(self, path: str) -> None:
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)

    @classmethod
    def load(cls, path: str) -> "BaseConfig":
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        return cls(**data)

    def __repr__(self) -> str:
        lines = [f"{self.__class__.__name__}("]
        for k, v in self.to_dict().items():
            lines.append(f"  {k}={v!r},")
        lines.append(")")
        return "\n".join(lines)


# ---------------------------------------------------------------------------
# SFT Config
# ---------------------------------------------------------------------------

@dataclass
class SFTConfig(BaseConfig):
    """
    Supervised Fine-Tuning konfigürasyonu.

    Kullanım::

        config = SFTConfig(
            model_name="meta-llama/Llama-3-8B",
            output_dir="./sft_output",
            max_seq_length=2048,
            num_train_epochs=3,
            per_device_train_batch_size=4,
            learning_rate=2e-5,
        )
    """

    model_name: str = "gpt2"
    dataset_text_field: str = "text"
    max_seq_length: int = 2048
    packing: bool = True  # Sequence packing → daha hızlı eğitim
    num_train_epochs: int = 3
    per_device_train_batch_size: int = 4
    per_device_eval_batch_size: int = 8
    gradient_accumulation_steps: int = 1
    learning_rate: float = 2e-5
    weight_decay: float = 0.01
    warmup_ratio: float = 0.03
    lr_scheduler_type: str = "cosine"
    optim: str = "adamw_torch_fused"  # Fused optimizer = %20 hız artışı
    max_grad_norm: float = 1.0
    neftune_noise_alpha: Optional[float] = None  # NEFTune gürültü artırımı
    completion_only: bool = False  # Sadece completion token'larında loss hesapla
    response_template: Optional[str] = None  # completion_only için template
    chat_template: Optional[str] = None  # Özel chat template
    dataset_kwargs: Dict[str, Any] = field(default_factory=dict)
    eval_packing: bool = False

    def _validate(self) -> None:
        _validate_positive(self.learning_rate, "learning_rate")
        _validate_positive(self.max_seq_length, "max_seq_length")
        _validate_range(self.warmup_ratio, 0.0, 1.0, "warmup_ratio")
        if self.completion_only and self.response_template is None:
            raise ValueError(
                "completion_only=True iken `response_template` belirtilmeli."
            )


# ---------------------------------------------------------------------------
# PPO Config
# ---------------------------------------------------------------------------

@dataclass
class PPOConfig(BaseConfig):
    """
    Proximal Policy Optimization konfigürasyonu.

    Kullanım::

        config = PPOConfig(
            model_name="gpt2",
            learning_rate=1e-5,
            ppo_epochs=4,
            batch_size=64,
        )
    """

    model_name: str = "gpt2"
    learning_rate: float = 1e-5
    batch_size: int = 64  # Kaç prompt her PPO adımında işlenir
    mini_batch_size: int = 16  # PPO güncelleme mini-batch boyutu
    gradient_accumulation_steps: int = 1
    ppo_epochs: int = 4
    max_grad_norm: Optional[float] = 1.0
    # Clipping
    cliprange: float = 0.2
    cliprange_value: float = 0.2
    # Value function
    vf_coef: float = 0.1
    # Entropy bonus
    entropy_coef: float = 0.0
    # KL kontrolü
    kl_penalty: Literal["kl", "abs", "mse", "full"] = "kl"
    init_kl_coef: float = 0.2
    adap_kl_ctrl: bool = True
    target_kl: float = 6.0
    horizon: int = 10_000
    # Generasyon
    max_new_tokens: int = 256
    temperature: float = 1.0
    top_k: int = 0
    top_p: float = 1.0
    do_sample: bool = True
    # GAE
    gamma: float = 1.0
    lam: float = 0.95
    # Reward normalizasyonu
    use_score_scaling: bool = False
    use_score_norm: bool = False
    score_clip: Optional[float] = None
    whiten_rewards: bool = False
    # Diğer
    optim: str = "adamw_torch_fused"

    def _validate(self) -> None:
        _validate_positive(self.learning_rate, "learning_rate")
        _validate_range(self.cliprange, 0.0, 1.0, "cliprange")
        _validate_range(self.cliprange_value, 0.0, 1.0, "cliprange_value")
        if self.mini_batch_size > self.batch_size:
            raise ValueError(
                f"mini_batch_size ({self.mini_batch_size}) "
                f"batch_size ({self.batch_size})'dan büyük olamaz."
            )
        if self.batch_size % self.mini_batch_size != 0:
            raise ValueError(
                f"batch_size ({self.batch_size}), "
                f"mini_batch_size ({self.mini_batch_size})'a tam bölünebilmeli."
            )


# ---------------------------------------------------------------------------
# DPO Config
# ---------------------------------------------------------------------------

@dataclass
class DPOConfig(BaseConfig):
    """
    Direct Preference Optimization konfigürasyonu.

    Kullanım::

        config = DPOConfig(
            model_name="gpt2",
            beta=0.1,
            loss_type="sigmoid",
        )
    """

    model_name: str = "gpt2"
    beta: float = 0.1  # KL düzenleyici katsayı
    loss_type: Literal[
        "sigmoid",      # Orijinal DPO
        "hinge",        # Hinge loss
        "ipo",          # IPO
        "kto_pair",     # KTO çiftli
        "bco_pair",     # BCO çiftli
        "robust",       # Robust DPO
        "exo_pair",     # EXO çiftli
        "nca_pair",     # NCA çiftli
        "sppo_hard",    # SPPO hard
        "aot",          # AOT
        "aot_pair",     # AOT çiftli
        "apo_zero",     # APO zero
        "apo_down",     # APO down
    ] = "sigmoid"
    label_smoothing: float = 0.0
    max_length: int = 1024
    max_prompt_length: int = 512
    max_target_length: Optional[int] = None
    # Referans model
    reference_free: bool = False  # True = referans modelsiz DPO
    precompute_ref_log_probs: bool = False  # Hız için ref log-prob önhesaplama
    sync_ref_model: bool = False
    ref_model_mixup_alpha: float = 0.9
    ref_model_sync_steps: int = 64
    # Eğitim
    num_train_epochs: int = 3
    per_device_train_batch_size: int = 4
    per_device_eval_batch_size: int = 8
    gradient_accumulation_steps: int = 1
    learning_rate: float = 5e-7
    weight_decay: float = 0.0
    warmup_ratio: float = 0.1
    lr_scheduler_type: str = "cosine"
    optim: str = "adamw_torch_fused"
    max_grad_norm: float = 1.0
    # RPO
    rpo_alpha: Optional[float] = None  # RPO ek SFT loss ağırlığı
    # Truncation
    truncation_mode: Literal["keep_end", "keep_start"] = "keep_end"

    def _validate(self) -> None:
        _validate_positive(self.beta, "beta")
        _validate_range(self.label_smoothing, 0.0, 0.5, "label_smoothing")
        _validate_positive(self.learning_rate, "learning_rate")
        if self.max_prompt_length >= self.max_length:
            raise ValueError(
                f"max_prompt_length ({self.max_prompt_length}) "
                f"max_length ({self.max_length})'den küçük olmalı."
            )


# ---------------------------------------------------------------------------
# Reward Config
# ---------------------------------------------------------------------------

@dataclass
class RewardConfig(BaseConfig):
    """
    Reward Model eğitim konfigürasyonu.

    Kullanım::

        config = RewardConfig(
            model_name="gpt2",
            max_length=512,
        )
    """

    model_name: str = "gpt2"
    max_length: int = 512
    num_train_epochs: int = 3
    per_device_train_batch_size: int = 8
    per_device_eval_batch_size: int = 16
    gradient_accumulation_steps: int = 1
    learning_rate: float = 1e-5
    weight_decay: float = 0.001
    warmup_ratio: float = 0.03
    lr_scheduler_type: str = "linear"
    optim: str = "adamw_torch_fused"
    max_grad_norm: float = 1.0
    # Margin loss
    margin: float = 0.0  # Chosen - rejected margin
    # Center rewards_mean üzerinden normalizasyon
    center_rewards_coefficient: Optional[float] = None

    def _validate(self) -> None:
        _validate_positive(self.learning_rate, "learning_rate")
        _validate_positive(self.max_length, "max_length")


# ---------------------------------------------------------------------------
# GRPO Config
# ---------------------------------------------------------------------------

@dataclass
class GRPOConfig(BaseConfig):
    """
    Group Relative Policy Optimization konfigürasyonu.

    Kullanım::

        config = GRPOConfig(
            model_name="gpt2",
            num_generations=8,
        )
    """

    model_name: str = "gpt2"
    num_generations: int = 8  # Her prompt için üretilen yanıt sayısı (G)
    max_new_tokens: int = 256
    temperature: float = 0.9
    top_p: float = 1.0
    top_k: int = 0
    # Eğitim
    num_train_epochs: int = 1
    per_device_train_batch_size: int = 1
    gradient_accumulation_steps: int = 8
    learning_rate: float = 1e-6
    weight_decay: float = 0.0
    warmup_ratio: float = 0.03
    lr_scheduler_type: str = "cosine_with_restarts"
    optim: str = "adamw_torch_fused"
    max_grad_norm: float = 0.1
    max_prompt_length: int = 512
    # GRPO parametreleri
    epsilon: float = 0.2  # Clipping
    epsilon_high: Optional[float] = None  # Asimetrik clipping (opsiyonel)
    beta: float = 0.04  # KL katsayı
    # Reward normalizasyon
    normalize_reward: bool = True
    reward_clip: Optional[float] = None
    # Referans modeli
    sync_ref_model: bool = True
    ref_model_sync_steps: int = 20

    def _validate(self) -> None:
        _validate_positive(self.num_generations, "num_generations")
        _validate_positive(self.learning_rate, "learning_rate")
        _validate_range(self.epsilon, 0.0, 1.0, "epsilon")
        if self.epsilon_high is not None and self.epsilon_high < self.epsilon:
            raise ValueError(
                f"epsilon_high ({self.epsilon_high}) >= epsilon ({self.epsilon}) olmalı."
            )
