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


class MILHead(nn.Module):
    """
    MIL Head supporting max / LogSumExp / top-k / attention-gated pooling.

    Args:
        d_in (int): token feature dim
        num_classes (int): number of classes
        pooling (str): 'max' | 'lsep' | 'topk' | 'attn'
        lsep_temp (float): temperature for LogSumExp pooling
        lsep_norm (bool): whether to subtract log(N_eff)
        topk (int | None): fixed k for top-k pooling
        topk_ratio (float): ratio for top-k if topk is None
        attn_hidden (int): hidden dim for gated-attention
        attn_dropout (float): dropout for gated-attention
        attn_bag_mode (str): only used when pooling='attn'
            - 'logit_pool': bag_logits = sum(attn * inst_logits)
            - 'embed_cls' : bag_embed = sum(attn * feats), then bag_logits = bag_classifier(bag_embed)
    """

    def __init__(
        self,
        d_in: int = 1024,
        num_classes: int = 2,
        pooling: str = "lsep",
        lsep_temp: float = 1.0,
        lsep_norm: bool = True,
        topk: int = None,
        topk_ratio: float = 0.1,
        attn_hidden: int = 128,
        attn_dropout: float = 0.0,
        attn_bag_mode: str = "embed_cls",
    ):
        super().__init__()
        self.num_classes = num_classes
        self.pooling = pooling.lower()
        self.lsep_temp = float(lsep_temp)
        self.lsep_norm = bool(lsep_norm)
        self.topk = topk
        self.topk_ratio = float(topk_ratio)
        self.attn_bag_mode = attn_bag_mode.lower()

        if self.pooling not in {"max", "lsep", "topk", "attn"}:
            raise ValueError(f"Unknown pooling mode: {self.pooling}")

        if self.attn_bag_mode not in {"logit_pool", "embed_cls"}:
            raise ValueError(f"Unknown attn_bag_mode: {self.attn_bag_mode}")

        # Instance-level classifier: shared across tokens
        self.classifier = nn.Linear(d_in, num_classes)

        # Attention pooling modules
        if self.pooling == "attn":
            print(
                f"Using gated-attention pooling with hidden size {attn_hidden}, "
                f"bag mode = {self.attn_bag_mode}"
            )
            self.attn_V = nn.Linear(d_in, attn_hidden, bias=True)
            self.attn_U = nn.Linear(d_in, attn_hidden, bias=True)
            self.attn_w = nn.Linear(attn_hidden, 1, bias=False)
            self.attn_drop = nn.Dropout(attn_dropout)

            # Decoupled bag classifier
            if self.attn_bag_mode == "embed_cls":
                self.bag_classifier = nn.Linear(d_in, num_classes)
            else:
                self.bag_classifier = None
        else:
            self.bag_classifier = None

    @staticmethod
    def _masked_fill_min(x, mask, fill_value=None):
        """Set masked positions to a very small number before max/topk."""
        if mask is None:
            return x
        if fill_value is None:
            fill_value = torch.finfo(x.dtype).min
        return x.masked_fill(~mask, fill_value)

    def _pool_max(self, logits, mask=None):
        logits = self._masked_fill_min(
            logits, mask.unsqueeze(-1) if mask is not None else None
        )
        bag = logits.max(dim=1).values
        return bag

    def _pool_lsep(self, logits, mask=None):
        z = logits / self.lsep_temp
        if mask is not None:
            neg_inf = torch.finfo(z.dtype).min
            z = z.masked_fill(~mask.unsqueeze(-1), neg_inf)

        bag = torch.logsumexp(z, dim=1)

        if self.lsep_norm:
            if mask is None:
                n = logits.size(1)
                bag = bag - math.log(max(n, 1))
            else:
                n_eff = mask.sum(dim=1).clamp_min(1).unsqueeze(-1)
                bag = bag - n_eff.log()

        return bag

    def _pool_topk(self, logits, mask=None):
        B, N, C = logits.shape

        if mask is not None:
            logits = self._masked_fill_min(logits, mask.unsqueeze(-1))
            n_eff = mask.sum(dim=1)
        else:
            n_eff = logits.new_full((B,), N, dtype=torch.long)

        k = self.topk
        if k is None:
            k = torch.clamp((n_eff.float() * self.topk_ratio).long(), min=1)
        else:
            k = torch.clamp(torch.full_like(n_eff, k), min=1, max=N)

        scores = logits.transpose(1, 2)  # [B, C, N]
        topk_vals = []
        for b in range(B):
            kb = int(min(int(k[b].item()), int(n_eff[b].item())))
            vals, _ = scores[b].topk(kb, dim=-1)      # [C, kb]
            topk_vals.append(vals.mean(dim=-1))       # [C]

        bag = torch.stack(topk_vals, dim=0)           # [B, C]
        return bag

    def _compute_attn(self, feats, mask=None):
        """
        feats: [B, N, D]
        mask : [B, N]
        returns:
            attn: [B, N, 1]
        """
        Vh = self.attn_V(feats)                       # [B, N, L]
        Uh = self.attn_U(feats)                       # [B, N, L]
        gate = torch.tanh(Vh) * torch.sigmoid(Uh)
        gate = self.attn_drop(gate)
        attn = self.attn_w(gate)                      # [B, N, 1]

        if mask is not None:
            neg_inf = torch.finfo(attn.dtype).min
            attn = attn.masked_fill(~mask.unsqueeze(-1), neg_inf)

        attn = F.softmax(attn, dim=1)
        return attn

    def _pool_attn(self, feats, logits, mask=None):
        """
        Returns:
            bag_logits: [B, C]
            attn:       [B, N, 1]
            bag_embed:  [B, D]
        """
        attn = self._compute_attn(feats, mask)         # [B, N, 1]
        bag_embed = (attn * feats).sum(dim=1)          # [B, D]

        if self.attn_bag_mode == "logit_pool":
            bag_logits = (attn * logits).sum(dim=1)    # [B, C]
        elif self.attn_bag_mode == "embed_cls":
            bag_logits = self.bag_classifier(bag_embed)  # [B, C]
        else:
            raise ValueError(f"Unknown attn_bag_mode: {self.attn_bag_mode}")

        return bag_logits, attn, bag_embed

    def forward(self, feats, mask=None):
        """
        Args:
            feats: [B, N, D]
            mask:  [B, N], True means valid token
        Returns:
            dict{
                'inst_logits': [B, N, C],
                'bag_logits':  [B, C],
                'attn':        [B, N, 1] or None,
                'bag_embed':   [B, D] or None
            }
        """
        inst_logits = self.classifier(feats)  # [B, N, C]

        attn = None
        bag_embed = None

        if self.pooling == "max":
            bag_logits = self._pool_max(inst_logits, mask)

        elif self.pooling == "lsep":
            bag_logits = self._pool_lsep(inst_logits, mask)

        elif self.pooling == "topk":
            bag_logits = self._pool_topk(inst_logits, mask)

        elif self.pooling == "attn":
            bag_logits, attn, bag_embed = self._pool_attn(feats, inst_logits, mask)

        else:
            raise ValueError(f"Unknown pooling mode: {self.pooling}")

        return {
            "inst_logits": inst_logits,
            "bag_logits": bag_logits,
            "attn": attn,
            "bag_embed": bag_embed,
        }