import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, bins=32, sigma=0.1):
        super().__init__()
        self.bins = bins

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        b, n = x.shape

        # Min-Max Normalization per sample
        min_x = x.min(dim=1, keepdim=True)[0]
        max_x = x.max(dim=1, keepdim=True)[0]
        min_y = y.min(dim=1, keepdim=True)[0]
        max_y = y.max(dim=1, keepdim=True)[0]

        x_norm = (x - min_x) / (max_x - min_x + 1e-6)
        y_norm = (y - min_y) / (max_y - min_y + 1e-6)

        # Quantize
        x_bin = torch.clamp((x_norm * self.bins).long(), 0, self.bins - 1)
        y_bin = torch.clamp((y_norm * self.bins).long(), 0, self.bins - 1)

        # Joint Histogram
        joint_idx = x_bin * self.bins + y_bin
        # (B, bins^2)
        hist = torch.zeros(b, self.bins * self.bins, device=x.device, dtype=torch.float32)
        ones = torch.ones_like(joint_idx, dtype=torch.float32)
        hist.scatter_add_(1, joint_idx, ones)

        # Normalize to probability
        p_xy = hist / n
        p_xy = p_xy.view(b, self.bins, self.bins)

        # Marginals
        p_x = p_xy.sum(dim=2)  # (B, bins)
        p_y = p_xy.sum(dim=1)  # (B, bins)

        # Mutual Information: sum p_xy * log(p_xy / (p_x * p_y))
        # Add eps to avoid log(0)
        eps = 1e-12
        p_x_p_y = torch.bmm(p_x.unsqueeze(2), p_y.unsqueeze(1))  # (B, bins, bins)

        mi = p_xy * torch.log((p_xy + eps) / (p_x_p_y + eps))
        return mi.sum(dim=(1, 2))


batch_size = 128
feature_dim = 512


def get_inputs():
    x = torch.randn(batch_size, feature_dim, dtype=torch.float32)
    y = torch.randn(batch_size, feature_dim, dtype=torch.float32)
    return [x, y]


def get_init_inputs():
    return [32]