import cv2
import numpy as np
import torch
import json
import time
import re
import os
from process_modules import Get_Landmark, Process_Data
from utils import landmark_processor

class GenderClassifyProcessor(object):
    def __init__(self, gpu_id=0):
        model_path = "./weights/gender_models"
        self.output_img_size = 128
        if not os.path.exists(model_path):
            print("GenderClassifyProcessor don't have model!")

        if gpu_id == 'cpu':
            self.device = torch.device(gpu_id)
        else:
            self.device = torch.device('cuda:{0}'.format(gpu_id))

        self.gender_model = cv2.dnn.readNetFromCaffe(os.path.join(model_path, "gender.prototxt"), os.path.join(model_path, "gender.caffemodel"))

    def forward(self, img, landmark_137):

        image_to_face_mat = landmark_processor.get_transform_mat_sex(landmark_137, self.output_img_size)
        gender_img = cv2.warpAffine(img, image_to_face_mat, (self.output_img_size, self.output_img_size), cv2.INTER_LANCZOS4)
        inpBlob = cv2.dnn.blobFromImage(gender_img, 1.0, (self.output_img_size, self.output_img_size), (0, 0, 0), swapRB=False,
                                        crop=False)
        self.gender_model.setInput(inpBlob)
        output = self.gender_model.forward()

        is_female = True
        if output[0][0] > output[0][1]:
            is_female = False
        return is_female


