"""Calibration head used at inference (temperature as a function of the logits and the pooled answer-slot state)."""
import math
from typing import Optional

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


LO, HI = math.log(0.5), math.log(5.0)
MODES = {"yesno": 0, "digits": 1, "letters": 2, "parallel": 3}
N_FEAT = 3 + 4 + 5


# ------------------------------------------------------------------ features -----------------------------------------------------------

def make_features(logits: torch.Tensor, mask: torch.Tensor, qtype: torch.Tensor, mode: torch.Tensor) -> torch.Tensor:
    """(N, Kmax) padded logits -> (N, N_FEAT) features that do not depend on which option is which."""
    z = logits.masked_fill(~mask, -1e4)
    k = mask.sum(-1).clamp(min=2).float()
    logp = torch.log_softmax(z, -1)
    p = logp.exp()
    ent = -(p * logp).sum(-1) / torch.log(k)
    top = z.topk(2, -1).values
    mean = (logits * mask).sum(-1) / k
    std = (((logits - mean[:, None]) ** 2 * mask).sum(-1) / k).sqrt()
    return torch.cat([F.one_hot(qtype, 3).float(), F.one_hot(mode, 4).float(), torch.stack([torch.log(k) / 5, ent, p.max(-1).values, (top[:, 0] - top[:, 1]) / 10, std / 10], -1)], -1)


class SwiGLU(nn.Module):
    def __init__(self, d_in: int, d_hid: int, d_out: int):
        super().__init__()
        self.gate, self.up, self.down = nn.Linear(d_in, d_hid), nn.Linear(d_in, d_hid), nn.Linear(d_hid, d_out)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))


class TempHead(nn.Module):
    """log T = LO + (HI - LO) * sigmoid(raw + c) with c chosen so that raw = 0 gives T = 1; the last layer starts at zero (T = 1)."""

    def __init__(self, n_feat: int = N_FEAT, hidden: Optional[int] = None, d: int = 128):
        super().__init__()
        self.hidden_dim = hidden
        self.proj = nn.Sequential(nn.LayerNorm(hidden), nn.Linear(hidden, d)) if hidden else None
        self.ffn = SwiGLU(n_feat + (d if hidden else 0), 2 * d, 1)
        nn.init.zeros_(self.ffn.down.weight); nn.init.zeros_(self.ffn.down.bias)
        self.register_buffer("c", torch.tensor(math.log((0 - LO) / (HI - 0))))

    def log_t(self, feats: torch.Tensor, pooled: Optional[torch.Tensor] = None) -> torch.Tensor:
        x = feats if self.proj is None else torch.cat([feats, self.proj(pooled)], -1)
        return LO + (HI - LO) * torch.sigmoid(self.ffn(x).squeeze(-1) + self.c)


# ------------------------------------------------------------------ cache --------------------------------------------------------------