import torch
import torch.nn as nn

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

assert (C * H * W) % 4 == 0


class ICC(nn.Module):

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

    def forward(self, x, y):
        x_flat = x.reshape(x.size(0), -1)
        y_flat = y.reshape(x.size(0), -1)

        x_mean = x_flat.mean(dim=1, keepdim=True)
        y_mean = y_flat.mean(dim=1, keepdim=True)

        x_centered = x_flat - x_mean
        y_centered = y_flat - y_mean

        cov_sum = (x_centered * y_centered).sum(dim=1)

        x_var_sum = (x_centered ** 2).sum(dim=1)
        y_var_sum = (y_centered ** 2).sum(dim=1)

        numerator = 2 * cov_sum
        denominator = x_var_sum + y_var_sum

        return numerator / (denominator + self.eps)


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

    def forward(self, x, y):
        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():
    return []