# Copyright (C) 2018-2021 coneypo
# SPDX-License-Identifier: MIT

import os
import dlib
import csv
import numpy as np
import logging
import cv2
from PIL import Image

# 要读取人脸图像文件的路径
path_images_from_camera = "data/data_faces"

# Dlib 检测器
detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor('models/dlib/shape_predictor_68_face_landmarks.dat')
face_reco_model = dlib.face_recognition_model_v1("models/dlib/dlib_face_recognition_resnet_model_v1.dat")


def return_128d_features(path_img):
    """返回单张图像的 128D 特征"""
    try:
        img_pil = Image.open(path_img)
        img_np = np.array(img_pil)
        img_rd = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
        faces = detector(img_rd, 1)

        logging.info("%-40s %-20s", "检测到人脸的图像:", path_img)

        if len(faces) != 0:
            shape = predictor(img_rd, faces[0])
            face_descriptor = face_reco_model.compute_face_descriptor(img_rd, shape)
            return face_descriptor
        else:
            logging.warning("未检测到人脸: %s", path_img)
            return None
    except Exception as e:
        logging.error("处理图像时出错 %s: %s", path_img, e)
        return None


def return_features_mean_personX(path_face_personX):
    """返回 personX 的 128D 特征均值"""
    features_list_personX = []
    photos_list = os.listdir(path_face_personX)

    if photos_list:
        for photo in photos_list:
            photo_path = os.path.join(path_face_personX, photo)
            logging.info("正在读取图像: %s", photo_path)

            features_128d = return_128d_features(photo_path)
            if features_128d is not None:
                features_list_personX.append(features_128d)
    else:
        logging.warning("文件夹为空: %s", path_face_personX)

    # 计算 128D 特征的均值
    if features_list_personX:
        features_mean_personX = np.array(features_list_personX).mean(axis=0)
    else:
        features_mean_personX = np.zeros(128, dtype=np.float64)

    return features_mean_personX


def get_person_name_from_folder(folder_name):
    """从文件夹名称获取有意义的姓名"""
    # 常见的文件夹前缀
    prefixes = ['person_', 'face_', 'user_']

    for prefix in prefixes:
        if folder_name.startswith(prefix):
            name_part = folder_name[len(prefix):]
            # 如果剩下的部分是纯数字，使用完整文件夹名
            if name_part.isdigit():
                return folder_name
            else:
                return name_part

    return folder_name


def main():
    logging.basicConfig(level=logging.INFO)

    # 检查源文件夹是否存在
    if not os.path.exists(path_images_from_camera):
        logging.error("人脸图像文件夹不存在: %s", path_images_from_camera)
        return

    # 获取人脸文件夹列表
    person_list = os.listdir(path_images_from_camera)
    person_list.sort()

    if not person_list:
        logging.error("没有人脸文件夹可处理")
        return

    logging.info("找到 %d 个人脸文件夹: %s", len(person_list), person_list)

    # 创建CSV文件
    with open("data/features_all.csv", "w", newline="", encoding="utf-8") as csvfile:
        writer = csv.writer(csvfile)

        successful_count = 0
        for person_folder in person_list:
            folder_path = os.path.join(path_images_from_camera, person_folder)

            if not os.path.isdir(folder_path):
                continue

            logging.info("处理文件夹: %s", person_folder)

            # 提取特征
            features_mean = return_features_mean_personX(folder_path)

            # 获取有意义的姓名
            person_name = get_person_name_from_folder(person_folder)
            logging.info("使用姓名: %s", person_name)

            # 构建行数据：姓名 + 128维特征
            row_data = [person_name] + features_mean.tolist()
            writer.writerow(row_data)

            successful_count += 1
            logging.info("完成: %s", person_name)
            logging.info("-" * 50)

        logging.info("成功处理 %d/%d 个人脸文件夹", successful_count, len(person_list))
        logging.info("特征数据已保存到: data/features_all.csv")


if __name__ == '__main__':
    main()