import os
import time
import re
import torch
import cv2
import numpy as np

from anti import anti_spoofing, load_anti_model
from backbones import iresnet50, iresnet18, iresnet100
from retinaface_detect import load_retinaface_model, detect_one, detect_video, set_retinaface_conf
from torch2trt import torch2trt, TRTModule

threshold = 0.7


# 读取112x112的本地图片并变换通道位置归一化
def load_image(img_path):
    img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    img = img.transpose((2, 0, 1))
    img = img[np.newaxis, :, :, :]
    img = np.array(img, dtype=np.float32)
    img -= 127.5
    img /= 127.5
    return img


# 计算两个特征向量的欧式距离
def findEuclideanDistance(source_representation, test_representation):
    euclidean_distance = source_representation - test_representation
    euclidean_distance = np.sum(np.multiply(euclidean_distance, euclidean_distance))
    euclidean_distance = np.sqrt(euclidean_distance)
    return euclidean_distance


# 计算两个特征向量的余弦距离
def findCosineDistance(source_representation, test_representation):
    a = np.matmul(np.transpose(source_representation), test_representation)
    b = np.sum(np.multiply(source_representation, source_representation))
    c = np.sum(np.multiply(test_representation, test_representation))
    return 1 - (a / (np.sqrt(b) * np.sqrt(c)))


# 归一化欧氏距离
def l2_normalize(x):
    return x / np.sqrt(np.sum(np.multiply(x, x)))


# 归一化余弦距离
def cosin_metric(x1, x2):
    return np.dot(x1, x2) / (np.linalg.norm(x1) * np.linalg.norm(x2))


# 加载保存的姓名、人脸特征向量的人脸库
def load_npy(path):
    data = np.load(path, allow_pickle=True)
    data = data.item()
    return data


# 批量化生成人脸特征向量并保存到人脸库
def create_database_batch(path, model, database_path):
    name_list = os.listdir(path)
    k_v = {}
    if os.path.exists(database_path):
        k_v = np.load(database_path, allow_pickle=True)
        k_v = k_v.item()
    batch = 256
    order_name = []
    order_path = []
    emb_list = []
    for name in name_list[:]:
        img_path = os.path.join(path, name)
        # for img_name in img_path[:1]:
        order_name.append(name[:-4])
        order_path.append(img_path)
    order_img = np.zeros((len(order_path), 3, 112, 112), dtype=np.float32)
    for index, img_path in enumerate(order_path):
        order_img[index] = load_image(img_path)
    print(order_img.shape)
    order_img = torch.from_numpy(order_img)
    order_img = order_img.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
    now = 0
    number = len(order_img)
    with torch.no_grad():
        while now < number:
            if now + batch < number:
                emb = model(order_img[now:now + batch])
            else:
                emb = model(order_img[now:])
            now = now + batch
            emb = emb.cpu().numpy()
            for em in emb:
                emb_list.append(em)
            print("batch" + str(now))

    for i, emb in enumerate(emb_list):
        k_v[order_name[i]] = l2_normalize(emb)
    np.save(database_path, k_v)

def create_database_from_img(order_name, order_img, model, database_path, cpu_or_cuda):
    k_v = {}
    if os.path.exists(database_path):
        k_v = np.load(database_path, allow_pickle=True)
        k_v = k_v.item()
    batch = 256
    emb_list = []

    print(order_img.shape)
    order_img = torch.from_numpy(order_img)
    order_img = order_img.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
    now = 0
    number = len(order_img)
    with torch.no_grad():
        while now < number:
            if now + batch < number:
                emb = model(order_img[now:now + batch])
            else:
                emb = model(order_img[now:])
            now = now + batch
            emb = emb.cpu().numpy()
            for em in emb:
                emb_list.append(em)
            print("batch" + str(now))
    for i, emb in enumerate(emb_list):
        k_v[order_name[i]] = l2_normalize(emb)
    np.save(database_path, k_v)

# 向人脸库中新增一个人的姓名和人脸特征向量，若人脸库不存在则创建
def add_one_to_database(img, model, name, database_path, cpu_or_cuda):
    img = torch.from_numpy(img)
    img = img.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
    with torch.no_grad():
        pred = model(img)
        pred = pred.cpu().numpy()
        k_v = {}
        if os.path.exists(database_path):
            k_v = np.load(database_path, allow_pickle=True)
            k_v = k_v.item()
        k_v[name] = l2_normalize(pred)
        np.save(database_path, k_v)


