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


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

    def forward(self, x, y):
        C1 = 0.01 ** 2
        C2 = 0.03 ** 2

        mu_x = x.mean(dim=1, keepdim=True)
        mu_y = y.mean(dim=1, keepdim=True)

        sigma_x = ((x - mu_x) ** 2).mean(dim=1, keepdim=True)
        sigma_y = ((y - mu_y) ** 2).mean(dim=1, keepdim=True)
        sigma_xy = ((x - mu_x) * (y - mu_y)).mean(dim=1, keepdim=True)

        numerator = (2 * mu_x * mu_y + C1) * (2 * sigma_xy + C2)
        denominator = (mu_x ** 2 + mu_y ** 2 + C1) * (sigma_x + sigma_y + C2)

        ssim = numerator / denominator
        return F.softplus(ssim).mean()


batch_size = 16
input_dim = 1024


def get_inputs():
    x = torch.rand(batch_size, input_dim)
    y = torch.rand(batch_size, input_dim)
    return [x, y]


def get_init_inputs():
    return []