import cv2
import insightface
import numpy as np
import os
import datetime
import argparse
from insightface.app import FaceAnalysis

# 设备配置映射（NPU采用用户指定的完整参数）
DEVICE_CONFIG = {
    "cpu": (['CPUExecutionProvider'], -1),
	"gpu": (['CUDAExecutionProvider'], 0),
	"npu": (
	    [
		    (
			    "CANNExecutionProvider",
				{
				    "device_id": 1,
					"arena_extend_strategy": "kNextPowerOfTwo",
					"npu_mem_limit": 16*1024*1024*1024,
					"op_select_impl_mode": "high_precision",
					"precision_mode": "allow_fp32_to_fp16",
					"enable_cann_graph": True,

				},
			),
			"CPUExecutionProvider"
		],
		0
	)
}
#allow_fp32_to_fp16
#核心配置参数
THRESHOLD = 0.65
IMAGE_EXTENSIONS = ('.jpg','.jepg','.png','.bmp','.gif')
NPU_REQUIREMENTS = {
    "依赖包": "onnxruntime-cann（华为官方）+ onnxruntime（基础）",
	"驱动要求": "Ascend CANN Toolkit ≥ 5.0.3",
	"硬件要求": "华为昇腾芯片（如Ascend 310/910）",
	"文档链接": "https://onnxruntime.ai/docs/execution-providers/community-maintained/CANN-ExecutionProvider.html"
}

# Initialize face analysis model
#app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])  # Use 'CUDAExecutionProvider' for GPU
#app.prepare(ctx_id=-1)  # ctx_id=-1 for CPU, 0 for GPU

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="人脸特征提取与同一人照片分组工具（支持CPU/GPU/NPU）")
    parser.add_argument("-d","--device", type=str, choices=DEVICE_CONFIG.keys(),
        default="cpu", help=f"指定运行设备")
    parser.add_argument("-t","--threshold", type=float, default=THRESHOLD, help=f"相似度阈值，值越大匹配越严格")
    parser.add_argument("--npu-device-id", type=int, default="0", help=f"覆盖NPU设备ID")
    parser.add_argument("--npu-mem-limit", type=int, default="16", help=f"NPU内存限制")
    args = parser.parse_args()
    return args

def get_face_embedding(image_path, app):
    """Extract face embedding from an image"""
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"Could not read image: {image_path}")
    
    faces = app.get(img)
    
    if len(faces) < 1:
        raise ValueError("No faces detected in the image")
    if len(faces) > 1:
        print("Warning: Multiple faces detected. Using first detected face")
    
    return faces[0].embedding

def compare_faces(emb1, emb2, threshold): # Adjust this threshold according to your usecase.
    """Compare two embeddings using cosine similarity"""
    similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
    return similarity, similarity > threshold
	
def get_all_images_files(directory="."):
    """获取目录下所有支持的图片文件"""
    image_files	= []
    for filename in os.listdir(directory):
        file_path = os.path.abspath(filename)
        if filename.lower().endswith(IMAGE_EXTENSIONS) and os.path.isfile(file_path):
            image_files.append(file_path)
	
    return sorted(image_files)

def find_same_person_groups(embedding_dict, threshold):
    """基于特征向量分组同一人照片"""
    ungrouped_files = list(embedding_dict.keys())
    same_person_groups = []
    single_groups = []

    while ungrouped_files:
        current_file =ungrouped_files.pop(0)
        current_emb = embedding_dict[current_file]
        current_group = [current_file]
        to_remove = []
		
        for candidate_file in ungrouped_files:
            candidate_emb = embedding_dict[candidate_file]
            similarity, is_same = compare_faces(current_emb, candidate_emb, threshold)
            if is_same:
                current_group.append(candidate_file)
                to_remove.append(candidate_file)
                print(f"匹配成功：{os.path.basename(current_file)}和{os.path.basename(candidate_file)}（相似度：{similarity:.4f}）")
        for file in to_remove:
            ungrouped_files.remove(file)
		
        if len(current_group) >= 2:
            same_person_groups.append(current_group)
        else:
            single_groups.append(current_group)
	
    return same_person_groups, single_groups

