import torch
import torch.distributed.nn
from torch import distributed as dist, nn as nn
from torch.nn import functional as F
import numpy as np
from sklearn.metrics import average_precision_score, roc_auc_score, accuracy_score

try:
    import horovod.torch as hvd
except ImportError:
    hvd = None


def gather_features(
    audio_features,
    text_features,
    audio_features_mlp=None,
    text_features_mlp=None,
    local_loss=False,
    gather_with_grad=False,
    rank=0,
    world_size=1,
    use_horovod=False,
    mlp_loss=False,
):
    if use_horovod:
        assert hvd is not None, "Please install horovod"
        if gather_with_grad:
            all_audio_features = hvd.allgather(audio_features)
            all_text_features = hvd.allgather(text_features)
            if mlp_loss:
                all_audio_features_mlp = hvd.allgather(audio_features_mlp)
                all_text_features_mlp = hvd.allgather(text_features_mlp)
        else:
            with torch.no_grad():
                all_audio_features = hvd.allgather(audio_features)
                all_text_features = hvd.allgather(text_features)
                if mlp_loss:
                    all_audio_features_mlp = hvd.allgather(audio_features_mlp)
                    all_text_features_mlp = hvd.allgather(text_features_mlp)
            if not local_loss:
                # ensure grads for local rank when all_* features don't have a gradient
                gathered_audio_features = list(
                    all_audio_features.chunk(world_size, dim=0)
                )
                gathered_text_features = list(
                    all_text_features.chunk(world_size, dim=0)
                )
                gathered_audio_features[rank] = audio_features
                gathered_text_features[rank] = text_features
                all_audio_features = torch.cat(gathered_audio_features, dim=0)
                all_text_features = torch.cat(gathered_text_features, dim=0)
                if mlp_loss:
                    gathered_audio_features_mlp = list(
                        all_audio_features_mlp.chunk(world_size, dim=0)
                    )
                    gathered_text_features_mlp = list(
                        all_text_features_mlp.chunk(world_size, dim=0)
                    )
                    gathered_audio_features_mlp[rank] = audio_features_mlp
                    gathered_text_features_mlp[rank] = text_features_mlp
                    all_audio_features_mlp = torch.cat(
                        gathered_audio_features_mlp, dim=0
                    )
                    all_text_features_mlp = torch.cat(gathered_text_features_mlp, dim=0)
    else:
        # We gather tensors from all gpus
        if gather_with_grad:
            all_audio_features = torch.cat(
                torch.distributed.nn.all_gather(audio_features), dim=0
            )
            all_text_features = torch.cat(
                torch.distributed.nn.all_gather(text_features), dim=0
            )
            if mlp_loss:
                all_audio_features_mlp = torch.cat(
                    torch.distributed.nn.all_gather(audio_features_mlp), dim=0
                )
                all_text_features_mlp = torch.cat(
                    torch.distributed.nn.all_gather(text_features_mlp), dim=0
                )
        else:
            gathered_audio_features = [
                torch.zeros_like(audio_features) for _ in range(world_size)
            ]
            gathered_text_features = [
                torch.zeros_like(text_features) for _ in range(world_size)
            ]
            dist.all_gather(gathered_audio_features, audio_features)
            dist.all_gather(gathered_text_features, text_features)
            if mlp_loss:
                gathered_audio_features_mlp = [
                    torch.zeros_like(audio_features_mlp) for _ in range(world_size)
                ]
                gathered_text_features_mlp = [
                    torch.zeros_like(text_features_mlp) for _ in range(world_size)
                ]
                dist.all_gather(gathered_audio_features_mlp, audio_features_mlp)
                dist.all_gather(gathered_text_features_mlp, text_features_mlp)
            if not local_loss:
                # ensure grads for local rank when all_* features don't have a gradient
                gathered_audio_features[rank] = audio_features
                gathered_text_features[rank] = text_features
                if mlp_loss:
                    gathered_audio_features_mlp[rank] = audio_features_mlp
                    gathered_text_features_mlp[rank] = text_features_mlp

            all_audio_features = torch.cat(gathered_audio_features, dim=0)
            all_text_features = torch.cat(gathered_text_features, dim=0)
            if mlp_loss:
                all_audio_features_mlp = torch.cat(gathered_audio_features_mlp, dim=0)
                all_text_features_mlp = torch.cat(gathered_text_features_mlp, dim=0)
    if mlp_loss:
        return (
            all_audio_features,
            all_text_features,
            all_audio_features_mlp,
            all_text_features_mlp,
        )
    else:
        return all_audio_features, all_text_features


