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


def __init_weight(feature, conv_init, norm_layer, bn_eps, bn_momentum,
                  **kwargs):
    for name, m in feature.named_modules():
        if isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
            conv_init(m.weight, **kwargs)
        elif isinstance(m, norm_layer):
            m.eps = bn_eps
            m.momentum = bn_momentum
            nn.init.constant_(m.weight, 1)
            nn.init.constant_(m.bias, 0)


def init_weight(module_list, conv_init, norm_layer, bn_eps, bn_momentum,
                **kwargs):
    if isinstance(module_list, list):
        for feature in module_list:
            __init_weight(feature, conv_init, norm_layer, bn_eps, bn_momentum,
                          **kwargs)
    else:
        __init_weight(module_list, conv_init, norm_layer, bn_eps, bn_momentum,
                      **kwargs)
        

def unsupervised_loss(probability_map1, probability_map2):
    max1 = torch.max(probability_map1, dim=1)[1].long()
    max2 = torch.max(probability_map2, dim=1)[1].long()

    loss1 = F.binary_cross_entropy(probability_map1, max2)
    loss2 = F.binary_cross_entropy(probability_map2, max1)

    return loss1 + loss2