"""
Here, we ensure that contrastive loss
"""

import numpy as np
import torch
import torch.nn as nn


def nc2(n):
    return int(n * (n - 1) / 2)


def get_similarity_matrix(noise_levels):
    """
    noise_levels: N*1
    Returns:
        N*N boolean matrix X. X[i,j] is true iff noise_levels[i] == noise_levels[j]. For digonal elements, it has 0
    """
    batch_size = len(noise_levels)
    noise_levels = noise_levels.view(batch_size, -1)
    b = noise_levels[:, :, None].repeat(1, 1, batch_size)
    c = torch.swapaxes(b, 0, 2)
    sim_matrix = b == c
    assert sim_matrix.shape == (batch_size, 1, batch_size)
    sim_matrix = sim_matrix[:, 0, :]
    sim_matrix[np.arange(batch_size), np.arange(batch_size)] = 0
    return sim_matrix


def get_contrastive_loss_vectorized(representation, noise_levels, tau_pos=0.1, tau_neg=0.1):

    # compute pairwise MSE between all representations.
    batch_size = len(representation)
    representation = representation.view(batch_size, -1)
    b = representation[:, :, None].repeat(1, 1, batch_size)
    c = torch.swapaxes(b, 0, 2)
    pairwise_mse = torch.sum((b - c)**2, dim=1) / b.shape[1]

    # clamp mse values
    pairwise_sim = get_similarity_matrix(noise_levels)
    pairwise_mse[pairwise_sim] = torch.clamp(pairwise_mse[pairwise_sim] - tau_pos, min=0)
    pairwise_mse[~pairwise_sim] = torch.clamp(tau_neg - pairwise_mse[~pairwise_sim], min=0)
    pairwise_mse[np.arange(batch_size), np.arange(batch_size)] = 0

    # scale mse values
    c_pos = torch.sum(pairwise_sim) / 2
    c_neg = nc2(batch_size) - c_pos
    pairwise_mse[pairwise_sim] = pairwise_mse[pairwise_sim] * 1 / c_pos
    pairwise_mse[~pairwise_sim] = pairwise_mse[~pairwise_sim] * 1 / c_neg
    return pairwise_mse.sum() / 2


class ContrastiveLearninglossOnLatent(nn.Module):
    def __init__(self, latent_size_dict) -> None:
        super().__init__()
        self._lsizes = latent_size_dict

        self._tau_pos = 0.1
        self._tau_neg = 0.1

    def get_loss(self, latent_activations, noise_levels):
        latent_loss_dict = {}
        seen_loss_sz = []
        for one_layer_activation in latent_activations:
            sz = one_layer_activation.shape[-1]

            assert sz not in seen_loss_sz
            seen_loss_sz.append(sz)

            if self._lsizes[sz] == 0:
                continue
            N = self._lsizes[sz]
            noise_z = one_layer_activation[:, :N]
            latent_loss_dict[sz] = get_contrastive_loss_vectorized(noise_z,
                                                                   noise_levels,
                                                                   tau_pos=0,
                                                                   tau_neg=self._tau_neg)

        return latent_loss_dict

    def forward(self, latent_activations, noise_levels):
        return self.get_loss(latent_activations, noise_levels)


class DisentanglementModule(nn.Module):
    def __init__(self, latent_size_dict, image_model, cl_latent_weight=None) -> None:
        super().__init__()
        self._lsizes = latent_size_dict
        self._image_model = image_model
        for param in self._image_model.parameters():
            param.requires_grad = False

        self._tau_pos = 0.1
        self._tau_neg = 0.1
        assert cl_latent_weight is not None and cl_latent_weight <= 1 and cl_latent_weight >= 0
        self._w = cl_latent_weight

    def get_loss(self, activations, pred_image, noise_levels):
        latent_loss = 0
        cnt = 0
        for key in activations:
            if self._lsizes[key] == 0:
                continue
            N = self._lsizes[key]
            cnt += 1
            noise_z = activations[key][:, :N]
            latent_loss += get_contrastive_loss_vectorized(noise_z, noise_levels, tau_pos=0, tau_neg=self._tau_neg)

        latent_loss = latent_loss / cnt

        if self._w < 1:
            representation = self._image_model(pred_image)
            output_contrastive_loss = get_contrastive_loss_vectorized(representation,
                                                                      noise_levels,
                                                                      tau_pos=self._tau_pos,
                                                                      tau_neg=self._tau_neg)

            return self._w * latent_loss + (1 - self._w) * output_contrastive_loss
        return latent_loss

    def forward(self, activations, pred_image, noise_levels):
        return self.get_loss(activations, pred_image, noise_levels)


if __name__ == '__main__':
    rep = torch.Tensor(np.arange(40).reshape(8, 5))
    noise_level_list = torch.Tensor([10, 20, 10, 40, 30, 40, 20, 30])
    loss = get_contrastive_loss_vectorized(rep, noise_level_list).item()
    assert abs(loss - 262.4) < 1e-4, f'{loss - 262.4}'
