from typing import Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.models.llama.modeling_llama import (
    LlamaAttention,
    apply_rotary_pos_emb,
    math,
)


def calculate_dice_loss(
    predictions: torch.Tensor,
    ground_truth: torch.Tensor,
    mask_count: float,
    scale_factor=1000,
    eps=1e-6,
):
    predictions = predictions.sigmoid()
    predictions = predictions.flatten(1, 2)
    ground_truth = ground_truth.flatten(1, 2)

    intersection = 2 * (predictions / scale_factor * ground_truth).sum(dim=-1)
    union = (predictions / scale_factor).sum(dim=-1) + (
        ground_truth / scale_factor
    ).sum(dim=-1)

    dice_loss = 1 - (intersection + eps) / (union + eps)
    dice_loss = dice_loss.sum() / (mask_count + 1e-8)
    return dice_loss


def compute_sigmoid_cross_entropy(
    predictions: torch.Tensor, targets: torch.Tensor, mask_count: float
):
    loss = F.binary_cross_entropy_with_logits(predictions, targets, reduction="none")
    loss = loss.flatten(1, 2).mean(1)
    loss = loss.sum() / (mask_count + 1e-8)
    return loss