from lossfunction.softmaxproto import SoftmaxProto
import torch.nn as nn
import lossfunction.softmax as softmax
import torch
import torch.nn.functional as F
import numpy

class Unetloss(nn.Module):
    def __init__(self, nOut, nClasses):
        super(Unetloss, self).__init__()

        self.test_normalize = True
        self.softmax = SoftmaxProto(nOut, nClasses)
        self.mseloss = nn.MSELoss()
        print('Initialised Unet Loss')

    def forward(self, emb, spectrogram, x, label=None):

        nlossE, prec1 = self.softmax(emb, label)
        nlossS = self.mseloss(spectrogram, x)
        # print("\nnlossE:", nlossE,"nlossS:", nlossS)
        # nlossE: 13.1695 , nlossS:0.8902

        return nlossE+10*nlossS, prec1


class UnetMaskloss(nn.Module):
    def __init__(self, nOut, nClasses):
        super(UnetMaskloss, self).__init__()

        self.test_normalize = True
        self.softmax = softmax.Softmax(nOut, nClasses)
        self.mseloss = nn.MSELoss(reduction='sum')
        self.criterion = torch.nn.CrossEntropyLoss()
        print('Initialised UnetMask Loss')

    def forward(self, emb, spectrogram, label=None):

        assert emb.size()[1] >= 2
        nlossEd1 = self.mseloss(emb[:, 0, :], emb[:, 1, :])+self.mseloss(emb[:, 0, :], emb[:, 2, :])
        nlossEd2 = self.mseloss(emb[:, 3, :], emb[:, 4, :])+self.mseloss(emb[:, 3, :], emb[:, 5, :])

        emb_anchor = torch.mean(emb[:, 0:3, :], 1)
        emb_positive = torch.mean(emb[:, 3:6, :], 1)
        stepsize = emb_anchor.size()[0]
        output = -1 * (F.pairwise_distance(emb_positive.unsqueeze(-1), emb_anchor.unsqueeze(-1).transpose(0, 2)) ** 2)
        label0 = torch.from_numpy(numpy.asarray(range(0, stepsize))).cuda()
        nlossEP = self.criterion(output, label0)

        nlossEC, prec1 = self.softmax(emb.reshape(-1, emb.size()[-1]), label.repeat_interleave(emb.size()[1]))

        nlossSd1 = self.mseloss(spectrogram[:, 0, :, :], spectrogram[:, 1, :, :]) + self.mseloss(spectrogram[:, 0, :, :], spectrogram[:, 2, :, :])
        nlossSd2 = self.mseloss(spectrogram[:, 3, :, :], spectrogram[:, 4, :, :]) + self.mseloss(
            spectrogram[:, 3, :, :], spectrogram[:, 5, :, :])

        spec_anchor = torch.mean(spectrogram[:, 0:3, :, :], 1)
        spec_positive = torch.mean(spectrogram[:, 3:6, :, :], 1)

        nlossS = self.mseloss(spec_anchor, spec_positive)
        # print("\nnlossEd1:", nlossEd1, "nlossEd2:", nlossEd2,  "nlossEP:", nlossEP,  "nlossEC:", nlossEC)
        # print("nlossSd1:", nlossSd1, "nlossSd2:", nlossSd2, "nlossS:", nlossS)
        # nlossEd1: 3.9563, nlossEd2: 3.5833, nlossEP:0.6218,nlossEC: 8.7362,
        # nlossSd1: 3.4339,  nlossSd2: 30.1156,nlossS: 2.2820,


        loss = 100*(nlossEd1+nlossEd2)+10*nlossEP+nlossEC+nlossSd1+nlossSd2+10*nlossS
        return loss, prec1


if __name__ == "__main__":
    # a = torch.tensor([[[1, 2], [3, 4]], [[1, 2], [3, 4]]])
    # b = torch.tensor([[[2, 3], [4, 5]], [[1, 2], [3, 4]]])
    a = torch.randint(10,(1,2,3))
    b = torch.randint(10,(1,2,3))
    print(a)
    print(b)
    print(a.shape,a.shape)
    # loss_fn = torch.nn.MSELoss(reduce=False, size_average=True)
    # input = torch.autograd.Variable(torch.from_numpy(a))
    # target = torch.autograd.Variable(torch.from_numpy(b))
    # loss = loss_fn(input.float(), target.float())
    # print(loss)
    distance = F.pairwise_distance(a, b)
    print(distance.shape)
    print(distance)