import torch
import argparse
import segmentation_models_pytorch as smp
import torch.nn as nn


def sim(z_i, z_j):
    norm_dot_product = None
    norm_dot_product = torch.dot(z_i, z_j)/\
    (torch.linalg.norm(z_i) * torch.linalg.norm(z_j))
    return norm_dot_product

def sim_positive_pairs(out_left, out_right):
    """Normalized dot product between positive pairs.

    Inputs:
    - out_left: NxD tensor; output of the projection head g(), left branch in SimCLR model.
    - out_right: NxD tensor; output of the projection head g(), right branch in SimCLR model.
    Each row is a z-vector for an augmented sample in the batch.
    The same row in out_left and out_right form a positive pair.
    
    Returns:
    - A Nx1 tensor; each row k is the normalized dot product between out_left[k] and out_right[k].
    """
    pos_pairs = None
    
    pos_pairs = torch.sum(out_left * out_right, dim = 1) / \
    (torch.linalg.norm(out_left, dim = 1) * torch.linalg.norm(out_right, dim = 1))

    pos_pairs = pos_pairs.unsqueeze(1)
    return pos_pairs


def compute_sim_matrix(out):
    """Compute a 2N x 2N matrix of normalized dot products between all pairs of augmented examples in a batch.

    Inputs:
    - out: 2N x D tensor; each row is the z-vector (output of projection head) of a single augmented example.
    There are a total of 2N augmented examples in the batch.
    
    Returns:
    - sim_matrix: 2N x 2N tensor; each element i, j in the matrix is the normalized dot product between out[i] and out[j].
    """
    sim_matrix = None

    sim_matrix = torch.mm(out, out.T) 
    sim_matrix/= torch.linalg.norm(out, dim = 1).unsqueeze(1)
    sim_matrix/= torch.linalg.norm(out, dim = 1).unsqueeze(1).T

    return sim_matrix

def simclr_loss_vectorized(out_left, out_right, tau):

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    N = out_left.shape[0]
    
    # Concatenate out_left and out_right into a 2*N x D tensor.
    out = torch.cat([out_left, out_right], dim=0)  # [2*N, D]
    
    # Compute similarity matrix between all pairs of augmented examples in the batch.
    sim_matrix = compute_sim_matrix(out)  # [2*N, 2*N]
    
    # Step 1: Use sim_matrix to compute the denominator value for all augmented samples.
    # Hint: Compute e^{sim / tau} and store into exponential, which should have shape 2N x 2N.
    exponential = torch.exp(sim_matrix/tau)
    
    # This binary mask zeros out terms where k=i.
    mask = (torch.ones_like(exponential, device=device) - torch.eye(2 * N, device=device)).to(device).bool()
    
   #print((sim_matrix * mask).argmax(dim = 1))

    # We apply the binary mask.
    exponential = exponential.masked_select(mask).view(2 * N, -1)
    
    # Hint: Compute the denominator values for all augmented samples. This should be a 2N x 1 vector.
    denom = torch.sum(exponential, dim = 1).unsqueeze(1)
 
    # Step 2: Compute similarity between positive pairs.

    sim_pos = sim_positive_pairs(out_left, out_right)

    numerator = None

    numerator = torch.exp(sim_pos/tau)
    
    # Step 4: Now that you have the numerator and denominator for all augmented samples, compute the total loss.
    loss = None

    loss = torch.sum(-torch.log(numerator/denom[:N]) - torch.log(numerator/denom[N:])) / (2*N)

    #print(numerator/denom[:N])
    #print(numerator/denom[N:])
    
    return loss

class CombinedLoss(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, pred, truth, aux = None, aux_features = None):

        if pred is not None:
            DiceLoss = smp.losses.DiceLoss(mode = 'multilabel', from_logits = True)
            BCELoss = smp.losses.SoftBCEWithLogitsLoss()
            Loss = 0.4 * DiceLoss(pred, truth) + 0.6 * BCELoss(pred,truth)

        if aux == "simclr":
            N, D = aux_features.shape
            if (N % 2) == 1:
                aux_features = aux_features[:-1, :]
                N, D = aux_features.shape

            aux_features = aux_features.reshape(N//2, 2, D)

            SimClrLoss = simclr_loss_vectorized(aux_features[:, 0, :], aux_features[:, 1, :],  tau = 0.1)

            if pred is None:
                return SimClrLoss
            
            else:
                Loss = 0.9 * Loss + 0.1 * SimClrLoss
                #print('simclr_loss:', SimClrLoss)
                return Loss, SimClrLoss

        elif aux == "reg":
            #Add reg loss code here
            pass
        
        else:
            return Loss

def get_loss_fn(loss_args):
    loss_args_ = loss_args
    if isinstance(loss_args, argparse.Namespace):
        loss_args_ = vars(loss_args)
    loss_fn = loss_args_.get("loss_fn")

    if loss_fn == "BCE":
        return torch.nn.BCEWithLogitsLoss()
    elif loss_fn == "CE":
        return torch.nn.CrossEntropyLoss()
    elif loss_fn == 'DBE':
        return smp.losses.DiceLoss(mode = 'binary', from_logits = True)
    elif loss_fn == 'DLE':
        return smp.losses.DiceLoss(mode = 'multilabel', from_logits = True)
    elif loss_fn == 'Combined':
        return CombinedLoss()
    else:
        raise ValueError(f"loss_fn {loss_args.loss_fn} not supported.")