"""
FastTRL Utils - Yardımcı fonksiyonlar
=====================================
Tüm matematiksel işlemler, yardımcı araçlar ve ölçüm fonksiyonları.
"""

from __future__ import annotations

import os
import random
import logging
import warnings
import time
import functools
from contextlib import contextmanager
from typing import Any, Dict, List, Optional, Tuple, Union, Callable
from collections import defaultdict

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Reproducibilite
# ---------------------------------------------------------------------------

def set_seed(seed: int, deterministic: bool = False) -> None:
    """Tüm RNG'leri aynı anda sıfırlar."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
        torch.use_deterministic_algorithms(True)


# ---------------------------------------------------------------------------
# Tensor işlemleri
# ---------------------------------------------------------------------------

def masked_mean(
    values: Tensor,
    mask: Tensor,
    axis: Optional[int] = None,
    eps: float = 1e-8,
) -> Tensor:
    """Maske uygulanmış ortalama. Sıfır bölünmesini önler."""
    masked = values * mask
    if axis is not None:
        return masked.sum(axis) / (mask.sum(axis) + eps)
    return masked.sum() / (mask.sum() + eps)


def masked_var(
    values: Tensor,
    mask: Tensor,
    unbiased: bool = True,
    eps: float = 1e-8,
) -> Tensor:
    """Maske uygulanmış varyans."""
    mean = masked_mean(values, mask)
    centered = (values - mean) ** 2
    variance = masked_mean(centered, mask)
    if unbiased:
        n = mask.sum()
        if n > 1:
            variance = variance * n / (n - 1)
    return variance


def masked_whiten(
    values: Tensor,
    mask: Tensor,
    shift_mean: bool = True,
    eps: float = 1e-8,
) -> Tensor:
    """Maske uygulanmış beyazlatma (whitening)."""
    mean = masked_mean(values, mask)
    var = masked_var(values, mask)
    whitened = (values - mean) * torch.rsqrt(var + eps)
    if not shift_mean:
        whitened += mean
    return whitened


def clip_by_value(
    x: Tensor,
    lo: Union[float, Tensor],
    hi: Union[float, Tensor],
) -> Tensor:
    """Elementwise kırpma — torch.clamp'ın tensor sınır desteğiyle versiyonu."""
    return torch.max(torch.min(x, hi), lo)


def logprobs_from_logits(
    logits: Tensor,
    labels: Tensor,
    gather: bool = True,
) -> Tensor:
    """
    Logit'lerden log-olasılık hesaplar.

    Args:
        logits: [batch, seq_len, vocab_size]
        labels: [batch, seq_len]
        gather: True → her token'ın log-prob'u döner [batch, seq_len]
                False → tüm vocab log-prob dağılımı [batch, seq_len, vocab]
    """
    logp = F.log_softmax(logits, dim=-1)
    if not gather:
        return logp
    logp = torch.gather(logp, dim=2, index=labels.unsqueeze(2)).squeeze(2)
    return logp


def entropy_from_logits(logits: Tensor) -> Tensor:
    """
    Logit'lerden Shannon entropisi hesaplar.

    Args:
        logits: [..., vocab_size]
    Returns:
        Entropi [...] (vocab boyutu yok edilmiş)
    """
    p = F.softmax(logits, dim=-1)
    logp = F.log_softmax(logits, dim=-1)
    return -(p * logp).sum(dim=-1)


def get_kl_penalty(
    logprob: Tensor,
    ref_logprob: Tensor,
    penalty_type: str = "kl",
    mask: Optional[Tensor] = None,
    eps: float = 1e-8,
) -> Tensor:
    """
    KL ceza hesaplar.

    Desteklenen tipler:
    - "kl"   : KL(pi || ref) = logprob - ref_logprob (yaklaşım)
    - "abs"  : |KL| = |logprob - ref_logprob|
    - "mse"  : (logprob - ref_logprob)²
    - "full" : Tam KL divergence
    """
    if penalty_type == "kl":
        kl = logprob - ref_logprob
    elif penalty_type == "abs":
        kl = (logprob - ref_logprob).abs()
    elif penalty_type == "mse":
        kl = (logprob - ref_logprob).pow(2)
    elif penalty_type == "full":
        # Tam KL: log(p/q) = log(p) - log(q), beklenti p üzerinden
        log_ratio = logprob - ref_logprob
        # p * log(p/q) - (p - q) = tam KL yaklaşımı
        kl = torch.exp(ref_logprob) * (ref_logprob - logprob) + torch.exp(logprob) - torch.exp(ref_logprob)
        kl = kl.clamp(min=0)
    else:
        raise ValueError(
            f"Geçersiz kl_penalty tipi: '{penalty_type}'. "
            f"Biri olmalı: 'kl', 'abs', 'mse', 'full'"
        )

    if mask is not None:
        kl = kl * mask

    return kl


