#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
简化版图像相似度检测器
直接计算指定目录中图片之间的相似度，输出表格结果
"""

import os
import sys
import hashlib
import argparse
import time
import numpy as np
import cv2
from PIL import Image
import matplotlib.pyplot as plt
import pandas as pd
from skimage.metrics import structural_similarity as ssim
import imagehash
from datetime import datetime

class SimpleImageDupeChecker:
    """简化版图像相似度检测器"""
    
    def __init__(self):
        """初始化检测器"""
        # 算法列表
        self.algorithms = {
            'md5': self.compare_md5,
            'phash': self.compare_phash,
            'ahash': self.compare_ahash,
            'dhash': self.compare_dhash,
            'color_hist': self.compare_color_histogram,
            'sift': self.compare_sift,
            'ssim': self.compare_ssim
        }
        
        # 初始化SIFT检测器
        self.sift = cv2.SIFT_create()
        # FLANN匹配器参数
        FLANN_INDEX_KDTREE = 1
        index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
        search_params = dict(checks=50)
        self.flann = cv2.FlannBasedMatcher(index_params, search_params)

    def get_md5(self, image_path):
        """计算图像的MD5哈希值"""
        with open(image_path, 'rb') as f:
            md5hash = hashlib.md5(f.read()).hexdigest()
        return md5hash
    
    def compare_md5(self, image1_path, image2_path):
        """比较两张图像的MD5哈希值相似度"""
        hash1 = self.get_md5(image1_path)
        hash2 = self.get_md5(image2_path)
        
        # 如果哈希完全相同返回1.0，否则返回0.0
        return 1.0 if hash1 == hash2 else 0.0
    
    def get_phash(self, image_path, hash_size=8):
        """计算图像的感知哈希(pHash)"""
        img = Image.open(image_path).convert('L').resize((hash_size * 4, hash_size * 4), Image.LANCZOS)
        return imagehash.phash(img, hash_size=hash_size)
    
    def compare_phash(self, image1_path, image2_path):
        """比较两张图像的感知哈希(pHash)相似度"""
        hash1 = self.get_phash(image1_path)
        hash2 = self.get_phash(image2_path)
        
        # 计算哈希相似度（64位哈希的汉明距离）
        hash_diff = hash1 - hash2
        max_diff = 64.0  # 8x8哈希的最大汉明距离
        similarity = 1.0 - (hash_diff / max_diff)
        return max(0.0, similarity)  # 确保相似度不低于0
    
    def get_ahash(self, image_path, hash_size=8):
        """计算图像的平均哈希(aHash)"""
        img = Image.open(image_path).convert('L').resize((hash_size, hash_size), Image.LANCZOS)
        return imagehash.average_hash(img, hash_size=hash_size)
    
    def compare_ahash(self, image1_path, image2_path):
        """比较两张图像的平均哈希(aHash)相似度"""
        hash1 = self.get_ahash(image1_path)
        hash2 = self.get_ahash(image2_path)
        
        # 计算哈希相似度
        hash_diff = hash1 - hash2
        max_diff = 64.0  # 8x8哈希的最大汉明距离
        similarity = 1.0 - (hash_diff / max_diff)
        return max(0.0, similarity)
    
    def get_dhash(self, image_path, hash_size=8):
        """计算图像的差值哈希(dHash)"""
        img = Image.open(image_path).convert('L').resize((hash_size + 1, hash_size), Image.LANCZOS)
        return imagehash.dhash(img, hash_size=hash_size)
    
    def compare_dhash(self, image1_path, image2_path):
        """比较两张图像的差值哈希(dHash)相似度"""
        hash1 = self.get_dhash(image1_path)
        hash2 = self.get_dhash(image2_path)
        
        # 计算哈希相似度
        hash_diff = hash1 - hash2
        max_diff = 64.0  # 8x8哈希的最大汉明距离
        similarity = 1.0 - (hash_diff / max_diff)
        return max(0.0, similarity)
    
    def get_color_histogram(self, image_path):
        """计算图像的颜色直方图"""
        img = cv2.imread(image_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        
        # 计算HSV颜色空间的直方图
        hist = cv2.calcHist([img], [0, 1, 2], None, [8, 8, 8], [0, 180, 0, 256, 0, 256])
        cv2.normalize(hist, hist, 0, 1.0, cv2.NORM_MINMAX)
        return hist.flatten()
    
    def compare_color_histogram(self, image1_path, image2_path):
        """比较两张图像的颜色直方图相似度"""
        hist1 = self.get_color_histogram(image1_path)
        hist2 = self.get_color_histogram(image2_path)
        
        # 计算直方图相似度
        similarity = cv2.compareHist(hist1, hist2, cv2.HISTCMP_CORREL)
        return max(0.0, similarity)  # 确保相似度不低于0
    
    def get_sift_features(self, image_path):
        """获取图像的SIFT特征点和描述符"""
        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        keypoints, descriptors = self.sift.detectAndCompute(img, None)
        return keypoints, descriptors
    
    def compare_sift(self, image1_path, image2_path):
        """比较两张图像的SIFT特征点相似度"""
        kp1, des1 = self.get_sift_features(image1_path)
        kp2, des2 = self.get_sift_features(image2_path)
        
        # 如果没有足够的特征点，返回低相似度
        if des1 is None or des2 is None or len(des1) < 2 or len(des2) < 2:
            return 0.0
        
        # 使用FLANN匹配器查找最佳匹配
        matches = self.flann.knnMatch(des1, des2, k=2)
        
        # 应用比率测试，筛选好的匹配
        good_matches = []
        for m, n in matches:
            if m.distance < 0.7 * n.distance:
                good_matches.append(m)
        
        # 计算相似度
        similarity = len(good_matches) / max(len(kp1), len(kp2)) if max(len(kp1), len(kp2)) > 0 else 0
        return min(1.0, similarity)  # 确保相似度不超过1
    
    def compare_ssim(self, image1_path, image2_path):
        """比较两张图像的结构相似性(SSIM)"""
        img1 = cv2.imread(image1_path, cv2.IMREAD_GRAYSCALE)
        img2 = cv2.imread(image2_path, cv2.IMREAD_GRAYSCALE)
        
        # 确保两张图像尺寸相同
        h, w = min(img1.shape[0], img2.shape[0]), min(img1.shape[1], img2.shape[1])
        img1 = cv2.resize(img1, (w, h))
        img2 = cv2.resize(img2, (w, h))
        
        # 计算SSIM
        similarity = ssim(img1, img2)
        return max(0.0, similarity)  # 确保相似度不低于0

    def check_images(self, input_dir, output_dir=None):
        """
        检查目录中所有图像之间的相似度
        
        Args:
            input_dir: 输入图像目录
            output_dir: 输出结果目录，如果为None则使用input_dir
            
        Returns:
            包含相似度信息的DataFrame
        """
        # 如果未指定输出目录，使用输入目录
        if output_dir is None:
            output_dir = input_dir
        os.makedirs(output_dir, exist_ok=True)
        
        # 获取所有图像文件
        print("正在读取图像文件...")
        image_files = []
        for file in os.listdir(input_dir):
            file_path = os.path.join(input_dir, file)
            if os.path.isfile(file_path) and file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.gif')):
                image_files.append(file_path)
        
        # 检查图像数量
        if len(image_files) == 0:
            print(f"错误: 在目录 '{input_dir}' 中没有找到图像文件")
            return None
        
        print(f"找到 {len(image_files)} 个图像文件，开始计算相似度...")
        
        # 结果列表
        results = []
        
        # 计算所有图像对的相似度
        total_pairs = len(image_files) * (len(image_files) - 1) // 2
        pair_count = 0
        
        for i, img1 in enumerate(image_files):
            img1_name = os.path.basename(img1)
            
            for j, img2 in enumerate(image_files):
                if j <= i:  # 跳过重复比较
                    continue
                    
                img2_name = os.path.basename(img2)
                pair_count += 1
                print(f"比较图像对 {pair_count}/{total_pairs}: {img1_name} vs {img2_name}")
                
                # 对每种算法计算相似度
                for alg_name, compare_func in self.algorithms.items():
                    try:
                        start_time = time.time()
                        similarity = compare_func(img1, img2)
                        compute_time = time.time() - start_time
                        
                        results.append({
                            'image1': img1_name,
                            'image2': img2_name,
                            'algorithm': alg_name,
                            'similarity': similarity,
                            'compute_time': compute_time
                        })
                        
                    except Exception as e:
                        print(f"  算法 {alg_name} 比较出错: {e}")
                        results.append({
                            'image1': img1_name,
                            'image2': img2_name,
                            'algorithm': alg_name,
                            'similarity': None,
                            'compute_time': None,
                            'error': str(e)
                        })
        
        # 创建DataFrame
        df = pd.DataFrame(results)
        
        # 保存结果
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        csv_path = os.path.join(output_dir, f"similarity_results_{timestamp}.csv")
        df.to_csv(csv_path, index=False)
        print(f"结果已保存至: {csv_path}")
        
        # 创建透视表
        pivot_df = df.pivot_table(
            index=['image1', 'image2'],
            columns='algorithm',
            values='similarity',
            aggfunc='first'
        ).reset_index()
        
        # 保存透视表
        pivot_path = os.path.join(output_dir, f"similarity_pivot_{timestamp}.csv")
        pivot_df.to_csv(pivot_path, index=False)
        print(f"透视表已保存至: {pivot_path}")
        
        # 显示结果摘要
        self.print_summary(df)
        
        # 生成可视化
        self.visualize_results(df, output_dir)
        
        return df
    
    def print_summary(self, df):
        """打印结果摘要"""
        print("\n===== 图像相似度检测结果摘要 =====")
        
        # 按算法分组计算平均相似度
        alg_summary = df.groupby('algorithm')['similarity'].agg(['mean', 'min', 'max', 'std']).reset_index()
        alg_summary = alg_summary.sort_values('mean', ascending=False)
        
        print("\n各算法平均相似度:")
        for _, row in alg_summary.iterrows():
            print(f"  {row['algorithm']:<12}: 平均值 = {row['mean']:.4f} (最小值: {row['min']:.4f}, 最大值: {row['max']:.4f}, 标准差: {row['std']:.4f})")
        
        # 显示最相似的图像对
        print("\n相似度最高的图像对:")
        for alg in df['algorithm'].unique():
            alg_df = df[df['algorithm'] == alg]
            if not alg_df.empty:
                max_idx = alg_df['similarity'].idxmax()
                max_row = alg_df.loc[max_idx]
                print(f"  {alg:<12}: {max_row['image1']} 和 {max_row['image2']} (相似度: {max_row['similarity']:.4f})")
        
        # 显示最不相似的图像对
        print("\n相似度最低的图像对:")
        for alg in df['algorithm'].unique():
            alg_df = df[df['algorithm'] == alg]
            if not alg_df.empty:
                min_idx = alg_df['similarity'].idxmin()
                min_row = alg_df.loc[min_idx]
                print(f"  {alg:<12}: {min_row['image1']} 和 {min_row['image2']} (相似度: {min_row['similarity']:.4f})")
    
    def visualize_results(self, df, output_dir):
        """可视化结果"""
        print("\n生成结果可视化...")
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 1. 绘制条形图 - 各算法的平均相似度
        plt.figure(figsize=(10, 6))
        avg_similarity = df.groupby('algorithm')['similarity'].mean().reset_index()
        avg_similarity = avg_similarity.sort_values('similarity', ascending=False)
        
        plt.bar(avg_similarity['algorithm'], avg_similarity['similarity'], color='skyblue')
        plt.xlabel('算法')
        plt.ylabel('平均相似度')
        plt.title('各算法的平均相似度')
        plt.xticks(rotation=45)
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.tight_layout()
        
        bar_chart_path = os.path.join(output_dir, f"algorithm_avg_similarity_{timestamp}.png")
        plt.savefig(bar_chart_path)
        
        # 2. 绘制热力图 - 各图像对的相似度矩阵
        # 为每个算法生成一个热力图
        image_names = sorted(list(set(df['image1'].unique()) | set(df['image2'].unique())))
        
        for alg in df['algorithm'].unique():
            plt.figure(figsize=(10, 8))
            
            # 创建热力图数据
            heatmap_data = np.zeros((len(image_names), len(image_names)))
            
            # 填充热力图数据
            for idx, row in df[df['algorithm'] == alg].iterrows():
                i = image_names.index(row['image1'])
                j = image_names.index(row['image2'])
                heatmap_data[i, j] = row['similarity']
                heatmap_data[j, i] = row['similarity']  # 对称矩阵
            
            # 热力图对角线设为1.0（自己与自己相似度为1）
            for i in range(len(image_names)):
                heatmap_data[i, i] = 1.0
            
            # 创建热力图
            plt.imshow(heatmap_data, cmap='viridis')
            plt.colorbar(label='相似度')
            plt.title(f'{alg} 算法的图像相似度热力图')
            
            # 设置坐标轴刻度
            plt.xticks(np.arange(len(image_names)), image_names, rotation=90)
            plt.yticks(np.arange(len(image_names)), image_names)
            
            # 添加文本标签
            for i in range(len(image_names)):
                for j in range(len(image_names)):
                    text = plt.text(j, i, f"{heatmap_data[i, j]:.2f}",
                                  ha="center", va="center", 
                                  color="w" if heatmap_data[i, j] < 0.7 else "black")
            
            plt.tight_layout()
            heatmap_path = os.path.join(output_dir, f"similarity_heatmap_{alg}_{timestamp}.png")
            plt.savefig(heatmap_path)
        
        # 3. 箱线图 - 各算法的相似度分布
        plt.figure(figsize=(10, 6))
        
        # 准备数据
        data = []
        labels = []
        
        for alg in sorted(df['algorithm'].unique()):
            alg_data = df[df['algorithm'] == alg]['similarity']
            if not alg_data.empty:
                data.append(alg_data)
                labels.append(alg)
        
        # 创建箱线图
        plt.boxplot(data, labels=labels)
        plt.ylabel('相似度')
        plt.title('各算法的相似度分布')
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.tight_layout()
        
        boxplot_path = os.path.join(output_dir, f"algorithm_boxplot_{timestamp}.png")
        plt.savefig(boxplot_path)
        
        # 关闭所有图形
        plt.close('all')
        
        print(f"可视化图表已保存到:\n  条形图: {bar_chart_path}\n  箱线图: {boxplot_path}")
        print(f"  各算法热力图已保存到输出目录")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='简化版图像相似度检测工具')
    parser.add_argument('--input_dir', '-i', type=str, default='./scripts/image_test/test_images',
                        help='输入图像目录路径')
    parser.add_argument('--output_dir', '-o', type=str, default=None,
                        help='输出结果目录路径，默认与输入目录相同')
    
    args = parser.parse_args()
    
    print("===== 简化版图像相似度检测工具 =====")
    print(f"输入目录: {args.input_dir}")
    print(f"输出目录: {args.output_dir or args.input_dir}")
    
    checker = SimpleImageDupeChecker()
    checker.check_images(args.input_dir, args.output_dir)
    
    print("\n===== 检测完成 =====")

if __name__ == '__main__':
    main() 