import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, bins=32, eps=1e-12):
        super().__init__()
        self.bins = bins
        self.eps = eps

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

        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)

        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_idx = x_bin * self.bins + y_bin
        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)

        p_xy = hist.view(b, self.bins, self.bins) / n

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

        h_x = -torch.sum(p_x * torch.log(p_x + self.eps), dim=1)
        h_y = -torch.sum(p_y * torch.log(p_y + self.eps), dim=1)

        p_x_p_y = torch.bmm(p_x.unsqueeze(2), p_y.unsqueeze(1))
        mi = torch.sum(p_xy * torch.log((p_xy + self.eps) / (p_x_p_y + self.eps)), dim=(1, 2))

        return h_x + h_y - 2 * mi


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]