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 HellingerDistance(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)

        sqrt_x = torch.sqrt(x + self.eps)
        sqrt_y = torch.sqrt(y + self.eps)

        diff_sq = torch.square(sqrt_x - sqrt_y)

        sum_sq = torch.sum(diff_sq, dim=[1, 2, 3])

        return torch.sqrt(sum_sq + self.eps) / 1.41421356  # 1/sqrt(2)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.op = HellingerDistance(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 []
