"""
浒苔预测脚本 - 256x256图像尺寸
使用训练好的模型进行浒苔分割预测 - 简化版本
"""
import os
import torch
import numpy as np
from pathlib import Path
import json
import time
from datetime import datetime
import cv2
from tqdm import tqdm

# 导入自定义模块
from dinov3_deeplabv3plus import DinoV3DeepLabV3Plus

def load_model(model_path, config, device):
    """加载训练好的模型"""
    print(f"正在加载模型: {model_path}")
    
    # 创建模型
    model = DinoV3DeepLabV3Plus(
        num_classes=config['num_classes'],
        backbone_name=config['backbone_name'],
        pretrained=False,  # 不加载预训练权重，使用我们自己的权重
        weights=config['backbone_weights'],
        use_4channel=config['use_4channel']
    ).to(device)
    
    # 加载权重
    checkpoint = torch.load(model_path, map_location=device)
    
    # 处理不同的checkpoint格式
    if 'model_state_dict' in checkpoint:
        model.load_state_dict(checkpoint['model_state_dict'])
        print(f"加载模型权重 (epoch {checkpoint.get('epoch', 'unknown')})")
    elif 'state_dict' in checkpoint:
        model.load_state_dict(checkpoint['state_dict'])
        print("加载模型权重")
    else:
        # 直接加载state dict
        model.load_state_dict(checkpoint)
        print("加载模型权重")
    
    model.eval()
    print("模型加载完成，进入评估模式")
    return model

def process_image_for_model(image_path, target_size=256, use_4channel=True):
    """处理图像用于模型预测 - 使用与训练相同的方式"""
    # 使用OpenCV读取16位TIF图像
    img_16bit = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
    
    if img_16bit is None:
        raise ValueError(f"无法读取图像文件: {image_path}")
    
    # 使用rasterio方式读取多波段数据
    try:
        import rasterio
        with rasterio.open(image_path) as src:
            # 读取所有波段
            image = src.read()
            # 转换为HWC格式
            image = np.transpose(image, (1, 2, 0))
    except Exception as e:
        print(f"rasterio读取失败，使用OpenCV: {e}")
        # 回退到OpenCV
        image = img_16bit
    
    # 处理通道数
    if use_4channel and image.shape[2] >= 4:
        # 使用4通道
        processed_image = image[:, :, :4]
    else:
        # 使用3通道假彩色（432波段）
        if image.shape[2] >= 4:
            # 选择第4,3,2波段（索引3,2,1）
            processed_image = image[:, :, [3, 2, 1]]
        else:
            # 如果通道数不足，使用可用通道
            processed_image = image[:, :, :3]
            if image.shape[2] < 3:
                # 如果还是不足3通道，重复填充
                while processed_image.shape[2] < 3:
                    processed_image = np.concatenate([processed_image, processed_image[:, :, -1:]], axis=2)
    
    # 调整尺寸
    if processed_image.shape[0] != target_size or processed_image.shape[1] != target_size:
        processed_image = cv2.resize(processed_image, (target_size, target_size), interpolation=cv2.INTER_LINEAR)
    
    # 转换为float32并标准化到0-1范围
    processed_image = processed_image.astype(np.float32)
    if processed_image.max() > 1.0:
        processed_image = processed_image / 65535.0  # 16位最大值
    
    # 标准化
    if use_4channel:
        # 4通道标准化
        mean = np.array([0.430, 0.411, 0.296, 0.350])
        std = np.array([0.213, 0.156, 0.143, 0.180])
    else:
        # 3通道标准化
        mean = np.array([0.430, 0.411, 0.296])
        std = np.array([0.213, 0.156, 0.143])
    
    # 应用标准化
    for i in range(processed_image.shape[2]):
        processed_image[:, :, i] = (processed_image[:, :, i] - mean[i]) / std[i]
    
    # 转换为tensor并添加batch维度
    img_tensor = torch.from_numpy(processed_image).permute(2, 0, 1).float()
    
    return img_tensor.unsqueeze(0)  # 添加batch维度

