#!/usr/bin/python
# -*- encoding: utf-8 -*-


import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class OhemCELoss(nn.Module):

    def __init__(self, thresh, ignore_lb=255):
        super(OhemCELoss, self).__init__()
        self.thresh = -torch.log(torch.tensor(thresh, requires_grad=False, dtype=torch.float)).cuda()
        self.ignore_lb = ignore_lb
        self.criteria = nn.CrossEntropyLoss(ignore_index=ignore_lb, reduction='none')

    def forward(self, logits, labels):
        n_min = labels[labels != self.ignore_lb].numel() // 16

        #print(logits.size())
        #print(np.unique(logits))

        loss = self.criteria(logits, labels).view(-1)
        
        loss_hard = loss[loss > self.thresh]
        if loss_hard.numel() < n_min:
            loss_hard, _ = loss.topk(n_min)
        return torch.mean(loss_hard)

class BiSeNetV2Loss(nn.Module):

    def __init__(self):
        super(BiSeNetV2Loss, self).__init__()
        self.criteria_pre = OhemCELoss(0.7)
        
        # Booster branch losses TODO dont hardcode num stages
        self.criteria_aux = [OhemCELoss(0.7) for _ in range(4)]


    def forward(self, logits, labels):

        loss_pre = self.criteria_pre(logits[0], labels)

        #loss_aux = [crit(lgt, labels) for crit, lgt in zip(self.criteria_aux, logits[1:])]

        loss = loss_pre #+ sum(loss_aux)

        return loss