class Prepare_Ref_HairStyle_Data(object):
    def __init__(self, gpu, device_id):
        if gpu and torch.cuda.is_available():
            self.device = torch.device("cuda:%d" % device_id if device_id >= 0 else "cpu")
            self.cuda = True
        else:
            self.device = torch.device("cpu")
            self.cuda = False

        self.process_data = Process_Data(gpu, device_id)
        self.get_landmark = Get_Landmark(gpu_id=device_id)
        self.gender_classify = GenderClassifyProcessor(gpu_id=device_id)

    def get_prepare_ref_768_color_data(self, ref_rgb_8uc3_orisize):
        ref_landmark_1k2_f_orisize = self.get_landmark.forward(ref_rgb_8uc3_orisize)
        ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768 = \
            self.process_data.get_prepare_ref_768_bald_data(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)

        return ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768

    def get_prepare_ref_768_color_data_landmark1k(self, ref_rgb_8uc3_orisize,ref_landmark_1k2_f_orisize):
        # ref_landmark_1k2_f_orisize = self.get_landmark.inference(ref_rgb_8uc3_orisize)
        ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768 = \
            self.process_data.get_prepare_ref_768_bald_data(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)

        return ref_rgb_8uc3_color_768, ref_matting_8uc3_color_768, ref_baldseg_8uc3_color_768, ref_landmark_f1k2_color_768

    def calculate_hair_ratio_after_align(self, hair_mask, origin_landmark1k, img_size=768):
        image_to_face_mat = landmark_processor.get_transform_mat_hair_ratio_v1(origin_landmark1k, 768, ratio=0.35, h_offset=0.32)
        hair_mask_align = cv2.warpAffine(hair_mask, image_to_face_mat, (img_size, img_size))

        hair_rect = cv2.boundingRect(hair_mask_align[:, :, :1])
        hair_mask_ratio = hair_rect[2] * hair_rect[3] / (img_size * img_size)
        return hair_mask_ratio

    def check_female_hair_ratio(self, origin_img_8uc3, landmark_1k2_f_orisize):
        ref_matte_fg_8uc3_orisize, ref_matte_pred_8uc1_orisize = self.process_data.generator_matte.matte_inference(origin_img_8uc3, landmark_1k2_f_orisize)
        ref_matte_pred_8uc3_orisize = np.repeat(ref_matte_pred_8uc1_orisize[:, :, np.newaxis], 3, axis=2)
        hairstyle_M = self.process_data.get_hair_M_girl_v1(landmark_1k2_f_orisize)
        ref_rgb_8uc3_768 = cv2.warpAffine(ref_matte_pred_8uc3_orisize, hairstyle_M, (768, 768))

        # cv2.imshow("ref_rgb_8uc3_768", ref_rgb_8uc3_768)
        # cv2.waitKey()

        edge_width = 5
        if (ref_rgb_8uc3_768[-edge_width:, :, :]).max() > 0 or (ref_rgb_8uc3_768[:, -edge_width:, :]).max() > 0 or (ref_rgb_8uc3_768[:, :edge_width, :]).max() > 0:
            return False
        else:
            return True

    def get_prepare_ref_768_data(self, ref_rgb_8uc3_orisize, long_flag=False):

        ref_landmark_1k2_f_orisize = self.get_landmark.forward(ref_rgb_8uc3_orisize)

        # ref_matte_fg_8uc3_orisize, ref_matte_pred_8uc1_orisize = self.process_data.generator_matte.matte_inference(
        #     ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
        ref_landmark_137kpts_f_orisize = landmark_processor.pts_1k_to_137(ref_landmark_1k2_f_orisize)
        gender_res = self.gender_classify.forward(ref_rgb_8uc3_orisize, ref_landmark_137kpts_f_orisize)

        # hair_ratio = self.calculate_hair_ratio_after_align(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize)
        # if hair_ratio > 0.3:
        #     ratio = 2
        # else:
        #     if gender_res:
        #         ratio = 1
        #     else:
        #         ratio = 0
        if not gender_res:
            ratio = 0
        else:
            check_res = self.check_female_hair_ratio(ref_rgb_8uc3_orisize, ref_landmark_137kpts_f_orisize)
            if check_res:
                ratio = 1
            else:
                ratio = 2
        if gender_res:
            gender = "girl"
        else:
            gender = "boy"
        print("gender: ", gender, " ratio: ", ratio)

        # show_concat = np.concatenate((ref_rgb_8uc3_orisize, ref_matte_fg_8uc3_orisize), axis=1)
        # resize_ratio = 1024. / max(show_concat.shape[:2])
        # show_concat = cv2.resize(show_concat, (0, 0), fx=resize_ratio, fy=resize_ratio)
        # print("gender_res: ", gender_res, " hair_ratio: ", hair_ratio)
        # cv2.imshow("show_concat", show_concat)
        # cv2.imshow("ref_matte_pred_8uc1_orisize", ref_matte_pred_8uc1_orisize)
        # cv2.waitKey()

        ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768 = \
            self.process_data.get_prepare_ref_768_data(ref_rgb_8uc3_orisize, ref_landmark_1k2_f_orisize, ratio, long=long_flag)

        # cv2.imshow("ref_rgb_8uc3_768", ref_rgb_8uc3_768)
        # cv2.waitKey()

        return ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768, gender, ratio

    def Generator_reftensor(self, ref_rgb_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768,
                                 ref_landmark_f1k2_768):

        """
        input：

        图像尺寸基于： 人脸 512， 图像均为3通道
        ref_rgb_8uc3_512: 参考图， size 512， uint8 （0-255）
        ref_matting_8uc3_512：参考图 matting alpha 值， uint8 （0-255）
        ref_baldseg_8uc3_512: 参考图 光头分割， uint8 （0-255）
        ref_landmark_f1k2_512： 参考图 关键点 1k*2 float32

        output：

        input_another_pose_hair_image: 参考图 条件图， float32 （0-255）

        """

        # 8 ********************** another pose hair_image **********************
        another_pose_image = ref_rgb_8uc3_768.copy()

        another_nohair_pose_mask = ref_baldseg_8uc3_768.copy()

        # cv2.imshow("another_nohair_pose_mask", another_nohair_pose_mask)
        # cv2.waitKey()

        another_nohair_pose_mask[(np.abs(another_nohair_pose_mask - [0, 0, 255]) < 50).all(axis=2)] = [0, 255, 255]
        another_nohair_pose_mask[(another_nohair_pose_mask == [0, 0, 0]).all(axis=2)] = [255, 255, 255]
        another_pose_pts137 = landmark_processor.pts_1k_to_137(ref_landmark_f1k2_768).astype(np.int32)
        cv2.fillPoly(another_nohair_pose_mask,
                     np.concatenate((another_pose_pts137[47:35:-1], another_pose_pts137[56:64],
                                     [another_pose_pts137[48], another_pose_pts137[22]]))[
                     np.newaxis, :, :], (255, 255, 0))

        cv2.fillPoly(another_nohair_pose_mask,
                     np.concatenate((another_pose_pts137[22:37], another_pose_pts137[56:47:-1]))[np.newaxis, :, :],
                     (255, 0, 128))

        cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[88:104][np.newaxis, :, :], (255, 0, 255))  # eye
        cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[105:121][np.newaxis, :, :], (255, 0, 255))  # eye
        # Label nose
        cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[64:79][np.newaxis, :, :], (255, 255, 255))
        # Label eyebrow
        cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[121:129][np.newaxis, :, :], (255, 128, 0))
        cv2.fillPoly(another_nohair_pose_mask, another_pose_pts137[129:137][np.newaxis, :, :], (255, 128, 0))

        another_pose_hair_alpha = ref_matting_8uc3_768.copy() / 255.
        another_pose_hair_image = (another_pose_image * another_pose_hair_alpha + another_nohair_pose_mask * (
                    1 - another_pose_hair_alpha)).astype(np.uint8)

        input_another_pose_hair_image = another_pose_hair_image.astype(np.float32) / 255

        return input_another_pose_hair_image

