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 MinkowskiDistance(nn.Module):

    def __init__(self, p=2.0, keepdim=False, eps=1e-6):
        super().__init__()
        self.p = p
        self.keepdim = keepdim
        self.eps = eps

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

        pow_diff = torch.pow(diff, self.p)

        sum_pow = torch.sum(pow_diff, dim=[1, 2, 3], keepdim=self.keepdim)

        output = torch.pow(sum_pow + self.eps, 1.0 / self.p)

        return output


class Model(nn.Module):

    def __init__(self, p=3.0):
        super().__init__()
        self.op = MinkowskiDistance(p=p, keepdim=False, eps=EPS)

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


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


def get_init_inputs():
    p = 3.0
    return [p]