def predict_image(model, image_path, device, config, threshold=0.5):
    """
    预测单张图像
    
    Args:
        model: 训练好的模型
        image_path: 图像路径
        device: 设备
        config: 配置字典
        threshold: 浒苔预测阈值，只有当浒苔概率超过此阈值时才预测为浒苔（默认0.5）
    """
    # 处理图像
    image_tensor = process_image_for_model(
        image_path, 
        target_size=config['image_size'], 
        use_4channel=config['use_4channel']
    ).to(device)
    
    # 预测
    with torch.no_grad():
        start_time = time.time()
        output = model(image_tensor)
        inference_time = time.time() - start_time
        
        # 处理输出格式
        if isinstance(output, dict):
            output = output['out']
        
        # 获取预测结果
        pred = torch.softmax(output, dim=1)
        
        # 计算浒苔概率
        seaweed_prob = pred[0, 1].squeeze(0).cpu().numpy()  # 类别1是浒苔
        
        # 使用阈值进行预测：只有当浒苔概率超过阈值时才预测为浒苔
        # 这样可以避免在无浒苔图片中，由于两个类别概率接近而误判
        pred_class = (seaweed_prob > threshold).astype(np.uint8)
        
        return {
            'prediction': pred_class,
            'seaweed_probability': seaweed_prob,
            'inference_time': inference_time
        }

def save_prediction_results(prediction, seaweed_prob, output_path, base_name):
    """保存预测结果到文件"""
    # 保存预测掩码 (0-1)
    pred_mask = prediction.astype(np.uint8)
    np.save(str(output_path / f"{base_name}_prediction.npy"), pred_mask)
    
    # 保存概率图 (0-1)
    np.save(str(output_path / f"{base_name}_probability.npy"), seaweed_prob)
    
    # 保存为PNG图像
    pred_png = (prediction * 255).astype(np.uint8)
    prob_png = (seaweed_prob * 255).astype(np.uint8)
    
    cv2.imwrite(str(output_path / f"{base_name}_prediction.png"), pred_png)
    cv2.imwrite(str(output_path / f"{base_name}_probability.png"), prob_png)

def calculate_metrics(prediction, target, num_classes=2):
    """计算评估指标 - 精确率、召回率、IoU、F1"""
    # 确保输入是numpy数组
    if isinstance(prediction, torch.Tensor):
        prediction = prediction.cpu().numpy()
    if isinstance(target, torch.Tensor):
        target = target.cpu().numpy()
    
    # 展平数组
    pred_flat = prediction.flatten()
    target_flat = target.flatten()
    
    # 计算混淆矩阵
    confusion_matrix = np.zeros((num_classes, num_classes))
    for i in range(num_classes):
        for j in range(num_classes):
            confusion_matrix[i, j] = np.sum((pred_flat == i) & (target_flat == j))
    
    # 计算每个类别的指标
    metrics_per_class = []
    for i in range(num_classes):
        tp = confusion_matrix[i, i]  # 真正例
        fp = confusion_matrix[i, :].sum() - tp  # 假正例
        fn = confusion_matrix[:, i].sum() - tp  # 假负例
        tn = confusion_matrix.sum() - tp - fp - fn  # 真负例
        
        # 计算指标
        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
        iou = tp / (tp + fp + fn) if (tp + fp + fn) > 0 else 0.0
        f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0.0
        
        metrics_per_class.append({
            'precision': precision,
            'recall': recall,
            'iou': iou,
            'f1': f1,
            'tp': int(tp),
            'fp': int(fp),
            'fn': int(fn),
            'tn': int(tn)
        })
    
    # 只计算前景（浒苔，类别1）的指标
    foreground_metrics = metrics_per_class[1] if len(metrics_per_class) > 1 else metrics_per_class[0]
    
    # 计算整体准确率
    accuracy = np.sum(pred_flat == target_flat) / len(pred_flat)
    
    return {
        'accuracy': float(accuracy),
        'foreground_precision': foreground_metrics['precision'],
        'foreground_recall': foreground_metrics['recall'],
        'foreground_iou': foreground_metrics['iou'],
        'foreground_f1': foreground_metrics['f1'],
        'foreground_tp': foreground_metrics['tp'],
        'foreground_fp': foreground_metrics['fp'],
        'foreground_fn': foreground_metrics['fn'],
        'foreground_tn': foreground_metrics['tn'],
        'metrics_per_class': metrics_per_class
    }

def calculate_statistics(prediction, seaweed_prob):
    """计算统计信息"""
    total_pixels = prediction.size
    seaweed_pixels = np.sum(prediction == 1)
    seaweed_ratio = seaweed_pixels / total_pixels
    
    # 计算平均概率
    avg_probability = np.mean(seaweed_prob)
    max_probability = np.max(seaweed_prob)
    
    return {
        'total_pixels': int(total_pixels),
        'seaweed_pixels': int(seaweed_pixels),
        'seaweed_ratio': float(seaweed_ratio),
        'avg_probability': float(avg_probability),
        'max_probability': float(max_probability)
    }

