import torch
import torch.nn as nn
from src.utils.geometry_utils import unit_normals


class L1Loss(nn.Module):
    def __init__(self, invalid_mask_weight=0.0):
        super(L1Loss, self).__init__()
        self.name = 'L1'
        self.invalid_mask_weight = invalid_mask_weight

    def forward(self, pred, target, mask):
        loss = nn.functional.l1_loss(pred[mask], target[mask])
        if self.invalid_mask_weight > 0.0:
            invalid_mask = ~mask
            if invalid_mask.sum() > 0:
                invalid_loss = nn.functional.l1_loss(pred[invalid_mask], target[invalid_mask])
                loss = loss + self.invalid_mask_weight * invalid_loss
        return loss



class GradL1Loss(nn.Module):
    def __init__(self):
        super().__init__()
        self.name = 'GradL1'

    def grad(self, x):
        dx = x[..., :-1, 1:] - x[..., :-1, :-1]
        dy = x[..., 1:, :-1] - x[..., :-1, :-1]
        return dx, dy

    def grad_mask(self, mask):
        return (mask[..., :-1, :-1] & mask[..., :-1, 1:] &
                mask[..., 1:, :-1] & mask[..., 1:, 1:])

    def forward(self, pred, target, mask):
        dx_p, dy_p = self.grad(pred)
        dx_t, dy_t = self.grad(target)
        mask_g = self.grad_mask(mask)

        loss_x = nn.functional.l1_loss(dx_p[mask_g], dx_t[mask_g], reduction='mean')
        loss_y = nn.functional.l1_loss(dy_p[mask_g], dy_t[mask_g], reduction='mean')

        return 0.5 * (loss_x + loss_y)


class CosineNormalLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.name = "CosineNormalLoss"

    def forward(self, pred: torch.Tensor,
                      target: torch.Tensor,
                      mask: torch.Tensor) -> torch.Tensor:
        assert pred.shape == target.shape, "pred and target must have same shape"

        pred = unit_normals(pred)
        target = unit_normals(target)

        dot = (pred * target).sum(dim=1, keepdim=True).clamp(-1.0, 1.0)
        cos_term = 1.0 - dot

        if mask is not None:
            loss = cos_term[mask].mean()
        else:
            loss = cos_term.mean()
        return loss