from collections import OrderedDict

import torch

import torch.nn.functional as F

def get_cross_entropy_loss(model_output, batch, model, args):

    logging_dict, predictions = OrderedDict(), OrderedDict()

    logit = model_output["logit"]

    loss = F.cross_entropy(logit, batch["y"].long())

    logging_dict["cross_entropy_loss"] = loss.detach()

    predictions["probs"] = F.softmax(logit, dim=-1).detach()

    predictions["golds"] = batch["y"]

    return loss, logging_dict, predictions

def get_survival_loss(model_output, batch, model, args):

    logging_dict, predictions = OrderedDict(), OrderedDict()

    logit = model_output["logit"]

    y_seq, y_mask = batch["y_seq"], batch["y_mask"]

    loss = F.binary_cross_entropy_with_logits(logit, y_seq.float(), weight=y_mask.float(), reduction='sum') / torch.sum(y_mask.float())

    logging_dict["survival_loss"] = loss.detach()

    predictions["probs"] = torch.sigmoid(logit).detach()

    predictions["golds"] = batch["y"]

    predictions["censors"] = batch["time_at_event"]

    return loss, logging_dict, predictions

def get_annotation_loss(model_output, batch, model, args):

    total_loss, logging_dict, predictions = 0, OrderedDict(), OrderedDict()

    B, _, N, H, W, = model_output["activ"].shape

    batch_mask = batch["has_annotation"]

    for attn_num in [1, 2]:

        side_attn = -1

        if model_output.get("image_attention_{}".format(attn_num), None) is not None:

            if len(batch["image_annotations"].shape) == 4:

                batch["image_annotations"] = batch["image_annotations"].unsqueeze(1)

            # resize annotation to 'activ' size

            annotation_gold = F.interpolate(

                batch["image_annotations"], (N, H, W), mode="area"

            )

            annotation_gold = annotation_gold * batch_mask[:, None, None, None, None]

            # renormalize scores

            mask_area = annotation_gold.sum(dim=(-1, -2)).unsqueeze(-1).unsqueeze(-1)

            mask_area[mask_area == 0] = 1

            annotation_gold /= mask_area

            # reshape annotation into 1D vector

            annotation_gold = annotation_gold.view(B, N, -1).float()

            # get mask over annotation boxes in order to weigh

            # non-annotated scores with zero when computing loss

            annotation_gold_mask = (annotation_gold > 0).float()

            num_annotated_samples = (annotation_gold.view(B * N, -1).sum(-1) > 0).sum()

            num_annotated_samples = max(1, num_annotated_samples)

            pred_attn = (

                model_output["image_attention_{}".format(attn_num)]

                * batch_mask[:, None, None]

            )

            kldiv = (

                F.kl_div(pred_attn, annotation_gold, reduction="none")

                * annotation_gold_mask

            )

            # sum loss per volume and average over batches

            loss = kldiv.sum() / num_annotated_samples

            logging_dict["image_attention_loss_{}".format(attn_num)] = loss.detach()

            total_loss += args.image_attention_loss_lambda * loss

            # attend to cancer side

            cancer_side_mask = (batch["cancer_laterality"][:, :2].sum(-1) == 1).float()[

                :, None

            ]  # only one side is positive

            cancer_side_gold = (

                batch["cancer_laterality"][:, 1].unsqueeze(1).repeat(1, N)

            )  # left side (seen as lung on right) is positive class

            num_annotated_samples = max(N * cancer_side_mask.sum(), 1)

            side_attn = torch.exp(model_output["image_attention_{}".format(attn_num)])

            side_attn = side_attn.view(B, N, H, W)

            side_attn = torch.stack(

                [

                    side_attn[:, :, :, : W // 2].sum((2, 3)),

                    side_attn[:, :, :, W // 2 :].sum((2, 3)),

                ],

                dim=-1,

            )

            side_attn_log = F.log_softmax(side_attn, dim=-1).transpose(1, 2)

            loss = (

                F.cross_entropy(side_attn_log, cancer_side_gold, reduction="none")

                * cancer_side_mask

            ).sum() / num_annotated_samples

            logging_dict[

                "image_side_attention_loss_{}".format(attn_num)

            ] = loss.detach()

            total_loss += args.image_attention_loss_lambda * loss

        if model_output.get("volume_attention_{}".format(attn_num), None) is not None:

            # find size of annotation box per slice and normalize

            annotation_gold = batch["annotation_areas"].float() * batch_mask[:, None]

            if N != args.num_images:

                annotation_gold = F.interpolate(annotation_gold.unsqueeze(1), (N), mode= 'linear', align_corners = True)[:,0]

            area_per_slice = annotation_gold.sum(-1).unsqueeze(-1)

            area_per_slice[area_per_slice == 0] = 1

            annotation_gold /= area_per_slice

            num_annotated_samples = (annotation_gold.sum(-1) > 0).sum()

            num_annotated_samples = max(1, num_annotated_samples)

            # find slices with annotation

            annotation_gold_mask = (annotation_gold > 0).float()

            pred_attn = (

                model_output["volume_attention_{}".format(attn_num)]

                * batch_mask[:, None]

            )

            kldiv = (

                F.kl_div(pred_attn, annotation_gold, reduction="none")

                * annotation_gold_mask

            )  # B, N

            loss = kldiv.sum() / num_annotated_samples

            logging_dict["volume_attention_loss_{}".format(attn_num)] = loss.detach()

            total_loss += args.volume_attention_loss_lambda * loss

            if isinstance(side_attn, torch.Tensor):

                # attend to cancer side

                cancer_side_mask = (

                    batch["cancer_laterality"][:, :2].sum(-1) == 1

                ).float()  # only one side is positive

                cancer_side_gold = batch["cancer_laterality"][

                    :, 1

                ]  # left side (seen as lung on right) is positive class

                num_annotated_samples = max(cancer_side_mask.sum(), 1)

                pred_attn = torch.exp(

                    model_output["volume_attention_{}".format(attn_num)]

                )

                side_attn = (side_attn * pred_attn.unsqueeze(-1)).sum(1)

                side_attn_log = F.log_softmax(side_attn, dim=-1)

                loss = (

                    F.cross_entropy(side_attn_log, cancer_side_gold, reduction="none")

                    * cancer_side_mask

                ).sum() / num_annotated_samples

                logging_dict[

                    "volume_side_attention_loss_{}".format(attn_num)

                ] = loss.detach()

                total_loss += args.volume_attention_loss_lambda * loss

    return total_loss * args.annotation_loss_lambda, logging_dict, predictions

def get_risk_factor_loss(model_output, batch, model, args):

    total_loss, logging_dict, predictions = 0, OrderedDict(), OrderedDict()

    for idx, key in enumerate(args.risk_factor_keys):

        logit = model_output["{}_logit".format(key)]

        gold_rf = batch["risk_factors"][idx]

        is_rf_known = (torch.sum(gold_rf, dim=-1) > 0).unsqueeze(-1).float()

        gold = torch.argmax(gold_rf, dim=-1).contiguous().view(-1)

        loss = (

            F.cross_entropy(logit, gold, reduction="none") * is_rf_known

        ).sum() / max(1, is_rf_known.sum())

        total_loss += loss

        logging_dict["{}_loss".format(key)] = loss.detach()

        probs = F.softmax(logit, dim=-1).detach()

        predictions["{}_probs".format(key)] = probs.detach()

        predictions["{}_golds".format(key)] = gold.detach()

        predictions["{}_risk_factor".format(key)] = batch["risk_factors"][idx]

        # preds = torch.argmax(probs, dim=-1).view(-1)

    return total_loss * args.primary_loss_lambda, logging_dict, predictions

def discriminator_loss(model_output, batch, model, args):

    logging_dict, predictions = OrderedDict(), OrderedDict()

    d_output = model.discriminator(model_output, batch)

    loss = F.cross_entropy(d_output['logit'], batch['origin_dataset'].long()) * args.adv_loss_lambda

    logging_dict['discrim_loss'] = loss.detach()

    predictions['discrim_probs'] = d_output['logit'].detach()

    predictions['discrim_golds'] = batch['origin_dataset']

    if model.reverse_discrim_loss:

        loss = -loss

    return loss, logging_dict, predictions