class ClipLoss(nn.Module):
    def __init__(
        self,
        local_loss=False,
        gather_with_grad=False,
        cache_labels=False,
        rank=0,
        world_size=1,
        use_horovod=False,
        mlp_loss=False,
        weight_loss_kappa=0,
    ):
        super().__init__()
        self.local_loss = local_loss
        self.gather_with_grad = gather_with_grad
        self.cache_labels = cache_labels
        self.rank = rank
        self.world_size = world_size
        self.use_horovod = use_horovod
        self.mlp_loss = mlp_loss
        self.weighted_loss = bool(weight_loss_kappa != 0)
        self.weight_loss_kappa = weight_loss_kappa
        # cache state
        self.prev_num_logits = 0
        self.labels = {}

    def forward(
        self,
        audio_features,
        text_features,
        logit_scale_a,
        logit_scale_t=None,
        audio_features_mlp=None,
        text_features_mlp=None,
    ):
        device = audio_features.device
        if self.mlp_loss:
            if self.world_size > 1:
                (
                    all_audio_features,
                    all_text_features,
                    all_audio_features_mlp,
                    all_text_features_mlp,
                ) = gather_features(
                    audio_features=audio_features,
                    text_features=text_features,
                    audio_features_mlp=audio_features_mlp,
                    text_features_mlp=text_features_mlp,
                    local_loss=self.local_loss,
                    gather_with_grad=self.gather_with_grad,
                    rank=self.rank,
                    world_size=self.world_size,
                    use_horovod=self.use_horovod,
                    mlp_loss=self.mlp_loss,
                )
                if self.local_loss:
                    a_logits_per_audio = (
                        logit_scale_a * audio_features @ all_text_features_mlp.T
                    )
                    a_logits_per_text = (
                        logit_scale_a * text_features_mlp @ all_audio_features.T
                    )
                    t_logits_per_audio = (
                        logit_scale_t * audio_features_mlp @ all_text_features.T
                    )
                    t_logits_per_text = (
                        logit_scale_t * text_features @ all_audio_features_mlp.T
                    )
                else:
                    a_logits_per_audio = (
                        logit_scale_a * all_audio_features @ all_text_features_mlp.T
                    )
                    a_logits_per_text = a_logits_per_audio.T
                    t_logits_per_audio = (
                        logit_scale_t * all_audio_features_mlp @ all_text_features.T
                    )
                    t_logits_per_text = t_logits_per_audio.T
            else:
                a_logits_per_audio = (
                    logit_scale_a * audio_features @ text_features_mlp.T
                )
                a_logits_per_text = logit_scale_a * text_features_mlp @ audio_features.T
                t_logits_per_audio = (
                    logit_scale_t * audio_features_mlp @ text_features.T
                )
                t_logits_per_text = logit_scale_t * text_features @ audio_features_mlp.T

            # calculated ground-truth and cache if enabled
            num_logits = a_logits_per_audio.shape[0]
            if self.prev_num_logits != num_logits or device not in self.labels:
                labels = torch.arange(num_logits, device=device, dtype=torch.long)
                if self.world_size > 1 and self.local_loss:
                    labels = labels + num_logits * self.rank
                if self.cache_labels:
                    self.labels[device] = labels
                    self.prev_num_logits = num_logits
            else:
                labels = self.labels[device]

            if not self.weighted_loss:
                total_loss = (
                    F.cross_entropy(a_logits_per_audio, labels)
                    + F.cross_entropy(a_logits_per_text, labels)
                    + F.cross_entropy(t_logits_per_audio, labels)
                    + F.cross_entropy(t_logits_per_text, labels)
                ) / 4
            else:
                audio_weight = (audio_features @ audio_features.T).detach()
                audio_weight = (
                    torch.exp(
                        torch.sum(audio_weight, axis=1)
                        / (self.weight_loss_kappa * len(audio_weight))
                    )
                ).detach()
                text_weight = (text_features @ text_features.T).detach()
                text_weight = (
                    torch.exp(
                        torch.sum(text_weight, axis=1)
                        / (self.weight_loss_kappa * len(text_features))
                    )
                ).detach()
                total_loss = (
                    F.cross_entropy(a_logits_per_audio, labels, weight=audio_weight)
                    + F.cross_entropy(a_logits_per_text, labels, weight=audio_weight)
                    + F.cross_entropy(t_logits_per_audio, labels, weight=text_weight)
                    + F.cross_entropy(t_logits_per_text, labels, weight=text_weight)
                ) / 4
        else:
            if self.world_size > 1:
                all_audio_features, all_text_features = gather_features(
                    audio_features=audio_features,
                    text_features=text_features,
                    local_loss=self.local_loss,
                    gather_with_grad=self.gather_with_grad,
                    rank=self.rank,
                    world_size=self.world_size,
                    use_horovod=self.use_horovod,
                    mlp_loss=self.mlp_loss,
                )

                if self.local_loss:
                    logits_per_audio = (
                        logit_scale_a * audio_features @ all_text_features.T
                    )
                    logits_per_text = (
                        logit_scale_a * text_features @ all_audio_features.T
                    )
                else:
                    logits_per_audio = (
                        logit_scale_a * all_audio_features @ all_text_features.T
                    )
                    logits_per_text = logits_per_audio.T
            else:
                logits_per_audio = logit_scale_a * audio_features @ text_features.T
                logits_per_text = logit_scale_a * text_features @ audio_features.T

            # calculated ground-truth and cache if enabled
            num_logits = logits_per_audio.shape[0]
            if self.prev_num_logits != num_logits or device not in self.labels:
                labels = torch.arange(num_logits, device=device, dtype=torch.long)
                if self.world_size > 1 and self.local_loss:
                    labels = labels + num_logits * self.rank
                if self.cache_labels:
                    self.labels[device] = labels
                    self.prev_num_logits = num_logits
            else:
                labels = self.labels[device]
            if not self.weighted_loss:
                total_loss = (
                    F.cross_entropy(logits_per_audio, labels)
                    + F.cross_entropy(logits_per_text, labels)
                ) / 2
            else:
                audio_weight = (all_audio_features @ all_audio_features.T).detach()
                audio_weight = (
                    torch.exp(
                        torch.sum(audio_weight, axis=1)
                        / (self.weight_loss_kappa * len(all_audio_features))
                    )
                ).detach()
                text_weight = (all_text_features @ all_text_features.T).detach()
                text_weight = (
                    torch.exp(
                        torch.sum(text_weight, axis=1)
                        / (self.weight_loss_kappa * len(all_text_features))
                    )
                ).detach()
                total_loss = (
                    F.cross_entropy(logits_per_audio, labels, weight=text_weight)
                    + F.cross_entropy(logits_per_text, labels, weight=audio_weight)
                ) / 2
        return total_loss