def prepare_single(img_path=None, dst_dir=None, long_flag=False):
    # img_path = "/home/yangchaojie/Desktop/hairstyle/test_data/1123/pics/test3.jpg"
    # img_path = "./test_data/ref_imgs/female1.jpg"
    # dst_dir = "./data/ref_hairstyle/tmp"

    # img_path = "/home/szlc/Downloads/111/111/ddbc6cb6-766c-4f08-b1fd-22c6bab86432.jpg"
    # dst_dir = "/home/szlc/Downloads/111/res"
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)
    prepare_data_process = Prepare_Ref_HairStyle_Data(True, device_id=0)
    ref_rgb_8uc3_orisize = cv2.imread(img_path)

    # debug = True
    debug = False

    version = "20221206"
    start_time = time.time()

    ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768, gender, ratio = \
        prepare_data_process.get_prepare_ref_768_data(ref_rgb_8uc3_orisize, long_flag)


    if not debug:
        cv2.imwrite(os.path.join(dst_dir, "ref_rgb_8uc3_768.png"), ref_rgb_8uc3_768)
        cv2.imwrite(os.path.join(dst_dir, "ref_matting_fg_8uc3_768.png"), ref_matting_fg_8uc3_768)
        cv2.imwrite(os.path.join(dst_dir, "ref_matting_8uc3_768.png"), ref_matting_8uc3_768)
        cv2.imwrite(os.path.join(dst_dir, "ref_baldseg_8uc3_768.png"), ref_baldseg_8uc3_768)
        np.savetxt(os.path.join(dst_dir, "ref_landmark_f1k2_768.txt"), ref_landmark_f1k2_768)

    input_another_pose_hair_image = prepare_data_process.Generator_reftensor(ref_rgb_8uc3_768, ref_matting_8uc3_768,
                                                                     ref_baldseg_8uc3_768,
                                                                     ref_landmark_f1k2_768)
    print("cost time: ", time.time() - start_time)
    if debug:
        print("gender: ", gender, " ratio: ", ratio)
        show_concat = np.concatenate((ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768), axis=1)
        # cv2.imshow("show_concat", cv2.resize(show_concat, (0, 0), fx=0.6, fy=0.6))
        # cv2.imshow("input_another_pose_hair_image", input_another_pose_hair_image)
        # cv2.waitKey()
    if not debug:
        np.save(os.path.join(dst_dir, "input_another_pose_hair_image.npy"), input_another_pose_hair_image)

        config_dict = {}

        config_dict['gender'] = gender
        config_dict['version'] = version
        config_dict['ratio'] = str(ratio)

        with open(os.path.join(dst_dir, "config.json"), "w") as f:
            json.dump(config_dict, f)
        print("写入文件完成...", dst_dir)


