import cv2
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import os
from znzh_x import PredictionVisualizer

class SimpleChangeVisualizer:
    def __init__(self):
        self.fig_size = (15, 10)
    
    def visualize_change_detection(self, img1_path, img2_path, mask_path, save_path=None):
        """
        简单的变化检测可视化
        
        Args:
            img1_path: 第一期影像路径
            img2_path: 第二期影像路径  
            mask_path: 预测掩膜路径
            save_path: 保存路径（可选）
        """
        # 读取图像
        img1 = cv2.imread(img1_path)
        img2 = cv2.imread(img2_path)
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        
        if img1 is None or img2 is None or mask is None:
            print("错误：无法读取图像文件")
            return
        
        # 转换颜色空间
        img1_rgb = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
        img2_rgb = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
        
        # 创建变化区域的彩色掩膜
        change_mask = np.zeros_like(img2_rgb)
        change_mask[mask > 127] = [255, 0, 0]  # 红色表示变化
        
        # 创建叠加图像
        overlay1 = cv2.addWeighted(img1_rgb, 0.7, change_mask, 0.3, 0)
        overlay2 = cv2.addWeighted(img2_rgb, 0.7, change_mask, 0.3, 0)
        
        # 创建subplot
        fig, axes = plt.subplots(2, 3, figsize=self.fig_size)
        fig.suptitle('变化检测结果可视化', fontsize=16, fontweight='bold')
        
        # 第一行：原始图像
        axes[0, 0].imshow(img1_rgb)
        axes[0, 0].set_title('第一期影像', fontsize=12)
        axes[0, 0].axis('off')
        
        axes[0, 1].imshow(img2_rgb)
        axes[0, 1].set_title('第二期影像', fontsize=12)
        axes[0, 1].axis('off')
        
        axes[0, 2].imshow(mask, cmap='gray')
        axes[0, 2].set_title('变化检测掩膜', fontsize=12)
        axes[0, 2].axis('off')
        
        # 第二行：叠加结果
        axes[1, 0].imshow(overlay1)
        axes[1, 0].set_title('第一期影像+变化区域', fontsize=12)
        axes[1, 0].axis('off')
        
        axes[1, 1].imshow(overlay2)
        axes[1, 1].set_title('第二期影像+变化区域', fontsize=12)
        axes[1, 1].axis('off')
        
        # 变化统计
        total_pixels = mask.size
        changed_pixels = np.sum(mask > 127)
        change_ratio = (changed_pixels / total_pixels) * 100
        
        stats_text = f"""变化统计信息：
        
总像素数: {total_pixels:,}
变化像素数: {changed_pixels:,}
变化比例: {change_ratio:.2f}%
        """
        
        axes[1, 2].text(0.1, 0.5, stats_text, fontsize=11, 
                        verticalalignment='center', transform=axes[1, 2].transAxes,
                        bbox=dict(boxstyle="round,pad=0.3", facecolor="lightblue"))
        axes[1, 2].set_title('统计信息', fontsize=12)
        axes[1, 2].axis('off')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"可视化结果已保存到: {save_path}")
        
        plt.show()
        
        return fig
    
    def batch_visualize(self, img1_dir, img2_dir, mask_dir, output_dir):
        """
        批量可视化处理
        """
        os.makedirs(output_dir, exist_ok=True)
        
        mask_files = [f for f in os.listdir(mask_dir) if f.endswith('.png')]
        
        for mask_file in mask_files:
            base_name = os.path.splitext(mask_file)[0]
            
            img1_path = os.path.join(img1_dir, base_name + '.png')
            img2_path = os.path.join(img2_dir, base_name + '.png')
            mask_path = os.path.join(mask_dir, mask_file)
            save_path = os.path.join(output_dir, f'{base_name}_visualization.png')
            
            if os.path.exists(img1_path) and os.path.exists(img2_path):
                print(f"处理: {base_name}")
                self.visualize_change_detection(img1_path, img2_path, mask_path, save_path)
            else:
                print(f"跳过: {base_name} (缺少对应的影像文件)")

# 使用示例
if __name__ == "__main__":
    visualizer = SimpleChangeVisualizer()
    
    # 单个文件可视化示例
    # visualizer.visualize_change_detection(
    #     img1_path="path/to/image1.png",
    #     img2_path="path/to/image2.png", 
    #     mask_path="path/to/mask.png",
    #     save_path="path/to/result.png"
    # )
    
    # 批量处理示例
    img1_dir = "/media/data0/HL/2025-7-6/val/image1"
    img2_dir = "/media/data0/HL/2025-7-6/val/image2"
    mask_dir = "/media/data0/HL/CropLand-CD-main_3/CropLand-CD-main/save"
    output_dir = "/media/data0/HL/CropLand-CD-main_3/CropLand-CD-main/visualizations"
    
    visualizer.batch_visualize(img1_dir, img2_dir, mask_dir, output_dir)