def lp_gather_features(pred, target, world_size=1, use_horovod=False):
    if use_horovod:
        assert hvd is not None, "Please install horovod"
        with torch.no_grad():
            all_preds = hvd.allgather(pred)
            all_targets = hvd.allgath(target)
    else:
        gathered_preds = [torch.zeros_like(pred) for _ in range(world_size)]
        gathered_targets = [torch.zeros_like(target) for _ in range(world_size)]

        dist.all_gather(gathered_preds, pred)
        dist.all_gather(gathered_targets, target)
        all_preds = torch.cat(gathered_preds, dim=0)
        all_targets = torch.cat(gathered_targets, dim=0)

    return all_preds, all_targets


def get_map(pred, target):
    pred = torch.sigmoid(pred).numpy()
    target = target.numpy()
    return np.mean(average_precision_score(target, pred, average=None))


def get_acc(pred, target):
    pred = torch.argmax(pred, 1).numpy()
    target = torch.argmax(target, 1).numpy()
    return accuracy_score(target, pred)


def get_mauc(pred, target):
    pred = torch.sigmoid(pred).numpy()
    target = target.numpy()
    return np.mean(roc_auc_score(target, pred, average=None))


class LPMetrics(object):
    def __init__(self, metric_names=["map", "acc", "mauc"]):
        self.metrics = []
        for name in metric_names:
            self.metrics.append(self.get_metric(name))
        self.metric_names = metric_names

    def get_metric(self, name):
        if name == "map":
            return get_map
        elif name == "acc":
            return get_acc
        elif name == "mauc":
            return get_mauc
        else:
            raise ValueError(f"the metric should be at least one of [map, acc, mauc]")

    def evaluate_mertics(self, pred, target):
        metric_dict = {}
        for i in range(len(self.metric_names)):
            metric_dict[self.metric_names[i]] = self.metrics[i](pred, target)
        return metric_dict


def calc_celoss(pred, target):
    target = torch.argmax(target, 1).long()
    return nn.CrossEntropyLoss()(pred, target)


class LPLoss(nn.Module):
    def __init__(self, loss_name):
        super().__init__()
        if loss_name == "bce":
            self.loss_func = nn.BCEWithLogitsLoss()
        elif loss_name == "ce":
            self.loss_func = calc_celoss
        elif loss_name == "mse":
            self.loss_func = nn.MSELoss()
        else:
            raise ValueError(f"the loss func should be at least one of [bce, ce, mse]")

    def forward(self, pred, target):
        loss = self.loss_func(pred, target)
        return loss