# Copyright 2019-present NAVER Corp.
# CC BY-NC-SA 3.0
# Available only for non-commercial use

import pdb

import torch
import torch.nn as nn
import torch.nn.functional as F

from nets.sampler import *
from nets.repeatability_loss import *
from nets.reliability_loss import *


class MultiLoss(nn.Module):
    """Combines several loss functions for convenience.
    *args: [loss weight (float), loss creator, ... ]

    Example:
        loss = MultiLoss( 1, MyFirstLoss(), 0.5, MySecondLoss() )
    """

    def __init__(self, *args, dbg=()):
        nn.Module.__init__(self)
        assert len(args) % 2 == 0, "args must be a list of (float, loss)"
        self.weights = []
        self.losses = nn.ModuleList()
        for i in range(len(args) // 2):
            weight = float(args[2 * i + 0])
            loss = args[2 * i + 1]
            assert isinstance(loss, nn.Module), "%s is not a loss!" % loss
            self.weights.append(weight)
            self.losses.append(loss)

    def forward(self, select=None, **variables):
        assert not select or all(1 <= n <= len(self.losses) for n in select)
        d = dict()
        cum_loss = 0
        for num, (weight, loss_func) in enumerate(zip(self.weights, self.losses), 1):
            if select is not None and num not in select:
                continue
            l = loss_func(**{k: v for k, v in variables.items()})
            if isinstance(l, tuple):
                assert len(l) == 2 and isinstance(l[1], dict)
            else:
                l = l, {loss_func.name: l}
            cum_loss = cum_loss + weight * l[0]
            for key, val in l[1].items():
                d["loss_" + key] = float(val)
        d["loss"] = float(cum_loss)
        return cum_loss, d