# 计算此特征向量与人脸库中的哪个人脸特征向量距离最近
def findmindistance(pred, threshold, k_v):
    distance = 10
    most_like = ""
    for name in k_v.keys():
        tmp = findEuclideanDistance(k_v[name], pred)
        if distance > tmp:
            distance = tmp
            most_like = name
    if distance < threshold:
        return most_like, distance
    else:
        return -1, distance


def faiss_find_face(pred, index, database_name_list):
    name_list = []
    start_time = time.time()
    D, I = index.search(pred, 1)
    end_time = time.time()
    # print("faiss cost %fs" % (end_time - start_time))
    # print(D, I)
    if len(pred) == 1:
        if D[0][0] < threshold:
            # print(database_name_list[I[0][0]])
            return database_name_list[I[0][0]], D[0][0]
        else:
            return "unknown", D[0][0]
    else:
        for i,index in enumerate(I):
            if D[i][0] < threshold:
                #print(database_name_list[I[0][0]])
                name_list.append(database_name_list[index[0]]+str(D[i][0]))
            else:
                name_list.append("unknown"+str(D[i][0]))
        return name_list


# 从人脸库中找到单个人脸
def findOne(img, model, index, database_name_list, cpu_or_cuda):
    img = torch.from_numpy(img)
    img = img.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
    with torch.no_grad():
        start_time = time.time()
        pred = model(img)
        end_time = time.time()
        print("predOne time: " + str(end_time - start_time))
        pred = pred.cpu().numpy()
        # start_time = time.time()
        # name, distance = findmindistance(l2_normalize(pred), threshold=threshold, k_v=k_v)
        # end_time = time.time()
        # print("baoli time: " + str(end_time - start_time))
        name, distance = faiss_find_face(l2_normalize(pred), index, database_name_list)
        print(pred.shape)
        if name != -1:
            mo = r'[\u4e00-\u9fa5_a-zA-Z0-9]*'
            name = re.match(mo, name)
            return name.group(0), distance
        else:
            return "unknown", distance


# 从人脸库中找到传入的人脸列表中的所有人脸
def findAll(imglist, model, index ,database_name_list, cpu_or_cuda):
    imglist = torch.from_numpy(imglist)
    imglist = imglist.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
    with torch.no_grad():
        name_list =[]
        start_time = time.time()
        pred = model(imglist)
        end_time = time.time()
        print("predOne time: " + str(end_time - start_time))
        pred = pred.cpu().numpy()
        start_time = time.time()
        #name_list = faiss_find_face(l2_normalize(pred), index, database_name_list)
        for pr in pred:
            pr = np.expand_dims(l2_normalize(pr), 0)
            #     #print(pr.shape)
            name, distance = faiss_find_face(l2_normalize(pr), index, database_name_list)
            #name_list.append(name+" "+str(distance))
            name_list.append(name)
        # for pr in pred:
        #     name, distance = findmindistance(l2_normalize(pr), threshold=threshold, k_v=k_v)
        #     if name != -1:
        #         mo = r'[\u4e00-\u9fa5_a-zA-Z]*'
        #         name = re.match(mo, name)
        #         name_list.append(name.group(0) + str(distance))
        #     else:
        #         name_list.append("unknown" + str(distance))
        end_time = time.time()
        print("searchALL time: " + str(end_time - start_time))
        return name_list


# 提取为512维特征向量
def embedding(order_img, model, cpu_or_cuda):
    number = len(order_img)
    order_img = torch.from_numpy(order_img)
    order_img = order_img.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
    batch = 64
    emb_list = []
    now = 0
    with torch.no_grad():
        while now < number:
            if now + batch < number:
                emb = model(order_img[now:now + batch])
            else:
                emb = model(order_img[now:])
            now = now + batch
            emb = emb.cpu().numpy()
            for em in emb:
                emb_list.append(l2_normalize(em))
            # print("batch" + str(now))
        emb_list = np.array(emb_list)
        return emb_list


# 处理聚类人脸文件夹，返回特征向量列表，文件名列表
def get_claster_tmp_file_embedding(file_path, retinaface_model, retinaface_args, arcface_model, cpu_or_cuda):
    img_name = os.listdir(file_path)
    img_list = []
    for name in img_name:
        all_face, box_and_point = detect_one(os.path.join(file_path, name), retinaface_model, retinaface_args)
        img_list.append(all_face[0])
    img_list = np.array(img_list)
    # print(img_list.shape)
    emb_list = embedding(img_list, arcface_model, cpu_or_cuda)
    return emb_list, img_name