def prepare_single_color(img_path=None, dst_dir=None, long_flag=False):

    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)
    prepare_data_process = Prepare_Ref_HairStyle_Data(True, device_id=0)
    ref_rgb_8uc3_orisize = cv2.imread(img_path)

    # debug = True
    debug = False

    version = "20221206"
    start_time = time.time()

    ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768, gender, ratio = \
        prepare_data_process.get_prepare_ref_768_data(ref_rgb_8uc3_orisize, long_flag)


    if not debug:
        cv2.imwrite(os.path.join(dst_dir, "ref_rgb_8uc3_color_768.png"), ref_rgb_8uc3_768)
        cv2.imwrite(os.path.join(dst_dir, "ref_matting_fg_8uc3_color_768.png"), ref_matting_fg_8uc3_768)
        cv2.imwrite(os.path.join(dst_dir, "ref_matting_8uc3_color_768.png"), ref_matting_8uc3_768)
        cv2.imwrite(os.path.join(dst_dir, "ref_baldseg_8uc3_color_768.png"), ref_baldseg_8uc3_768)
        np.savetxt(os.path.join(dst_dir, "ref_landmark_f1k2_color_768.txt"), ref_landmark_f1k2_768)

    input_another_pose_hair_image = prepare_data_process.Generator_reftensor(ref_rgb_8uc3_768, ref_matting_8uc3_768,
                                                                     ref_baldseg_8uc3_768,
                                                                     ref_landmark_f1k2_768)
    print("cost time: ", time.time() - start_time)
    if debug:
        print("gender: ", gender, " ratio: ", ratio)
        show_concat = np.concatenate((ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768), axis=1)
        # cv2.imshow("show_concat", cv2.resize(show_concat, (0, 0), fx=0.6, fy=0.6))
        # cv2.imshow("input_another_pose_hair_image", input_another_pose_hair_image)
        # cv2.waitKey()
    if not debug:
        np.save(os.path.join(dst_dir, "input_another_pose_hair_image.npy"), input_another_pose_hair_image)

        config_dict = {}

        config_dict['gender'] = gender
        config_dict['version'] = version
        config_dict['ratio'] = str(ratio)

        with open(os.path.join(dst_dir, "config.json"), "w") as f:
            json.dump(config_dict, f)
        print("写入文件完成...", dst_dir)

def prepare_multi():

    img_dir = "/home/yangchaojie/Desktop/hairstyle/test_data/testdata_tj_0309/raw_pics"
    save_dir = "/home/yangchaojie/Desktop/hairstyle/test_data/testdata_tj_0309/ref_haircolor"
    # pattern = re.compile(r'^[^\.].*\.jpg$')
    pattern = re.compile(r'^[^\.].*\.(jpg|JPG)$')
    for dirpath, dirnames, filenames in os.walk(img_dir, followlinks=True):
        for filename in filenames:
            if not pattern.match(filename): continue
            img_path = os.path.join(dirpath, filename)
            dst_dir = os.path.join(save_dir, filename[:-4])

            if not os.path.exists(dst_dir):
                os.makedirs(dst_dir)
            prepare_data_process = Prepare_Ref_HairStyle_Data(True, device_id=0)
            ref_rgb_8uc3_orisize = cv2.imread(img_path)

            # debug = True
            debug = False

            version = "20230309"
            start_time = time.time()

            ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768, ref_landmark_f1k2_768, gender, ratio = \
                prepare_data_process.get_prepare_ref_768_data(ref_rgb_8uc3_orisize)


            if not debug:
                cv2.imwrite(os.path.join(dst_dir, "ref_rgb_8uc3_768.png"), ref_rgb_8uc3_768)
                cv2.imwrite(os.path.join(dst_dir, "ref_matting_fg_8uc3_768.png"), ref_matting_fg_8uc3_768)
                cv2.imwrite(os.path.join(dst_dir, "ref_matting_8uc3_768.png"), ref_matting_8uc3_768)
                cv2.imwrite(os.path.join(dst_dir, "ref_baldseg_8uc3_768.png"), ref_baldseg_8uc3_768)
                np.savetxt(os.path.join(dst_dir, "ref_landmark_f1k2_768.txt"), ref_landmark_f1k2_768)

            input_another_pose_hair_image = prepare_data_process.Generator_reftensor(ref_rgb_8uc3_768, ref_matting_8uc3_768,
                                                                             ref_baldseg_8uc3_768,
                                                                             ref_landmark_f1k2_768)
            print("cost time: ", time.time() - start_time)
            if debug:
                print("gender: ", gender, " ratio: ", ratio)
                show_concat = np.concatenate((ref_rgb_8uc3_768, ref_matting_fg_8uc3_768, ref_matting_8uc3_768, ref_baldseg_8uc3_768), axis=1)
                cv2.imshow("show_concat", cv2.resize(show_concat, (0, 0), fx=0.6, fy=0.6))
                cv2.imshow("input_another_pose_hair_image", input_another_pose_hair_image)
                cv2.waitKey()
            if not debug:
                np.save(os.path.join(dst_dir, "input_another_pose_hair_image.npy"), input_another_pose_hair_image)

                config_dict = {}

                config_dict['gender'] = gender
                config_dict['version'] = version
                config_dict['ratio'] = str(ratio)

                with open(os.path.join(dst_dir, "config.json"), "w") as f:
                    json.dump(config_dict, f)
                print("写入文件完成...", dst_dir)

if __name__ == '__main__':
    prepare_single()
    # prepare_multi()