def whiten_advantages(
    advantages: Tensor,
    mask: Optional[Tensor] = None,
    eps: float = 1e-8,
) -> Tensor:
    """Avantaj değerlerini normalize eder."""
    if mask is not None:
        return masked_whiten(advantages, mask, eps=eps)
    mean = advantages.mean()
    std = advantages.std().clamp(min=eps)
    return (advantages - mean) / std


# ---------------------------------------------------------------------------
# GAE (Generalized Advantage Estimation)
# ---------------------------------------------------------------------------

def compute_gae(
    values: Tensor,
    rewards: Tensor,
    mask: Tensor,
    gamma: float = 0.99,
    lam: float = 0.95,
) -> Tuple[Tensor, Tensor]:
    """
    GAE ile avantaj ve getiri hesaplar.

    Args:
        values:  [batch, seq_len] — değer fonksiyonu tahmini
        rewards: [batch, seq_len] — token bazlı ödüller
        mask:    [batch, seq_len] — geçerli token maskesi
        gamma:   İndirim faktörü
        lam:     GAE lambda parametresi

    Returns:
        (advantages, returns) her biri [batch, seq_len]
    """
    batch_size, seq_len = values.shape
    device = values.device

    # Son adımdan geriye doğru hesapla
    advantages = torch.zeros_like(rewards)
    last_gae = torch.zeros(batch_size, device=device)

    # Son değer için sıfır bootstrap
    next_value = torch.zeros(batch_size, device=device)

    for t in reversed(range(seq_len)):
        non_terminal = mask[:, t]
        delta = rewards[:, t] + gamma * next_value * non_terminal - values[:, t]
        last_gae = delta + gamma * lam * non_terminal * last_gae
        advantages[:, t] = last_gae
        next_value = values[:, t]

    returns = advantages + values
    return advantages, returns


# ---------------------------------------------------------------------------
# Reward Processing
# ---------------------------------------------------------------------------

def clip_reward(reward: Tensor, clip: float) -> Tensor:
    """Ödülleri [-clip, clip] aralığına sıkıştırır."""
    return reward.clamp(-clip, clip)


def scale_reward(reward: Tensor) -> Tensor:
    """Ödülleri [0, 1] aralığına normalize eder."""
    rmin = reward.min()
    rmax = reward.max()
    if (rmax - rmin).abs() < 1e-8:
        return reward - rmin
    return (reward - rmin) / (rmax - rmin)


def normalize_reward(reward: Tensor, eps: float = 1e-8) -> Tensor:
    """Ödülleri ortalama 0, std 1'e normalize eder."""
    return (reward - reward.mean()) / (reward.std() + eps)


# ---------------------------------------------------------------------------
# Model yardımcıları
# ---------------------------------------------------------------------------

def print_trainable_parameters(model: nn.Module, verbose: bool = True) -> Dict[str, int]:
    """
    Eğitilebilir ve toplam parametre sayısını gösterir.

    Returns:
        {"trainable": ..., "total": ..., "ratio": ...}
    """
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    ratio = 100 * trainable / total if total > 0 else 0

    info = {
        "trainable": trainable,
        "total": total,
        "ratio": ratio,
    }

    if verbose:
        print(
            f"Eğitilebilir parametreler: {trainable:,} / {total:,} "
            f"({ratio:.2f}%)"
        )

    return info


def get_model_size_mb(model: nn.Module) -> float:
    """Modelin yaklaşık bellek kullanımını MB cinsinden döner."""
    param_size = sum(p.nelement() * p.element_size() for p in model.parameters())
    buffer_size = sum(b.nelement() * b.element_size() for b in model.buffers())
    return (param_size + buffer_size) / 1024 / 1024


def freeze_model(model: nn.Module) -> nn.Module:
    """Tüm parametreleri dondurur."""
    for param in model.parameters():
        param.requires_grad_(False)
    return model


def unfreeze_model(model: nn.Module) -> nn.Module:
    """Tüm parametreleri çözer."""
    for param in model.parameters():
        param.requires_grad_(True)
    return model


def get_layer_parameter_count(model: nn.Module) -> Dict[str, int]:
    """Her katman için parametre sayısını döner."""
    result = {}
    for name, module in model.named_modules():
        params = sum(p.numel() for p in module.parameters(recurse=False))
        if params > 0:
            result[name] = params
    return result


