import torch
import torch.nn as nn
from typing import Optional

class CrossTokenShuffle(nn.Module):
    """Replace bag tokens with real-image donor tokens using vectorized tensor operations.
    
    Real bags draw donors only from other real images. Generated bags may draw
    from the full real-token pool. min_keep_original limits replacement, and
    inserted tokens can be detached. The batch shape is preserved.
    """

    def __init__(self):
        super().__init__()

    @torch.no_grad()
    def forward(
        self,
        x: torch.Tensor,              # (B, N, D)
        y: torch.Tensor,              # (B,) 0/1
        replace_ratio: float = 0.5,
        with_replacement: bool = True,
        detach_inserted: bool = True,
        generator: Optional[torch.Generator] = None,
        min_keep_original: int = 0,
    ) -> torch.Tensor:
        """Return augmented tokens with the same shape as x, using the configured donor and replacement rules."""
        assert x.dim() == 3, "x must be (B, N, D)"
        B, N, D = x.shape
        if B == 0 or N == 0:
            return x

        device = x.device
        dtype = x.dtype

        
        K = int(N * replace_ratio)
        mk = int(min_keep_original)
        if mk > 0:
            K = min(K, max(0, N - mk))
        if K <= 0:
            return x  

        
        x_flat = x.reshape(B * N, D)
        y = y.to(device).view(B)  # (B,)
        token_labels = y.repeat_interleave(N)  # (B*N,)

        
        
        real_flat = torch.nonzero(token_labels == 0, as_tuple=False).view(-1)
        M = int(real_flat.numel())
        if M == 0:
            return x.clone()

        owner = real_flat // N  
        b_ar = torch.arange(B, device=device, dtype=torch.long).unsqueeze(1)  # (B, 1)
        rows_fake = (y == 1).unsqueeze(1)  # (B, 1)
        
        allowed_m = rows_fake | (owner.unsqueeze(0) != b_ar)  # (B, M)

        has_cand = allowed_m.any(dim=1)  # (B,)

        
        
        r_pos = torch.rand(B, N, device=device, generator=generator)
        replace_pos_idx = r_pos.argsort(dim=1)[:, :K].long()  

        
        
        
        

        
        cand_flat_idx = torch.zeros(B, K, device=device, dtype=torch.long)

        if with_replacement:
            rows_with = has_cand  
            if rows_with.any():
                weights = allowed_m[rows_with].float()  # (Rw, M)
                idx_m = torch.multinomial(
                    weights, num_samples=K, replacement=True, generator=generator
                )  # (Rw, K) in [0..M-1]
                cand_flat_idx[rows_with] = real_flat[idx_m]
            
        else:
            scores = torch.where(
                allowed_m,
                torch.zeros(B, M, device=device, dtype=torch.float32),
                torch.full((B, M), float("-inf"), device=device, dtype=torch.float32),
            )
            U = torch.rand(B, M, device=device, generator=generator).clamp_(1e-6, 1 - 1e-6)
            g = -torch.log(-torch.log(U))
            cand_m = (scores + g).topk(k=K, dim=1).indices  # (B, K) in [0..M-1]
            cand_flat_idx = real_flat[cand_m]

        
        inserted_tokens = x_flat.index_select(0, cand_flat_idx.reshape(-1)).reshape(B, K, D)
        if detach_inserted:
            inserted_tokens = inserted_tokens.detach()

        
        x_out = x.clone()
        if has_cand.any():
            
            rows = torch.where(has_cand)[0]                                   
            pos_idx = replace_pos_idx.index_select(0, rows)                   # (Rw, K)
            batch_idx = rows.unsqueeze(1).expand(-1, pos_idx.size(1))         
            ins = inserted_tokens.index_select(0, rows)                       # (Rw, K, D)

            
            x_out[batch_idx, pos_idx, :] = ins                                
            

        
        return x_out.to(dtype=dtype, device=device)


# -------------------------
# Quick self-check (toy)
# -------------------------
if __name__ == "__main__":
    torch.manual_seed(0)
    B, N, D = 256, 256, 1024
    x = torch.arange(B*N*D, dtype=torch.float32).reshape(B, N, D)
    y = torch.tensor([0, 1, 0, 1]*64)  # real, fake, real, fake

    aug = CrossTokenShuffle()
    out = aug(x, y, replace_ratio=0.5, with_replacement=False, detach_inserted=True)

    
    assert out.shape == x.shape
    
    
    print("OK, shape:", out.shape)