import os
import time
import torch
import torch.nn.parallel

import numpy as np
from utils import landmark_processor
import cv2

# modelRoot = "/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/weights"
modelRoot = "./weights"

label_map = [
    [0, 0, 0],
    [255, 0, 0],
    [0, 0, 255],
    [0, 255, 0],
    [0, 255, 255],
    # [0, 255, 0]
]
class Generator_BaldSeg_5c(object):
    def __init__(self, gpu_flag, gpu_id):

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

        # load seg model
        self.model_dir = modelRoot
        self.pre_trained_model = os.path.join(self.model_dir, "ori_hair_checkpoint_7660_0611.pt")
        self.net = torch.jit.load(self.pre_trained_model, map_location=self.device).to(self.device)
        self.net.eval()

        self.output_img_size = 512
        self.img_ratio = 0.4

    def label_to_mask(self, label_np):
        label_np = label_np.astype(np.int32)[:, :, np.newaxis]
        mask = np.zeros((label_np.shape[0], label_np.shape[1], 3), dtype=np.uint8)
        for id, color in enumerate(label_map):
            index = (label_np == id).all(axis=2)
            mask[index] = color
        return mask
    def forward(self, image, alpha, landmarks1k):
        image_to_face_mat = landmark_processor.get_transform_mat_full_face_ratio_deeplab(landmarks1k, self.output_img_size, self.img_ratio)
        img_ = (image * (1 - alpha[:, :, np.newaxis].astype(np.float32) / 255)).astype(np.uint8)
        img_cuted = cv2.warpAffine(img_, image_to_face_mat, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_CONSTANT, borderValue=[255, 255, 255])
        inv_image_to_face_mat = cv2.invertAffineTransform(image_to_face_mat)

        img = (img_cuted.astype(np.float32) / 255).transpose((2, 0, 1))
        img = np.expand_dims(img, axis=0)
        img = torch.from_numpy(img).to(self.device)  #cuda(self.gpu_id)

        with torch.no_grad():
            output = self.net(img)

        pred = output.detach().cpu().numpy().squeeze().astype(np.float32) #torch.max(output[:1], 1)[1].detach().cpu().numpy().squeeze().astype(np.float32)
        mask = self.label_to_mask(pred)

        # cv2.imshow("img_cuted: ", img_cuted)
        # cv2.imshow("mask: ", mask)
        # cv2.waitKey()

        black_img = np.zeros(image.shape).astype(np.uint8)
        cv2.warpAffine(mask, inv_image_to_face_mat, (image.shape[1], image.shape[0]),
                       dst=black_img, flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_TRANSPARENT)

        return black_img