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

class MILHead(nn.Module):
    """Instance classification and MIL bag pooling.
    
    Supports max, log-sum-exp, top-k and gated-attention pooling. d_in is the
    token dimension; num_classes determines the classifier output size.
    Attention can pool instance logits or classify the pooled bag embedding.
    """
    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,
    ):
        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)

        # per-token classifier (shared across tokens)
        self.classifier = nn.Linear(d_in, num_classes)

        # components for gated-attention MIL (Ilse et al., 2018)
        if self.pooling == "attn":
            print(f'Using gated-attention pooling with hidden size {attn_hidden}')
            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)

    @staticmethod
    def _masked_fill_min(x, mask, fill_value=None):
        """Utility: 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: [B, N, C] -> [B, C]
        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):
        # logits: [B, N, C] -> [B, C]
        # temperature T: divide logits by T, smaller T -> sharper
        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)  # [B,1]
                bag = bag - n_eff.log()
        return bag

    def _pool_topk(self, logits, mask=None):
        # logits: [B, N, C] -> [B, C]
        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)  # [B]
        else:
            n_eff = logits.new_full((B,), N)

        k = self.topk
        if k is None:
            # ratio -> at least 1
            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)
            topk_vals.append(vals.mean(dim=-1, keepdim=True))  # [C,1]
        bag = torch.cat(topk_vals, dim=1).transpose(0,1).contiguous().squeeze(1)  # [B, C]
        return bag
    # v1
    # def _pool_attn(self, feats, logits, mask=None):
    #     """
    #     Gated-attention pooling:
    #         a_i ∝ w^T ( tanh(V h_i) ⊙ sigmoid(U h_i) )
    #         bag = Σ a_i * logits_i
    #     feats:  [B, N, D]
    #     logits: [B, N, C]
    #     mask:   [B, N] (True for valid)
    #     Returns:
    #         bag_logits: [B, C]
    #         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)     # [B, N, 1]
    #     bag = (attn * logits).sum(dim=1)  # [B, C]
    #     return bag, attn
    
    def _pool_attn(self, feats, logits, mask=None):
        """Pool instance logits and features using normalized gated attention.
        
        The gate is tanh(Vh) times sigmoid(Uh), projected to one attention logit.
        Input features have shape [B, N, D]; instance logits are [B, N, C].
        The optional mask marks valid tokens. Return bag logits, attention and embedding.
        """
        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)        # [B, N, 1]

        
        bag_logits = (attn * logits).sum(dim=1)  # [B, C]

        
        bag_embed = (attn * feats).sum(dim=1)    # [B, D]

        return bag_logits, attn, bag_embed
    def forward(self, feats, mask=None):
        """Classify token features [B, N, D] and pool the bag.
        
        The optional [B, N] mask marks valid tokens. Return instance logits, bag
        logits, attention weights and a bag embedding when the selected mode provides them.
        """
        inst_logits = self.classifier(feats)  # [B, N, C]

        if self.pooling == "max":
            bag_logits = self._pool_max(inst_logits, mask)
            attn = None
        elif self.pooling == "lsep":
            bag_logits = self._pool_lsep(inst_logits, mask)
            attn = None
        elif self.pooling == "topk":
            bag_logits = self._pool_topk(inst_logits, mask)
            attn = None
        elif self.pooling == "attn":
            # v1
            # bag_logits, attn = self._pool_attn(feats, inst_logits, mask)
            # v2
            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 if self.pooling == "attn" else None,
        }