import torch
from torch import nn

class PinballLoss(nn.Module):
    def __init__(self, quantiles):
        super(PinballLoss, self).__init__()
        self.quantiles_tensor = torch.tensor(quantiles, dtype=torch.float32)
        
    def forward(self, pred, target):
        error = target - pred
        upper = self.quantiles_tensor * error
        lower = (self.quantiles_tensor - 1) * error 
        losses = torch.max(lower, upper)
        loss = torch.mean(torch.mean(losses, dim=0))
        return loss
    

class NonAutoRegressivePinballLoss(nn.Module):
    def __init__(self, quantiles):
        super(NonAutoRegressivePinballLoss, self).__init__()
        self.quantiles_tensor = torch.tensor(quantiles, dtype=torch.float32)

    def forward(self, pred, target):
        pred = pred.reshape(-1, 96, len(self.quantiles_tensor))
        target_expanded = target.unsqueeze(2)
        error = target_expanded - pred
        upper = self.quantiles_tensor * error
        lower = (self.quantiles_tensor - 1) * error
        losses = torch.max(lower, upper)
        loss = torch.mean(losses)
        return loss
