"""
改进的损失函数 - 更好地处理类别不平衡
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class ImprovedFocalLoss(nn.Module):
    """
    改进的Focal Loss，支持类别权重
    """
    
    def __init__(self, alpha=None, gamma=2, ignore_index=255, reduction='mean', class_weights=None):
        """
        Args:
            alpha: 类别权重列表，如果为None则使用平衡权重
            gamma: 聚焦参数
            ignore_index: 忽略的索引
            reduction: 降维方式
            class_weights: 类别权重tensor，shape为(num_classes,)
        """
        super().__init__()
        if alpha is None:
            alpha = [1.0, 1.0]  # 默认平衡权重
        if isinstance(alpha, (list, tuple)):
            alpha = torch.tensor(alpha, dtype=torch.float32)
        self.alpha = alpha
        self.gamma = gamma
        self.ignore_index = ignore_index
        self.reduction = reduction
        self.class_weights = class_weights
        
    def forward(self, inputs, targets):
        """
        Args:
            inputs: 预测值 (N, C, H, W)
            targets: 目标值 (N, H, W)
        """
        # 处理多输出格式
        if isinstance(inputs, dict):
            inputs = inputs['out']
        
        # 忽略指定索引
        if self.ignore_index is not None:
            mask = targets != self.ignore_index
            targets = targets[mask]
            inputs = inputs.permute(0, 2, 3, 1)[mask]
        else:
            inputs = inputs.permute(0, 2, 3, 1).contiguous().view(-1, inputs.size(1))
            targets = targets.view(-1)
        
        # 计算交叉熵
        ce_loss = F.cross_entropy(inputs, targets, reduction='none', weight=self.class_weights)
        pt = torch.exp(-ce_loss)
        
        # 应用alpha权重
        if self.alpha is not None:
            if self.alpha.device != targets.device:
                self.alpha = self.alpha.to(targets.device)
            alpha_t = self.alpha[targets]
            focal_loss = alpha_t * (1 - pt) ** self.gamma * ce_loss
        else:
            focal_loss = (1 - pt) ** self.gamma * ce_loss
        
        if self.reduction == 'mean':
            return focal_loss.mean()
        elif self.reduction == 'sum':
            return focal_loss.sum()
        else:
            return focal_loss


class WeightedDiceLoss(nn.Module):
    """
    加权Dice Loss，给不同类别不同的权重
    """
    
    def __init__(self, num_classes=2, smooth=1e-6, class_weights=None):
        """
        Args:
            num_classes: 类别数
            smooth: 平滑系数
            class_weights: 类别权重tensor，shape为(num_classes,)
        """
        super().__init__()
        self.num_classes = num_classes
        self.smooth = smooth
        self.class_weights = class_weights
        
    def forward(self, inputs, targets):
        """
        Args:
            inputs: 预测值 (N, C, H, W)
            targets: 目标值 (N, H, W)
        """
        # 处理多输出格式
        if isinstance(inputs, dict):
            inputs = inputs['out']
        
        # 将预测转换为概率
        inputs = torch.softmax(inputs, dim=1)
        
        # 创建one-hot编码
        targets_one_hot = F.one_hot(targets, num_classes=self.num_classes)
        targets_one_hot = targets_one_hot.permute(0, 3, 1, 2).float()
        
        # 计算每个类别的dice系数
        intersection = (inputs * targets_one_hot).sum(dim=(2, 3))  # (N, C)
        union = inputs.sum(dim=(2, 3)) + targets_one_hot.sum(dim=(2, 3))  # (N, C)
        
        dice_score = (2. * intersection + self.smooth) / (union + self.smooth)  # (N, C)
        
        # 计算每个类别的dice loss
        dice_loss_per_class = 1 - dice_score  # (N, C)
        
        # 应用类别权重
        if self.class_weights is not None:
            if self.class_weights.device != dice_loss_per_class.device:
                self.class_weights = self.class_weights.to(dice_loss_per_class.device)
            # 对每个类别应用权重
            weighted_dice_loss = dice_loss_per_class * self.class_weights.unsqueeze(0)
            dice_loss = weighted_dice_loss.mean()
        else:
            dice_loss = dice_loss_per_class.mean()
        
        return dice_loss


class ImprovedSeaweedSegmentationLoss(nn.Module):
    """
    改进的浒苔分割损失函数
    更好地处理类别不平衡，特别是背景像素远多于浒苔像素的情况
    """
    
    def __init__(self, num_classes=2, focal_alpha=None, focal_gamma=2, 
                 dice_weight=0.5, focal_weight=1.0, 
                 background_weight=1.0, foreground_weight=2.0):
        """
        Args:
            num_classes: 类别数
            focal_alpha: Focal Loss的alpha参数（类别权重）
            focal_gamma: Focal Loss的gamma参数
            dice_weight: Dice Loss的权重
            focal_weight: Focal Loss的权重
            background_weight: 背景类别的权重（类别0）
            foreground_weight: 前景类别（浒苔）的权重（类别1）
        """
        super().__init__()
        
        # 创建类别权重
        class_weights = torch.tensor([background_weight, foreground_weight], dtype=torch.float32)
        
        # 创建Focal Loss的alpha权重
        if focal_alpha is None:
            # 默认给前景更高的权重
            focal_alpha = [background_weight, foreground_weight]
        
        self.focal_loss = ImprovedFocalLoss(
            alpha=focal_alpha, 
            gamma=focal_gamma,
            class_weights=class_weights
        )
        
        self.dice_loss = WeightedDiceLoss(
            num_classes=num_classes,
            class_weights=class_weights
        )
        
        self.dice_weight = dice_weight
        self.focal_weight = focal_weight
        self.num_classes = num_classes
        
    def forward(self, inputs, targets):
        """计算组合损失"""
        focal_loss = self.focal_loss(inputs, targets)
        dice_loss = self.dice_loss(inputs, targets)
        
        total_loss = self.focal_weight * focal_loss + self.dice_weight * dice_loss
        
        return {
            'total_loss': total_loss,
            'focal_loss': focal_loss,
            'dice_loss': dice_loss
        }


def calculate_class_weights_from_dataset(dataset, num_classes=2):
    """
    从数据集中计算类别权重
    用于平衡类别不平衡问题
    """
    print("正在计算类别权重...")
    total_pixels = 0
    class_counts = torch.zeros(num_classes)
    
    for i in range(len(dataset)):
        sample = dataset[i]
        mask = sample['mask']
        
        # 统计每个类别的像素数
        for c in range(num_classes):
            class_counts[c] += (mask == c).sum().item()
        
        total_pixels += mask.numel()
    
    # 计算权重：总像素数 / (类别数 * 该类别的像素数)
    class_weights = total_pixels / (num_classes * class_counts + 1e-6)
    
    # 归一化权重
    class_weights = class_weights / class_weights.sum() * num_classes
    
    print(f"类别像素统计: {class_counts.tolist()}")
    print(f"类别权重: {class_weights.tolist()}")
    
    return class_weights