import torch
import torch.nn as nn
import torch.nn.functional as F

N, C, H, W = 32, 64, 56, 56
EPS = 1e-6


class BhattacharyyaDistance(nn.Module):

    def __init__(self, eps=1e-6):
        super().__init__()
        self.eps = eps

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        x = torch.relu(x)
        y = torch.relu(y)

        product = x * y
        sqrt_prod = torch.sqrt(product + self.eps)

        bc = torch.sum(sqrt_prod, dim=[1, 2, 3])

        distance = -torch.log(bc + self.eps)

        return distance


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.op = BhattacharyyaDistance(EPS)

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return self.op(x, y)


def get_inputs():
    x = torch.rand(N, C, H, W, dtype=torch.float32)
    y = torch.rand(N, C, H, W, dtype=torch.float32)
    return [x, y]


def get_init_inputs():
    return []
