"""Segmentation losses: Lovasz-Softmax, Focal, Dice, and a combined loss."""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, List
import numpy as np


def lovasz_grad(gt_sorted):
    """
    Compute gradient of the Lovasz extension w.r.t sorted errors.
    See Algorithm 1 in paper.
    """
    p = len(gt_sorted)
    gts = gt_sorted.sum()
    intersection = gts - gt_sorted.float().cumsum(0)
    union = gts + (1 - gt_sorted).float().cumsum(0)
    jaccard = 1.0 - intersection / union
    if p > 1:
        jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
    return jaccard


def lovasz_softmax_flat(probas, labels, classes='present', ignore_index=255):
    """Multi-class Lovasz-Softmax loss over flat [P,C] probabilities and [P] labels."""
    if probas.numel() == 0:
        return probas * 0.0
    
    C = probas.size(1)
    losses = []
    
    class_to_sum = list(range(C)) if classes == 'all' else []
    
    for c in range(C):
        fg = (labels == c).float()  # foreground for class c
        if classes == 'present' and fg.sum() == 0:
            continue
        if c == ignore_index:
            continue
        
        class_to_sum.append(c) if classes != 'all' else None
        errors = (fg - probas[:, c]).abs()
        errors_sorted, perm = torch.sort(errors, 0, descending=True)
        perm = perm.data
        fg_sorted = fg[perm]
        losses.append(torch.dot(errors_sorted, lovasz_grad(fg_sorted)))
    
    if len(losses) == 0:
        return probas.sum() * 0.0
    
    return torch.stack(losses).mean()


def flatten_probas(probas, labels, ignore_index=255):
    """Flatten [B,C,H,W] preds and [B,H,W] labels to [P,C]/[P], dropping ignored pixels."""
    B, C, H, W = probas.size()
    probas = probas.permute(0, 2, 3, 1).contiguous().view(-1, C)  # [B*H*W, C]
    labels = labels.view(-1)  # [B*H*W]
    
    if ignore_index is not None:
        valid = (labels != ignore_index)
        probas = probas[valid]
        labels = labels[valid]
    
    return probas, labels


class LovaszSoftmaxLoss(nn.Module):
    """
    Lovasz-Softmax loss for multi-class semantic segmentation.
    Directly optimizes the mean IoU (Jaccard index).
    """
    
    def __init__(self, classes='present', ignore_index=255):
        super().__init__()
        self.classes = classes
        self.ignore_index = ignore_index
    
    def forward(self, logits, labels):
        probas = F.softmax(logits, dim=1)
        probas, labels = flatten_probas(probas, labels, self.ignore_index)
        return lovasz_softmax_flat(probas, labels, self.classes, self.ignore_index)


class FocalLoss(nn.Module):
    """
    Focal Loss for multi-class classification.
    FL(p_t) = -alpha_t * (1 - p_t)^gamma * log(p_t)
    """
    
    def __init__(
        self,
        gamma: float = 2.0,
        alpha: Optional[torch.Tensor] = None,
        ignore_index: int = 255,
        reduction: str = 'mean'
    ):
        super().__init__()
        self.gamma = gamma
        self.alpha = alpha
        self.ignore_index = ignore_index
        self.reduction = reduction
    
    def forward(self, logits, labels):
        B, C, H, W = logits.shape
        
        ce_loss = F.cross_entropy(
            logits, labels,
            weight=self.alpha,
            ignore_index=self.ignore_index,
            reduction='none'
        )  # [B, H, W]
        
        logits_flat = logits.permute(0, 2, 3, 1).contiguous().view(-1, C)  # [B*H*W, C]
        labels_flat = labels.view(-1)  # [B*H*W]
        
        valid_mask = (labels_flat != self.ignore_index)
        
        probs = F.softmax(logits_flat, dim=1)  # [B*H*W, C]
        labels_clamped = labels_flat.clamp(0, C-1)  # clamp for safe indexing of ignored pixels
        p_t = probs.gather(1, labels_clamped.unsqueeze(1)).squeeze(1)  # [B*H*W]
        
        focal_weight = (1 - p_t) ** self.gamma  # [B*H*W]
        focal_weight = focal_weight.view(B, H, W)  # [B, H, W]
        
        focal_loss = focal_weight * ce_loss
        
        valid_mask_2d = (labels != self.ignore_index)
        
        if self.reduction == 'mean':
            return focal_loss[valid_mask_2d].mean() if valid_mask_2d.sum() > 0 else focal_loss.sum() * 0.0
        elif self.reduction == 'sum':
            return focal_loss[valid_mask_2d].sum()
        else:
            return focal_loss


