"""
DINOv3 + DeepLabV3+ 网络架构
用于浒苔分割任务
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, List
import numpy as np
from pathlib import Path

# 导入DINOv3模型
import sys
sys.path.append(str(Path(__file__).resolve().parent / 'dinov3-main'))
from dinov3.hub.backbones import dinov3_vitl16, Weights

class ASPPModule(nn.Module):
    """Atrous Spatial Pyramid Pooling模块"""
    
    def __init__(self, in_channels: int, out_channels: int = 256, rates: List[int] = [6, 12, 18]):
        super().__init__()
        
        # 1x1卷积
        self.conv1x1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
        
        # 3x3卷积 with different dilation rates
        self.atrous_convs = nn.ModuleList()
        for rate in rates:
            self.atrous_convs.append(nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 3, padding=rate, dilation=rate, bias=False),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True)
            ))
        
        # Global average pooling
        self.global_avg_pool = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
        
        # 输出卷积
        self.output_conv = nn.Sequential(
            nn.Conv2d(out_channels * (len(rates) + 2), out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5)
        )
    
    def forward(self, x):
        size = x.shape[2:]
        
        # 1x1卷积
        conv1x1 = self.conv1x1(x)
        
        # Atrous卷积
        atrous_outputs = []
        for atrous_conv in self.atrous_convs:
            atrous_outputs.append(atrous_conv(x))
        
        # Global average pooling
        global_feat = self.global_avg_pool(x)
        global_feat = F.interpolate(global_feat, size=size, mode='bilinear', align_corners=False)
        
        # 拼接所有特征
        concat_feat = torch.cat([conv1x1] + atrous_outputs + [global_feat], dim=1)
        
        # 输出卷积
        output = self.output_conv(concat_feat)
        
        return output

class DeepLabV3PlusDecoder(nn.Module):
    """DeepLabV3+解码器"""
    
    def __init__(self, low_level_channels: int, high_level_channels: int, num_classes: int):
        super().__init__()
        
        # 低层特征处理
        self.low_level_conv = nn.Sequential(
            nn.Conv2d(low_level_channels, 48, 1, bias=False),
            nn.BatchNorm2d(48),
            nn.ReLU(inplace=True)
        )
        
        # ASPP模块
        self.aspp = ASPPModule(high_level_channels, out_channels=256)
        
        # 解码器卷积
        self.decoder_conv = nn.Sequential(
            nn.Conv2d(256 + 48, 256, 3, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            
            nn.Conv2d(256, 256, 3, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.1)
        )
        
        # 分类头
        self.classifier = nn.Conv2d(256, num_classes, 1)
    
    def forward(self, low_level_feat, high_level_feat):
        # 处理低层特征
        low_level_feat = self.low_level_conv(low_level_feat)
        low_level_size = low_level_feat.shape[2:]
        
        # 处理高层特征
        high_level_feat = self.aspp(high_level_feat)
        
        # 上采样高层特征
        high_level_feat = F.interpolate(
            high_level_feat, 
            size=low_level_size, 
            mode='bilinear', 
            align_corners=False
        )
        
        # 拼接特征
        concat_feat = torch.cat([low_level_feat, high_level_feat], dim=1)
        
        # 解码器卷积
        decoder_feat = self.decoder_conv(concat_feat)
        
        # 分类
        output = self.classifier(decoder_feat)
        
        return output

class DinoV3DeepLabV3Plus(nn.Module):
    """DINOv3 + DeepLabV3+ 用于浒苔分割"""
    
    def __init__(self, num_classes=2, backbone_name='dinov3_vitl16', 
                 pretrained=True, weights='SAT493M', use_4channel=True, freeze_backbone=False):
        """
        Args:
            num_classes: 类别数（浒苔/背景）
            backbone_name: DINOv3 backbone名称
            pretrained: 是否使用预训练权重
            weights: 预训练权重类型
            use_4channel: 是否使用4通道输入
            freeze_backbone: 是否冻结backbone参数
        """
        super().__init__()
        
        self.num_classes = num_classes
        self.use_4channel = use_4channel
        self.backbone_name = backbone_name
        self.freeze_backbone = freeze_backbone
        
        # 加载DINOv3 backbone
        if backbone_name == 'dinov3_vitl16':
            try:
                # 处理权重参数
                if isinstance(weights, str):
                    if weights.lower() == 'sat493m':
                        weights_enum = Weights.SAT493M
                    elif weights.lower() == 'lvd1689m':
                        weights_enum = Weights.LVD1689M
                    else:
                        # 如果是文件路径
                        weights_enum = weights
                else:
                    weights_enum = weights
                
                # 加载预训练模型
                self.backbone = dinov3_vitl16(pretrained=pretrained, weights=weights_enum)
            except Exception as e:
                print(f"Warning: Failed to load DINOv3 backbone with error: {e}")
                print("Creating backbone without pretrained weights...")
                # 创建不带预训练权重的模型
                self.backbone = dinov3_vitl16(pretrained=False)
                
            # 注意：DINOv3默认使用3通道，4通道需要修改backbone的第一层
            if use_4channel:
                # 修改第一层以支持4通道输入
                if hasattr(self.backbone, 'patch_embed'):
                    # 获取原始权重
                    original_weight = self.backbone.patch_embed.proj.weight
                    # 创建新的4通道权重
                    new_weight = torch.cat([original_weight, original_weight[:, :1, :, :]], dim=1)
                    # 修改投影层
                    self.backbone.patch_embed.proj = nn.Conv2d(
                        4, original_weight.shape[0], 
                        kernel_size=self.backbone.patch_embed.proj.kernel_size,
                        stride=self.backbone.patch_embed.proj.stride,
                        padding=self.backbone.patch_embed.proj.padding
                    )
                    # 加载新权重
                    with torch.no_grad():
                        self.backbone.patch_embed.proj.weight = nn.Parameter(new_weight)
        else:
            raise ValueError(f"Unsupported backbone: {backbone_name}")
        
        # 获取backbone信息
        self.embed_dim = self.backbone.embed_dim
        self.patch_size = self.backbone.patch_size
        # 从patch_embed获取图像大小信息
        if hasattr(self.backbone, 'patch_embed'):
            # 假设patch_embed有img_size属性，或者我们可以推断它
            if hasattr(self.backbone.patch_embed, 'img_size'):
                self.img_size = self.backbone.patch_embed.img_size
            else:
                # 默认使用224作为图像大小
                self.img_size = 224
        else:
            self.img_size = 224
        
        # 构建特征提取层
        # 将ViT特征转换为适合分割的格式
        self.feature_conv = nn.Sequential(
            nn.Conv2d(self.embed_dim, 256, 1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)
        )
        
        # 低层特征（这里使用backbone的早期层特征）
        self.low_level_channels = 256
        self.high_level_channels = 256
        
        # DeepLabV3+解码器
        self.decoder = DeepLabV3PlusDecoder(
            low_level_channels=self.low_level_channels,
            high_level_channels=self.high_level_channels,
            num_classes=num_classes
        )
        
        # 辅助分类头（用于深度监督）
        self.aux_classifier = nn.Sequential(
            nn.Conv2d(256, 256, 3, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Conv2d(256, num_classes, 1)
        )
        
        # 冻结backbone参数（如果指定）
        if self.freeze_backbone:
            self._freeze_backbone()
    
    def forward(self, x):
        batch_size = x.shape[0]
        input_size = x.shape[2:]
        
        # DINOv3 backbone前向传播
        if hasattr(self.backbone, 'forward_features'):
            # 获取patch特征
            features = self.backbone.forward_features(x)
            
            # DINOv3返回字典，我们需要x_norm_patchtokens
            if isinstance(features, dict) and 'x_norm_patchtokens' in features:
                patch_features = features['x_norm_patchtokens']
            else:
                patch_features = features
            
            # 处理ViT输出格式
            if len(patch_features.shape) == 3:  # (B, N, C)格式
                # 重塑为2D特征图
                feat_h = feat_w = int(np.sqrt(patch_features.shape[1]))
                patch_features = patch_features.transpose(1, 2).view(
                    batch_size, self.embed_dim, feat_h, feat_w
                )
        else:
            # 备用方案
            patch_features = self.backbone(x)
        
        # 特征卷积
        high_level_feat = self.feature_conv(patch_features)
        
        # 创建低层特征（这里简化处理，实际可以从backbone的不同层获取）
        low_level_feat = F.interpolate(
            high_level_feat, 
            scale_factor=4, 
            mode='bilinear', 
            align_corners=False
        )
        
        # DeepLabV3+解码
        output = self.decoder(low_level_feat, high_level_feat)
        
        # 上采样到输入尺寸
        output = F.interpolate(
            output, 
            size=input_size, 
            mode='bilinear', 
            align_corners=False
        )
        
        # 辅助输出（用于训练时的深度监督）
        if self.training:
            aux_output = F.interpolate(
                self.aux_classifier(high_level_feat),
                size=input_size,
                mode='bilinear',
                align_corners=False
            )
            return {'out': output, 'aux': aux_output}
        else:
            return output
    
    def get_backbone_params(self):
        """获取backbone参数"""
        return self.backbone.parameters()
    
    def get_decoder_params(self):
        """获取decoder参数"""
        decoder_params = []
        decoder_params.extend(self.feature_conv.parameters())
        decoder_params.extend(self.decoder.parameters())
        decoder_params.extend(self.aux_classifier.parameters())
        return decoder_params
    
    def _freeze_backbone(self):
        """冻结backbone参数"""
        print("冻结DINOv3 backbone参数...")
        for param in self.backbone.parameters():
            param.requires_grad = False
        
        # 如果使用了4通道，需要确保patch_embed的权重是可训练的
        # 因为这是我们修改过的层
        if self.use_4channel and hasattr(self.backbone, 'patch_embed'):
            for param in self.backbone.patch_embed.proj.parameters():
                param.requires_grad = True
            print("保持patch_embed.proj参数可训练（4通道适配层）")
        
        print(f"已冻结{sum(1 for p in self.backbone.parameters() if not p.requires_grad)}个backbone参数")
    
    def unfreeze_backbone(self):
        """解冻backbone参数"""
        print("解冻DINOv3 backbone参数...")
        for param in self.backbone.parameters():
            param.requires_grad = True
        print("所有backbone参数已解冻")
    
    def get_trainable_parameters(self):
        """获取所有可训练参数"""
        return [p for p in self.parameters() if p.requires_grad]

class FocalLoss(nn.Module):
    """Focal Loss for addressing class imbalance"""
    
    def __init__(self, alpha=1, gamma=2, ignore_index=255, reduction='mean', class_weights=None):
        """
        Args:
            alpha: 平衡参数（可以是标量或列表，如果是列表则按类别应用）
            gamma: 聚焦参数
            ignore_index: 忽略的索引
            reduction: 降维方式
            class_weights: 类别权重tensor，shape为(num_classes,)
        """
        super().__init__()
        # 如果alpha是列表或tensor，注册为buffer（自动移动到正确设备）
        if isinstance(alpha, (list, tuple)):
            self.register_buffer('alpha', torch.tensor(alpha, dtype=torch.float32))
        elif isinstance(alpha, torch.Tensor):
            self.register_buffer('alpha', alpha)
        else:
            self.alpha = alpha
        self.gamma = gamma
        self.ignore_index = ignore_index
        self.reduction = reduction
        # 如果class_weights是tensor，注册为buffer
        if isinstance(class_weights, torch.Tensor):
            self.register_buffer('class_weights', class_weights)
        else:
            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)
        
        # 确保class_weights在正确的设备上
        class_weights = self.class_weights
        if class_weights is not None and isinstance(class_weights, torch.Tensor):
            if class_weights.device != inputs.device:
                class_weights = class_weights.to(inputs.device)
        
        # 计算交叉熵（使用类别权重）
        ce_loss = F.cross_entropy(inputs, targets, reduction='none', weight=class_weights)
        pt = torch.exp(-ce_loss)
        
        # 应用alpha权重（如果是tensor则按类别应用）
        if isinstance(self.alpha, torch.Tensor):
            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 = self.alpha * (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 SeaweedSegmentationLoss(nn.Module):
    """浒苔分割专用损失函数"""
    
    def __init__(self, num_classes=2, focal_alpha=1, 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__()
        
        # 创建类别权重tensor（注册为buffer，会自动移动到正确的设备）
        self.register_buffer('class_weights', torch.tensor([background_weight, foreground_weight], dtype=torch.float32))
        
        # 处理focal_alpha：如果是标量，转换为列表
        if isinstance(focal_alpha, (int, float)):
            focal_alpha = [background_weight, foreground_weight]
        
        # 如果focal_alpha是列表，也注册为buffer
        if isinstance(focal_alpha, (list, tuple)):
            self.register_buffer('focal_alpha_tensor', torch.tensor(focal_alpha, dtype=torch.float32))
            focal_alpha_for_loss = self.focal_alpha_tensor
        else:
            focal_alpha_for_loss = focal_alpha
        
        self.focal_loss = FocalLoss(alpha=focal_alpha_for_loss, gamma=focal_gamma, class_weights=self.class_weights)
        self.dice_weight = dice_weight
        self.focal_weight = focal_weight
        self.num_classes = num_classes
    
    def dice_loss(self, inputs, targets, smooth=1e-6):
        """Dice loss - 改进版本，对前景类别给予更高权重"""
        # 处理多输出格式
        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 + smooth) / (union + smooth)  # (N, C)
        dice_loss_per_class = 1 - dice_score  # (N, C)
        
        # 应用类别权重：给前景类别更高的权重
        # 确保class_weights在正确的设备上
        class_weights = self.class_weights
        if class_weights.device != dice_loss_per_class.device:
            class_weights = class_weights.to(dice_loss_per_class.device)
        weighted_dice_loss = dice_loss_per_class * class_weights.unsqueeze(0)
        
        # 计算加权平均
        dice_loss = weighted_dice_loss.mean()
        
        return dice_loss
    
    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
        }

# 测试函数
if __name__ == "__main__":
    # 测试网络架构
    print("测试DINOv3 + DeepLabV3+网络架构...")
    
    # 测试4通道
    model_4ch = DinoV3DeepLabV3Plus(
        num_classes=2,
        backbone_name='dinov3_vitl16',
        pretrained=False,  # 不加载预训练权重进行测试
        weights='SAT493M',
        use_4channel=True
    )
    
    # 测试输入
    x_4ch = torch.randn(2, 4, 512, 512)
    with torch.no_grad():
        output_4ch = model_4ch(x_4ch)
        if isinstance(output_4ch, dict):
            print(f"4通道训练模式 - 主输出形状: {output_4ch['out'].shape}")
            print(f"4通道训练模式 - 辅助输出形状: {output_4ch['aux'].shape}")
        else:
            print(f"4通道推理模式 - 输出形状: {output_4ch.shape}")
    
    # 测试3通道
    model_3ch = DinoV3DeepLabV3Plus(
        num_classes=2,
        backbone_name='dinov3_vitl16',
        pretrained=False,
        weights='SAT493M',
        use_4channel=False
    )
    
    x_3ch = torch.randn(2, 3, 512, 512)
    with torch.no_grad():
        output_3ch = model_3ch(x_3ch)
        if isinstance(output_3ch, dict):
            print(f"3通道训练模式 - 主输出形状: {output_3ch['out'].shape}")
            print(f"3通道训练模式 - 辅助输出形状: {output_3ch['aux'].shape}")
        else:
            print(f"3通道推理模式 - 输出形状: {output_3ch.shape}")
    
    # 测试损失函数
    print("\n测试损失函数...")
    criterion = SeaweedSegmentationLoss(num_classes=2)
    
    # 模拟预测和标签
    pred = torch.randn(2, 2, 512, 512)
    target = torch.randint(0, 2, (2, 512, 512))
    
    losses = criterion(pred, target)
    print(f"总损失: {losses['total_loss'].item():.4f}")
    print(f"Focal损失: {losses['focal_loss'].item():.4f}")
    print(f"Dice损失: {losses['dice_loss'].item():.4f}")
    
    print("\n网络架构测试完成！")