from __future__ import annotations
import numpy as np
import torch
import torch.nn as nn

def mean_pool(hidden_states: np.ndarray, valid_length: int) -> np.ndarray:
    valid_length = max(1, min(valid_length, hidden_states.shape[0]))
    return hidden_states[-valid_length:].mean(axis=0)

class CrossAttentionPool(nn.Module):

    def __init__(self, hidden_dim: int):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.query = nn.Parameter(torch.randn(hidden_dim) * hidden_dim ** (-0.5))

    def forward(self, hidden_states: torch.Tensor, valid_lengths: torch.Tensor) -> torch.Tensor:
        batch, seq_len, dim = hidden_states.shape
        positions = torch.arange(seq_len, device=hidden_states.device).unsqueeze(0)
        start_idx = (seq_len - valid_lengths).unsqueeze(1)
        mask = positions >= start_idx
        scores = hidden_states @ self.query / dim ** 0.5
        scores = scores.masked_fill(~mask, float('-inf'))
        weights = torch.softmax(scores, dim=-1)
        pooled = torch.einsum('bt,btd->bd', weights, hidden_states)
        return pooled