import torch.nn as nn
import torch
import random
from utils import IOU


class Loss(nn.Module):
    def __init__(self):
        super().__init__()
        self.mse_loss_fn = nn.MSELoss()
        self.bce_loss_fn = nn.BCEWithLogitsLoss()
        self.entropy_loss_fn = nn.CrossEntropyLoss()
        self.sigmoid_fn = nn.Sigmoid()

        # Constants signifying how much to pay for each respective part of the loss
        self.lam_class = 1
        self.lam_noobj = 10
        self.lam_obj = 1
        self.lam_box = 10

    def forward(self, preds, destination, anchors):
        # Check where obj and noobj (we ignore if target == -1)
        obj = destination[..., 0] == 1  # in paper this is Iobj_i
        noobj = destination[..., 0] == 0  # in paper this is Inoobj_i

        # obj losses
        anchors = anchors.reshape(1, 3, 1, 1, 2)
        frame_preds = torch.cat([self.sigmoid_fn(preds[..., 1:3]), torch.exp(preds[..., 3:5]) * anchors], dim=-1)
        intersec_over_union = IOU(frame_preds[obj], destination[..., 1:5][obj]).detach()
        object_loss = self.mse_loss_fn(self.sigmoid_fn(preds[..., 0:1][obj]),
                                       intersec_over_union * destination[..., 0:1][obj])

        # losses for coordinates
        preds[..., 1:3] = self.sigmoid_fn(preds[..., 1:3])  # x,y coordinates
        destination[..., 3:5] = torch.log(
            (1e-16 + destination[..., 3:5] / anchors)
        )  # width, height coordinates
        box_loss = self.mse_loss_fn(preds[..., 1:5][obj], destination[..., 1:5][obj])

        # noobj losses
        noobj_loss = self.bce_loss_fn(
            (preds[..., 0:1][noobj]), (destination[..., 0:1][noobj]),
        )

        # class losses
        class_loss = self.entropy_loss_fn(
            (preds[..., 5:][obj]), (destination[..., 5][obj].long()),
        )

        return (
                self.lam_box * box_loss
                + self.lam_obj * object_loss
                + self.lam_noobj * noobj_loss
                + self.lam_class * class_loss
        )