# ---------------------------------------------------------------------------
# Metrik ve değerlendirme
# ---------------------------------------------------------------------------

def compute_accuracy(
    logits: Tensor,
    labels: Tensor,
    ignore_index: int = -100,
) -> float:
    """Token-level doğruluk hesaplar."""
    predictions = logits.argmax(dim=-1)
    valid = labels != ignore_index
    correct = (predictions == labels) & valid
    return (correct.sum() / valid.sum()).item() if valid.sum() > 0 else 0.0


def compute_perplexity(loss: float) -> float:
    """Cross-entropy loss'tan perplexity hesaplar."""
    import math
    return math.exp(min(loss, 100))  # Overflow koruması


# ---------------------------------------------------------------------------
# Dict yardımcıları
# ---------------------------------------------------------------------------

def flatten_dict(
    d: Dict[str, Any],
    parent_key: str = "",
    sep: str = "/",
) -> Dict[str, Any]:
    """İç içe dict'i düzleştirir."""
    items: List = []
    for k, v in d.items():
        new_key = f"{parent_key}{sep}{k}" if parent_key else k
        if isinstance(v, dict):
            items.extend(flatten_dict(v, new_key, sep).items())
        else:
            items.append((new_key, v))
    return dict(items)


def stack_dicts(dicts: List[Dict[str, Any]]) -> Dict[str, List]:
    """Dict listesini liste dict'e dönüştürür."""
    result = defaultdict(list)
    for d in dicts:
        for k, v in d.items():
            result[k].append(v)
    return dict(result)


# ---------------------------------------------------------------------------
# Zamanlama & loglama
# ---------------------------------------------------------------------------

@contextmanager
def timer(name: str = "", logger: Optional[logging.Logger] = None):
    """Kod bloğunun çalışma süresini ölçer."""
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        msg = f"[{name}] {elapsed:.4f}s" if name else f"{elapsed:.4f}s"
        if logger:
            logger.debug(msg)
        else:
            print(msg)


def get_logger(name: str, level: int = logging.INFO) -> logging.Logger:
    """Tutarlı formatlı logger oluşturur."""
    log = logging.getLogger(name)
    if not log.handlers:
        handler = logging.StreamHandler()
        formatter = logging.Formatter(
            "[%(asctime)s] [%(levelname)s] %(name)s: %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S",
        )
        handler.setFormatter(formatter)
        log.addHandler(handler)
    log.setLevel(level)
    return log


# ---------------------------------------------------------------------------
# Retry decorator
# ---------------------------------------------------------------------------

def retry(
    max_attempts: int = 3,
    delay: float = 1.0,
    exceptions: Tuple = (Exception,),
    backoff: float = 2.0,
) -> Callable:
    """
    Başarısız işlemleri otomatik tekrar dener.

    Kullanım::

        @retry(max_attempts=3, delay=0.5)
        def load_model():
            ...
    """
    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            attempt = 0
            wait = delay
            while attempt < max_attempts:
                try:
                    return func(*args, **kwargs)
                except exceptions as e:
                    attempt += 1
                    if attempt >= max_attempts:
                        raise
                    warnings.warn(
                        f"{func.__name__} başarısız (deneme {attempt}/{max_attempts}): "
                        f"{e}. {wait:.1f}s sonra tekrar denenecek..."
                    )
                    time.sleep(wait)
                    wait *= backoff
        return wrapper
    return decorator


# ---------------------------------------------------------------------------
# Device yardımcıları
# ---------------------------------------------------------------------------

def get_best_device() -> torch.device:
    """Kullanılabilir en iyi cihazı döner."""
    if torch.cuda.is_available():
        return torch.device("cuda")
    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def get_dtype(fp16: bool = False, bf16: bool = False) -> torch.dtype:
    """Config'e göre torch dtype döner."""
    if bf16:
        return torch.bfloat16
    elif fp16:
        return torch.float16
    return torch.float32


def move_to_device(
    obj: Union[Tensor, Dict, List],
    device: Union[str, torch.device],
) -> Any:
    """Tensor, dict veya listeden oluşan yapıyı cihaza taşır."""
    if isinstance(obj, Tensor):
        return obj.to(device)
    elif isinstance(obj, dict):
        return {k: move_to_device(v, device) for k, v in obj.items()}
    elif isinstance(obj, (list, tuple)):
        moved = [move_to_device(x, device) for x in obj]
        return type(obj)(moved)
    return obj