def generate_log_filename(device):
    """生成带设备类型和时间戳的日志文件名"""
    timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
    return f"face_embeddings_{device}_{timestamp}.log"
def write_embedding_log(log_entries, args, npu_config, log_filename):
    """写入详细日志"""
    with open(log_filename, 'w', encoding='utf-8') as f:
        f.write("="*80 + "\n")
        f.write(f"人脸特征向量分析日志\n")
        f.write(f"生成时间： {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"运行设备：{args.device}\n")
        if args.device == "npu":
            f.write(f"NPU配置：{npu_config}\n")
        f.write(f"相似度阈值：{args.threshold}\n")
        f.write(f"支持图片格式：{','.join(IMAGE_EXTENSIONS)}\n")
        f.write(f"处理图片总数：{len(log_entries)}\n")
        f.write("="*80 + "\n\n")
		
        for entry in log_entries:
            f.write(entry + "\n" + "-"*60 + "\n\n")
			
    print(f"\n 日志文件已保存至：{os.path.abspath(log_filename)}")
	
def check_npu_environment():
    """检查NPU环境是否满足要求（仅检查核心依赖，不验证环境变量）"""
    try:
        import onnxruntime as ort
        available_providers = ort.get_available_providers()
        if 'CANNExecutionProvider' not in available_providers:
            return False, "未检测到CANNExecutionProvider（请安装onnxruntime-cann包）"

        return True,"NPU核心依赖检查通过"
    except ImportError:
        return False, "未安装onnxruntime（基础依赖）"
    except Exception as e:
        return False, f"NPU环境检查失败： {str(e)}"

def main():
    args = parse_args()
    device = args.device
    threshold = args.threshold
    npu_device_id = args.npu_device_id	
    npu_mem_limit = args.npu_mem_limit*1024*1024*1024
    #生成带设备和时间戳的日志文件名
    log_filename = generate_log_filename(device)

    #打印启动信息
    print("人脸特征提取与同一人照片分组工具（支持CPU/GPU/NPU）")
    print("="*80)
    print(f"核心配置：")
    print(f"运行设备：{device}")
    print(f"相似度阈值：{threshold}")
    print(f"搜索目录：{os.getcwd()}")
    print(f"支持格式：{','.join(IMAGE_EXTENSIONS)}")
    print(f"日志文件：{log_filename}")
    if device == "npu":
        print(f"NPU设备ID：{npu_device_id}")
        print(f"NPU内存限制：{args.npu_mem_limit}GB")
    print("="*80 + "\n")
	
    #设备环境预检查
    if device == "gpu":
        available_providers = insightface.utils.get_available_providers()
        if 'CUDAExecutionProvider' not in available_providers:
            print(" 警告：未检测到CUDA环境，GPU模式可能运行失败！")
            print(" 解决方案：1.安装CUDA≥11.0 + cuDNN 2.安装onnxruntime-gpu 3. 切换至cpu模式")
        elif device == "npu":
            npu_ok,npu_msg = check_npu_environment()
            print(f" NPU环境检查：{npu_msg}")
            if not npu_ok:
                print(" NPU环境不满足要求，建议按以下步骤配置：")
                for key, value in NPU_REQUIREMENTS.items():
                    print(f"  - {key}:{value}")
                return
				
				
    #初始化模型
    try:
        providers, ctx_id = DEVICE_CONFIG[device]
        npu_config = None
		
        if device == "npu":
            npu_provider = list(providers[0])
            npu_provider[1]["device_id"] = npu_device_id
            npu_provider[1]["npu_mem_limit"] = npu_mem_limit
            providers[0] = tuple(npu_provider)
            npu_config = providers[0][1]
            ctx_id = npu_device_id
            app = FaceAnalysis(name='buffalo_l', providers=providers)
        else:
            app = FaceAnalysis(name='buffalo_l', providers=providers)
			
        app.prepare(ctx_id=ctx_id)
        print(f" 模型初始化成功（设备：{device}, ctx_id: {ctx_id}）")
        if device == "npu":
            print(f" NPU最终配置：")
            for key, value in npu_config.items():
                if key == "npu_mem_limit":
                    print(f"  - {key}: {value/(1024*1024*1024)}GB")
                else:
                    print(f"  - {key}: {value}")
    except Exception as e:
        print(f"模型初始化失败！")
        print(f"Error:{str(e)}")
        print(f" 解决方案:")
        if device == "gpu":
            print("   1.确认CUDA驱动已安装 2.确认onnxruntime-gpu版本与CUDA匹配 3.尝试切换CPU模式")
        elif device == "npu":
            print("   1.确认CANN Toolkit已正确安装 2.确认onnxruntime-cann版本兼容 3.检查设备ID和内存限制是否合理")
        return

    #1. 获取所有图片文件
    image_files = get_all_images_files()
    if not image_files:
        print("未找到任何支持的图片文件（检查目录下是否有jpg/png等格式图片）")
        return
    print(f"  找到{len(image_files)}个图片文件，开始提取特征向量...\n")

    #2. 提取特征向量并记录日志
    embedding_dict = {}
    log_entries = []
    for img_path in image_files:
        img_name = os.path.basename(img_path)
        try:
            emb = get_face_embedding(img_path, app)
            embedding_dict[img_name] = emb
            log_entry = f"【文件】：{img_name}\n" \
                        f"【路径】：{img_path}\n" \
                        f"【状态】：成功\n" \
                        f"【特征向量维度】：{len(emb)}\n" \
                        f"【特征向量】：{emb.tolist()}"
            log_entries.append(log_entry)
            print(f"  处理成功： {img_name}（特征向量维度：{len(emb)})")
        except Exception as e:
            log_entry = f"【文件】：{img_name}\n" \
                        f"【路径】：{img_path}\n" \
                        f"【状态】：失败\n" \
                        f":【错误信息】:{str(e)}"
            log_entries.append(log_entry)
            print(f"  处理失败： {img_name} - 原因：{str(e)}")

    #3. 写入日志文件（使用动态生成的文件名）
    write_embedding_log(log_entries, args, npu_config, log_filename)

    #4. 同一人分组分析    
    if not embedding_dict:
        print("\n 没有成功提取到人脸特征向量，无法进行分组分析")
        return
    
    print(f"\n 开始分组分析（有效人脸数：{len(embedding_dict)}）...")
    same_groups, single_groups = find_same_person_groups(embedding_dict, threshold)
    
    #5. 输出分组结果
    print("\n" + "="*80)
    print("同一人照片分组结果（每组≥2张）:")
    print("="*80)
    if same_groups:
        for i, group in enumerate(same_groups, 1):
            group_names = [os.path.basename(file) for file in group]
            print(f"组{i}:{', '.join(group_names)}")
    else:
        print(f"  未发现同一人的多张照片")
        
    print("\n" + "="*80)
    print("无匹配的单独照片：")
    #print("\n"*80)
    if single_groups:
        for group in single_groups:
            print(f"  - {os.path.basename(group[0])}")
    else:
        print(f"所有照片均已分组（无单独照片）")
    print("\n" + "="*80)
    print(f" 处理完成！详细日志请查看：{log_filename}")
    
if __name__ == "__main__":
    main()
# Paths to your Indian face images
#image1_path = "path/to/face1.jpg"
#image2_path = "path/to/face2.jpg"

#try:
    # Get embeddings
#    emb1 = get_face_embedding(image1_path)
#    emb2 = get_face_embedding(image2_path)
    
    # Compare faces
#    similarity_score, is_same_person = compare_faces(emb1, emb2)
    
#    print(f"Similarity Score: {similarity_score:.4f}")
#    print(f"Same person? {'YES' if is_same_person else 'NO'}")
    
#except Exception as e:
#    print(f"Error: {str(e)}")