def load_label_mask(mask_path, target_size=256):
    """加载标签掩码"""
    try:
        # 尝试使用OpenCV读取
        mask = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
        if mask is None:
            # 如果OpenCV失败，尝试使用rasterio
            import rasterio
            with rasterio.open(mask_path) as src:
                mask = src.read(1)  # 读取第一个波段
        else:
            # 如果OpenCV成功但有多通道，只取第一个通道
            if len(mask.shape) > 2:
                mask = mask[:, :, 0]
        
        # 调整尺寸
        if mask.shape[0] != target_size or mask.shape[1] != target_size:
            mask = cv2.resize(mask, (target_size, target_size), interpolation=cv2.INTER_NEAREST)
        
        # 确保是二值掩码 (0, 1)
        mask = (mask > 0).astype(np.uint8)
        
        return mask
        
    except Exception as e:
        print(f"加载标签失败 {mask_path}: {str(e)}")
        return None

def main():
    """主预测函数"""
    print("=" * 60)
    print("浒苔预测系统 - 256x256图像 (带评估指标)")
    print("=" * 60)
    
    # 配置
    model_path = "outputs/seaweed_segmentation_improved_epoch500/best_checkpoint.pth"
    config_path = "outputs/seaweed_segmentation_improved_epoch500/config.json"
    test_image_dir = "data/test/images"
    test_mask_dir = "data/test/masks"  # 添加标签目录
    output_dir = "outputs/predictions_256x256_results"
    
    # 预测阈值：只有当浒苔概率超过此阈值时才预测为浒苔
    # 建议值：0.5-0.7，可以根据验证集调整
    prediction_threshold = 0.5
    
    # 创建设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")
    
    # 加载配置
    with open(config_path, 'r', encoding='utf-8') as f:
        config = json.load(f)
    
    print(f"模型配置:")
    print(f"  - 输入通道: {4 if config['use_4channel'] else 3}")
    print(f"  - 图像尺寸: {config['image_size']}x{config['image_size']}")
    print(f"  - 类别数: {config['num_classes']}")
    print(f"  - 预测阈值: {prediction_threshold}")
    
    # 加载模型
    model = load_model(model_path, config, device)
    
    # 创建输出目录
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # 获取测试图像
    test_images = [f for f in os.listdir(test_image_dir) if f.endswith('.TIF')]
    print(f"\n找到 {len(test_images)} 张测试图像")
    
    # 检查是否有标签数据
    has_labels = os.path.exists(test_mask_dir)
    if has_labels:
        print(f"找到标签目录，将进行评估指标计算")
    else:
        print(f"未找到标签目录，仅进行预测")
    
    # 预测结果统计
    results = []
    evaluation_results = []
    total_inference_time = 0
    
    print("\n开始预测...")
    
    # 使用tqdm显示进度条
    for i, image_file in enumerate(tqdm(test_images, desc="预测进度"), 1):
        image_path = os.path.join(test_image_dir, image_file)
        
        try:
            # 预测（使用阈值）
            result = predict_image(model, image_path, device, config, threshold=prediction_threshold)
            
            # 计算统计信息
            stats = calculate_statistics(result['prediction'], result['seaweed_probability'])
            
            # 如果有标签，计算评估指标
            if has_labels:
                # 构建对应的标签文件名
                mask_file = image_file.replace('.TIF', '.png')  # 假设标签是PNG格式
                mask_path = os.path.join(test_mask_dir, mask_file)
                
                if os.path.exists(mask_path):
                    # 加载真实标签
                    true_mask = load_label_mask(mask_path, config['image_size'])
                    
                    if true_mask is not None:
                        # 计算评估指标
                        metrics = calculate_metrics(result['prediction'], true_mask)
                        
                        # 保存评估结果
                        eval_result = {
                            'image_file': image_file,
                            'accuracy': metrics['accuracy'],
                            'precision': metrics['foreground_precision'],
                            'recall': metrics['foreground_recall'],
                            'iou': metrics['foreground_iou'],
                            'f1': metrics['foreground_f1'],
                            'tp': metrics['foreground_tp'],
                            'fp': metrics['foreground_fp'],
                            'fn': metrics['foreground_fn'],
                            'tn': metrics['foreground_tn']
                        }
                        evaluation_results.append(eval_result)
                        
                        # 在统计信息中添加评估指标
                        stats.update({
                            'accuracy': metrics['accuracy'],
                            'precision': metrics['foreground_precision'],
                            'recall': metrics['foreground_recall'],
                            'iou': metrics['foreground_iou'],
                            'f1': metrics['foreground_f1']
                        })
            
            # 保存预测结果
            base_name = Path(image_file).stem
            save_prediction_results(result['prediction'], result['seaweed_probability'], output_path, base_name)
            
            # 保存结果
            result_info = {
                'image_file': image_file,
                'inference_time': float(result['inference_time']),
                'statistics': stats
            }
            results.append(result_info)
            
            total_inference_time += result['inference_time']
            
        except Exception as e:
            print(f"  - 预测失败: {str(e)}")
            continue
    
    # 保存预测结果
    results_file = output_path / "prediction_results.json"
    with open(results_file, 'w', encoding='utf-8') as f:
        json.dump(results, f, indent=2, ensure_ascii=False, default=str)
    
    # 保存评估结果
    if evaluation_results:
        eval_file = output_path / "evaluation_results.json"
        with open(eval_file, 'w', encoding='utf-8') as f:
            json.dump(evaluation_results, f, indent=2, ensure_ascii=False, default=str)
    
    # 生成总结报告
    if results:
        avg_inference_time = total_inference_time / len(results)
        
        # 计算平均统计
        all_ratios = [r['statistics']['seaweed_ratio'] for r in results]
        avg_seaweed_ratio = np.mean(all_ratios)
        
        print("\n" + "=" * 60)
        print("预测完成总结")
        print("=" * 60)
        print(f"总图像数: {len(results)}")
        print(f"平均推理时间: {avg_inference_time:.3f}s")
        print(f"平均浒苔比例: {avg_seaweed_ratio:.2%}")
        
        # 如果有评估结果，显示评估指标
        if evaluation_results:
            avg_accuracy = np.mean([r['accuracy'] for r in evaluation_results])
            avg_precision = np.mean([r['precision'] for r in evaluation_results])
            avg_recall = np.mean([r['recall'] for r in evaluation_results])
            avg_iou = np.mean([r['iou'] for r in evaluation_results])
            avg_f1 = np.mean([r['f1'] for r in evaluation_results])
            
            print(f"\n评估指标:")
            print(f"  - 准确率 (Accuracy): {avg_accuracy:.4f}")
            print(f"  - 精确率 (Precision): {avg_precision:.4f}")
            print(f"  - 召回率 (Recall): {avg_recall:.4f}")
            print(f"  - IoU: {avg_iou:.4f}")
            print(f"  - F1分数: {avg_f1:.4f}")
            
            # 计算总体混淆矩阵
            total_tp = sum([r['tp'] for r in evaluation_results])
            total_fp = sum([r['fp'] for r in evaluation_results])
            total_fn = sum([r['fn'] for r in evaluation_results])
            total_tn = sum([r['tn'] for r in evaluation_results])
            
            print(f"\n总体混淆矩阵:")
            print(f"  - 真正例 (TP): {total_tp:,}")
            print(f"  - 假正例 (FP): {total_fp:,}")
            print(f"  - 假负例 (FN): {total_fn:,}")
            print(f"  - 真负例 (TN): {total_tn:,}")
        
        print(f"\n结果保存目录: {output_path}")
        print(f"详细结果文件: {results_file}")
        
        # 保存总结报告
        summary = {
            'timestamp': datetime.now().isoformat(),
            'total_images': len(results),
            'avg_inference_time': avg_inference_time,
            'avg_seaweed_ratio': float(avg_seaweed_ratio),
            'results': results
        }
        
        # 添加评估指标到总结报告
        if evaluation_results:
            summary.update({
                'avg_accuracy': float(avg_accuracy),
                'avg_precision': float(avg_precision),
                'avg_recall': float(avg_recall),
                'avg_iou': float(avg_iou),
                'avg_f1': float(avg_f1),
                'total_confusion_matrix': {
                    'tp': total_tp,
                    'fp': total_fp,
                    'fn': total_fn,
                    'tn': total_tn
                },
                'evaluation_results': evaluation_results
            })
        
        summary_file = output_path / "prediction_summary.json"
        with open(summary_file, 'w', encoding='utf-8') as f:
            json.dump(summary, f, indent=2, ensure_ascii=False, default=str)
        
        print(f"总结报告: {summary_file}")
        
        # 显示一些统计信息
        ratios = [r['statistics']['seaweed_ratio'] for r in results]
        print(f"\n浒苔分布统计:")
        print(f"  - 最小浒苔比例: {min(ratios):.2%}")
        print(f"  - 最大浒苔比例: {max(ratios):.2%}")
        print(f"  - 浒苔比例标准差: {np.std(ratios):.2%}")
        
        # 列出保存的文件类型
        print(f"\n保存的文件类型:")
        print(f"  - .npy文件: NumPy数组格式的预测结果")
        print(f"  - .png文件: 可视化的图像格式")
        print(f"  - .json文件: 预测结果统计信息和评估指标")
    
    print("\n🎉 浒苔预测和评估完成！")
    print(f"预测结果保存在: {output_dir}")

if __name__ == "__main__":
    main()