import torch
import torch.nn as nn
from torch import Tensor


class MCRMSELoss(nn.Module):

    def __init__(self):
        super(MCRMSELoss, self).__init__()
        self.mse = nn.MSELoss(reduction='none')

    def forward(self, y_pred: Tensor, y_true: Tensor):
        """Calculate mean column-wise rmse on columns

        :param y_pred: tensor of shape (bs, 6)
        :param y_true: tensor of shape (bs, 6)
        :return: tensor of shape 0 (scalar with grad)
        """

        mse = self.mse(y_pred, y_true).mean(0)  # column-wise mean
        rmse = torch.sqrt(mse + 1e-7)

        return rmse.mean()

    def class_mcrmse(self, y_pred: Tensor, y_true: Tensor):
        mse = self.mse(y_pred, y_true).mean(0)  # column-wise mean
        rmse = torch.sqrt(mse + 1e-7)

        return rmse.squeeze()