class DiceLoss(nn.Module):
    """
    Multi-class Dice loss with optional class weighting.
    """
    
    def __init__(
        self,
        num_classes: int = 21,
        ignore_index: int = 255,
        smooth: float = 1e-6,
        class_weights: Optional[torch.Tensor] = None,
        reduction: str = 'mean'
    ):
        super().__init__()
        self.num_classes = num_classes
        self.ignore_index = ignore_index
        self.smooth = smooth
        self.class_weights = class_weights
        self.reduction = reduction
    
    def forward(self, logits, labels):
        probs = F.softmax(logits, dim=1)  # [B, C, H, W]
        valid_mask = (labels != self.ignore_index).float()  # [B, H, W]
        
        dice_loss = 0.0
        valid_classes = 0
        class_losses = []
        
        for c in range(self.num_classes):
            target_c = (labels == c).float()  # [B, H, W]
            pred_c = probs[:, c, :, :]  # [B, H, W]
            
            target_c = target_c * valid_mask
            pred_c = pred_c * valid_mask
            
            intersection = (pred_c * target_c).sum()
            union = pred_c.sum() + target_c.sum()
            
            if union > 0:
                dice = (2.0 * intersection + self.smooth) / (union + self.smooth)
                class_loss = 1.0 - dice
                
                if self.class_weights is not None:
                    class_loss = class_loss * self.class_weights[c]
                
                class_losses.append(class_loss)
                valid_classes += 1
        
        if valid_classes > 0:
            if self.reduction == 'mean':
                return torch.stack(class_losses).mean()
            else:
                return torch.stack(class_losses).sum()
        
        return logits.sum() * 0.0


class CombinedSegmentationLoss(nn.Module):
    """Weighted combination of CE/Focal, Dice, and Lovasz-Softmax losses."""
    
    def __init__(
        self,
        num_classes: int = 21,
        ignore_index: int = 255,
        weight_ce: float = 1.0,
        weight_dice: float = 1.0,
        weight_lovasz: float = 1.0,
        use_focal: bool = True,
        focal_gamma: float = 2.0,
        class_weights: Optional[torch.Tensor] = None,
    ):
        super().__init__()
        
        self.weight_ce = weight_ce
        self.weight_dice = weight_dice
        self.weight_lovasz = weight_lovasz
        self.use_focal = use_focal
        
        if use_focal:
            self.ce_loss = FocalLoss(
                gamma=focal_gamma,
                alpha=class_weights,
                ignore_index=ignore_index,
            )
        else:
            self.ce_loss = nn.CrossEntropyLoss(
                weight=class_weights,
                ignore_index=ignore_index,
            )
        
        self.dice_loss = DiceLoss(
            num_classes=num_classes,
            ignore_index=ignore_index,
            class_weights=class_weights,
        )
        
        self.lovasz_loss = LovaszSoftmaxLoss(
            classes='present',
            ignore_index=ignore_index,
        )
    
    def forward(self, logits, labels):
        loss_ce = self.ce_loss(logits, labels)
        loss_dice = self.dice_loss(logits, labels)
        loss_lovasz = self.lovasz_loss(logits, labels)
        
        loss_total = (
            self.weight_ce * loss_ce +
            self.weight_dice * loss_dice +
            self.weight_lovasz * loss_lovasz
        )
        
        return {
            'loss_ce': loss_ce,
            'loss_dice': loss_dice,
            'loss_lovasz': loss_lovasz,
            'loss_total': loss_total,
        }