"""rlhf_mapping.py — V6: learnable possibility mapping + RLHF.

Implementa (item do PowerMachine/gru-ring-v13-9-2):
  - **Learnable possibility mapping**: mapeia logits do modelo para uma
    distribuição de "possibilidades" aprendível. Substitui o softmax fixo
    por uma transformação aprendível que pode capturar distribuições não
    uniformes (heavy-tailed, multimodais).
  - **RLHF integration**: integra feedback humano (preference pairs) via
    DPO-style loss. Diferente do DPO clássico, mantém a possibility mapping
    aprendível durante o alignment.

Matemática
──────────
  Possibility Mapping:
    p_poss(y|x) = softmax(T_poss * logits + b_poss)
    onde T_poss (temperatura aprendível) e b_poss (bias aprendível) são
    parâmetros treináveis. T_poss permite que o modelo aprenda a "calibrar"
    sua própria confiança.

  RLHF (DPO-style):
    L_DPO = -log σ(β * (log p_θ(y_w|x) / p_ref(y_w|x)
                       - log p_θ(y_l|x) / p_ref(y_l|x)))
    onde y_w = winner, y_l = loser, β = temperatura DPO.
"""
from __future__ import annotations
import math
from typing import Optional, Tuple, Dict, List
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class RLHFConfig:
    """Configuração RLHF V6."""
    # Learnable possibility mapping
    init_temp: float = 1.0
    init_bias: float = 0.0
    temp_min: float = 0.1
    temp_max: float = 10.0
    # DPO
    dpo_beta: float = 0.1
    dpo_loss_weight: float = 0.5
    # PPO-like penalty
    kl_penalty_weight: float = 0.01


class LearnablePossibilityMapping(nn.Module):
    """V6: mapeamento aprendível de logits para possibilidades.

    p_poss(y|x) = softmax(T_poss * logits + b_poss)

    T_poss é aprendível mas clipado em [temp_min, temp_max] via sigmoid.
    """
    def __init__(
        self,
        init_temp: float = 1.0,
        init_bias: float = 0.0,
        temp_min: float = 0.1,
        temp_max: float = 10.0,
    ):
        super().__init__()
        self.temp_min = temp_min
        self.temp_max = temp_max
        # Parametrização: T = temp_min + (temp_max - temp_min) * sigmoid(raw_T)
        # Inicializa raw_T tal que T(0) = init_temp
        init_raw = math.log(
            (init_temp - temp_min) / max(temp_max - init_temp, 1e-8)
        ) if temp_min < init_temp < temp_max else 0.0
        self.raw_temp = nn.Parameter(torch.tensor(float(init_raw)))
        self.bias = nn.Parameter(torch.tensor(float(init_bias)))

    @property
    def temperature(self) -> torch.Tensor:
        return self.temp_min + (self.temp_max - self.temp_min) * torch.sigmoid(self.raw_temp)

    def forward(self, logits: torch.Tensor) -> torch.Tensor:
        """logits: (..., V) → possibilities: (..., V) — distribuição suavizada."""
        T = self.temperature
        return F.softmax(T * logits + self.bias, dim=-1)


class RLHFMapping(nn.Module):
    """V6: RLHF com learnable possibility mapping.

    Forward (preference):
        logits_w: (B, V) — logits do winner
        logits_l: (B, V) — logits do loser
        → dpo_loss: scalar, metrics: dict

    Forward (standard):
        logits: (B, V) → possibilities: (B, V)
    """
    def __init__(self, config: Optional[RLHFConfig] = None):
        super().__init__()
        self.cfg = config or RLHFConfig()
        self.possibility = LearnablePossibilityMapping(
            init_temp=self.cfg.init_temp,
            init_bias=self.cfg.init_bias,
            temp_min=self.cfg.temp_min,
            temp_max=self.cfg.temp_max,
        )

    def forward(self, logits: torch.Tensor) -> torch.Tensor:
        """Standard forward: logits → possibilities."""
        return self.possibility(logits)

    def dpo_loss(
        self,
        logits_w: torch.Tensor,
        logits_l: torch.Tensor,
        targets_w: torch.Tensor,
        targets_l: torch.Tensor,
        ref_logits_w: Optional[torch.Tensor] = None,
        ref_logits_l: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Dict[str, float]]:
        """Computa DPO loss.

        Se ref_logits não fornecido, usa logits_w/l como referência (toy mode).
        """
        # Log probs do winner e loser
        log_p_w = F.log_softmax(logits_w, dim=-1)
        log_p_l = F.log_softmax(logits_l, dim=-1)
        log_p_w_target = log_p_w.gather(-1, targets_w.unsqueeze(-1)).squeeze(-1).sum(-1)
        log_p_l_target = log_p_l.gather(-1, targets_l.unsqueeze(-1)).squeeze(-1).sum(-1)

        if ref_logits_w is None or ref_logits_l is None:
            # Toy mode: assume ref = uniform
            ref_log_p_w = torch.full_like(log_p_w_target, -math.log(logits_w.size(-1)))
            ref_log_p_l = torch.full_like(log_p_l_target, -math.log(logits_l.size(-1)))
        else:
            ref_log_p_w = F.log_softmax(ref_logits_w, dim=-1).gather(
                -1, targets_w.unsqueeze(-1)
            ).squeeze(-1).sum(-1)
            ref_log_p_l = F.log_softmax(ref_logits_l, dim=-1).gather(
                -1, targets_l.unsqueeze(-1)
            ).squeeze(-1).sum(-1)

        # DPO: -log σ(β * (log p_w/p_ref_w - log p_l/p_ref_l))
        delta = self.cfg.dpo_beta * (
            (log_p_w_target - ref_log_p_w) - (log_p_l_target - ref_log_p_l)
        )
        loss = -F.logsigmoid(delta).mean()

        metrics = {
            "dpo_loss": float(loss.item()),
            "dpo_delta": float(delta.mean().item()),
            "temperature": float(self.possibility.temperature.item()),
            "bias": float(self.possibility.bias.item()),
            "accuracy": float((delta > 0).float().mean().item()),
        }
        return self.cfg.dpo_loss_weight * loss, metrics

    def kl_penalty(
        self,
        current_logits: torch.Tensor,
        ref_logits: torch.Tensor,
    ) -> torch.Tensor:
        """KL(p_current || p_ref) — penalty PPO-style."""
        p_cur = F.softmax(current_logits, dim=-1)
        log_p_cur = F.log_softmax(current_logits, dim=-1)
        log_p_ref = F.log_softmax(ref_logits, dim=-1)
        kl = (p_cur * (log_p_cur - log_p_ref)).sum(dim=-1).mean()
        return self.cfg.kl_penalty_weight * kl


__all__ = ["RLHFMapping", "LearnablePossibilityMapping", "RLHFConfig"]