# 同一个人聚为一类
def cluster(emb_list, name_list):
    all_claster = []
    cla = []
    in_claster_name = []
    img_number = len(emb_list)
    for index, emb in enumerate(emb_list):
        if name_list[index] in in_claster_name:
            continue
        for j in range(img_number - index - 1):
            if findEuclideanDistance(emb, emb_list[index + 1 + j]) < threshold:
                if name_list[index + 1 + j] not in in_claster_name:
                    cla.append(name_list[index + 1 + j])
                    in_claster_name.append(name_list[index + 1 + j])
        cla.append(name_list[index])
        in_claster_name.append(name_list[index])
        all_claster.append(cla)
        cla = []
    return all_claster


# 加载人脸识别模型
def load_arcface_model(model_path, cpu_or_cuda):
    if cpu_or_cuda == "trt":
        model = TRTModule()
        model.load_state_dict(torch.load('./model/arcface_trt.pth'))
    elif cpu_or_cuda == "trt_new":
        model = iresnet100()
        model.load_state_dict(torch.load(model_path, map_location="cuda"))
        model = model.eval()
        model.to(torch.device("cuda"))
        x = torch.ones((1, 3, 112, 112)).to(torch.device("cuda"))
        model = torch2trt(model, [x], max_batch_size=4)
        torch.save(model.state_dict(), './model/arcface_trt.pth')
    else:
        model = iresnet100()
        model.load_state_dict(torch.load(model_path, map_location=cpu_or_cuda))
        model = model.eval()
        model.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
    return model


# 对比两张人脸是否相同
def face_verification(img1, img2, model, cpu_or_cuda):
    img_list = np.concatenate((img1, img2), axis=0)
    img_list = torch.from_numpy(img_list)
    img_list = img_list.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
    with torch.no_grad():
        pred = model(img_list)
        pred = pred.cpu().numpy()
        distance = findEuclideanDistance(l2_normalize(pred[0]), l2_normalize(pred[1]))
        # print("EuclideanDistance is :" + str(distance))
        if distance < threshold:
            return 'same ',distance
        else:
            return 'different ', distance


if __name__ == '__main__':
    cpu_or_cuda = "cuda" if torch.cuda.is_available() else "cpu"
    arcface_model = load_arcface_model("./model/backbone100.pth", cpu_or_cuda=cpu_or_cuda)
    # retinaface_args = set_retinaface_conf(cpu_or_cuda=cpu_or_cuda)
    # retinaface_model = load_retinaface_model(retinaface_args)
    #
    # anti_spoofing_model_path = "model/anti_spoof_models"
    # anti_model = load_anti_model(anti_spoofing_model_path, 0)
    #
    # k_v = load_npy("./Database/student.npy")
    # 对比两张人脸
    # img1, box_and_point = detect_one("D:\Download\lfw\lfw\Aaron_Peirsol\Aaron_Peirsol_0001.jpg", retinaface_model, retinaface_args)
    # img2, box_and_point = detect_one("D:\Download\lfw\lfw\Aaron_Peirsol\Aaron_Peirsol_0002.jpg", retinaface_model, retinaface_args)
    # print(face_verification(img1, img2, arcface_model))

    # img3 = load_image(r"D:\Download\out\alig_students\student.jpg")
    # img3 = torch.from_numpy(img3)
    # 单张人脸活体检测
    # img3, b_p = detect_one(r"C:\Users\ASUS\Desktop\face\IMG_20210525_113950.jpg", retinaface_model, retinaface_args)
    # b = b_p[0]
    # w = b[2] - b[0]
    # h = b[3] - b[1]
    # b[2] = w
    # b[3] = h
    # label, value = anti_spoofing("./img/recognition/000_0.bmp", "model/anti_spoof_models", 0, np.array(b[:4], int), anti_model)
    # print(label,value)
    # name = findOne(img3, arcface_model, k_v, cpu_or_cuda)
    # print(name)

    # 人脸聚类
    # emb_list, name_list = get_claster_tmp_file_embedding("./img/cluster_tmp_file/face", retinaface_model,
    #                                                      retinaface_args, arcface_model, cpu_or_cuda)
    # print(cluster(emb_list, name_list))

    # img3, box_and_point = detect_one("D:\Download\out\students\student.jpg", retinaface_model, retinaface_args)
    # print(embedding(img3,arcface_model).shape)

    # 人脸库中增加一张人脸
    # add_one_to_database(img1,arcface_model,"Aaron_Peirsol","./Database/student.npy")
    # name = findOne(img1, arcface_model, k_v)
    # print(name)

    # 人脸库中批量增加人脸
    create_database_batch(r"D:\Download\out\alig_students_all", arcface_model, "./Database/sfz.npy")

    # 识别视频中的人脸
    # detect_video("software.mp4","out.avi",retinaface_model,arcface_model,k_v,retinaface_args)