import json
import os
import time

import dlib
import faiss
import numpy as np

from utils.es import es_query
from utils.images import face_to_vec, url_to_ndarray


def save_diary_image_info(save_file, face_to_vec_f):
    q = {
        "query": {
            "bool": {
                "filter": [{
                    "term": {
                        "is_online": True
                    }
                }, {
                    "term": {
                        "has_before_cover": True
                    }
                }, {
                    "term": {
                        "has_after_cover": True
                    }
                }, {
                    "exists": {
                        "field": "before_cover_url"
                    }
                }]
            }
        },
        "_source": {
            "include": ["id", "before_cover_url", "after_cover_url"]
        }
    }

    with open(save_file, "w") as f:
        step = 100
        # for i in range(0, 500000, step):
        res_dict = es_query("diary", q, 0, step)

        for item in res_dict["hits"]["hits"]:
            diary_id = item["_source"]["id"]
            before_cover_url = item["_source"]["before_cover_url"] + "-w"
            after_cover_url = item["_source"]["after_cover_url"] + "-w"
            img = url_to_ndarray(after_cover_url)
            if img.any():
                faces = face_to_vec_f(img)
                for face in faces:
                    line = str(diary_id) + "\t" + face["feature"] + "\n"
                    # print(line)
                    f.write(line)


def save_faiss_index(load_file, save_path):
    with open(load_file, "r") as f:
        ids = []
        features = []
        for line in f.readlines():
            tmp = line.split("\t")
            ids.append(tmp[0])
            features.append(np.array(json.loads(tmp[1])))

        ids_np = np.array(ids).astype("float32")
        features_np = np.array(features).astype("int")
        print(ids_np)
        print(ids.dtype)
        print(features_np)
        print(features_np.dtype)
        index = faiss.IndexHNSW(128, 32)
        print(index.is_trained)
        index2 = faiss.IndexIDMap(index)
        index2.add_with_ids(features_np, ids_np)
        faiss.write_index(index2, save_path)


def main():
    base_dir = os.getcwd()
    print("base_dir: " + base_dir)
    model_diry = os.path.join(base_dir, "_models")
    facerec_model_path = os.path.join(model_diry, "dlib_face_recognition_resnet_model_v1.dat")
    shape_model_path = os.path.join(model_diry, "shape_predictor_68_face_landmarks.dat")
    faiss_index_path = os.path.join(base_dir, "_index", "current.index")

    diary_after_cover_vec_file = "./diary_after_cover_vec.txt"

    face_rec = dlib.face_recognition_model_v1(facerec_model_path)
    face_detector = dlib.get_frontal_face_detector()
    shape_predictor = dlib.shape_predictor(shape_model_path)
    face_to_vec_f = lambda img: face_to_vec(img, face_rec, face_detector, shape_predictor)

    # save_diary_image_info(diary_after_cover_vec_file, face_to_vec_f)
    save_faiss_index(diary_after_cover_vec_file, faiss_index_path)

    # img_url = "https://pic.igengmei.com/2020/07/03/1437/1b9975bb0b81-w"
    # img = url_to_ndarray(img_url)
    # if img.any():
    #     faces = face_to_vec_f(img)

    #     if len(faces) > 0:
    #         for face in faces:
    #             print(face)
    #             print("\n")
    #             print("-------------")
    #             face_feature = np.array(json.loads(face["feature"]))

    #         ids = [17418645]
    #         ids_np = np.array(ids).astype("int")
    #         faces_np = np.array([face_feature]).astype("float32")
    #         print(ids_np)
    #         print(faces_np)

    #         index = faiss.IndexHNSWFlat(128, 32)
    #         print("index: ")
    #         print(index)
    #         index2 = faiss.IndexIDMap(index)
    #         print("index2: ")
    #         print(index2)
    #         index2.add_with_ids(faces_np, ids_np)
    #         print("index2: ")
    #         print(index2)
    #     else:
    #         print("no faces")


if __name__ == "__main__":
    begin_time = time.time()

    main()

    print("total cost: {:.2f}mins".format((time.time() - begin_time) / 60))
