"""Pooled feature extraction and the ternary scoring rule.

The classifier reads one 768-D vector per image: layernorm across the 768
channels of every patch token, then max-pool across the 2304 patches. Score is
the sum of the person-positive dims minus the sum of the person-negative dims.
"""
from typing import Optional, Sequence

import torch
import torch.nn.functional as F

D = 768
RES = 768


def pool(patch_tokens: torch.Tensor) -> torch.Tensor:
    """(N, D) or (B, N, D) patch tokens -> (D,) or (B, D) pooled vector."""
    ln = F.layer_norm(patch_tokens.float(), [D])
    return ln.max(dim=-2).values


@torch.inference_mode()
def backbone_pooled(backbone, x: torch.Tensor, autocast: bool = True) -> torch.Tensor:
    """Forward a normalized batch through the backbone and pool it."""
    if autocast:
        dev = 'cuda' if x.is_cuda else 'cpu'
        with torch.autocast(dev, dtype=torch.bfloat16):
            out = backbone.forward_features(x)
    else:
        out = backbone.forward_features(x)
    return pool(out['x_norm_patchtokens'].float())


def score(pooled: torch.Tensor, pos: Sequence[int], neg: Sequence[int]) -> torch.Tensor:
    """sum(pooled[pos]) - sum(pooled[neg]), over the last axis."""
    if not torch.is_tensor(pos):
        pos = torch.tensor(list(pos), dtype=torch.long, device=pooled.device)
    if not torch.is_tensor(neg):
        neg = torch.tensor(list(neg), dtype=torch.long, device=pooled.device)
    return pooled.index_select(-1, pos).sum(-1) - pooled.index_select(-1, neg).sum(-1)


def score_pool(backbone, loaded, pos, neg, target_dims: Optional[torch.Tensor] = None):
    """Score a pool; with `target_dims`, also return the pooled activations there."""
    scores, targets = [], []
    for x in loaded:
        pooled = backbone_pooled(backbone, x)[0]
        scores.append(score(pooled, pos, neg))
        if target_dims is not None:
            targets.append(pooled[target_dims])
    stacked = torch.stack(scores)
    return (stacked, torch.stack(targets)) if target_dims is not None else (stacked, None)