import os
import torch

from seg.networks.deeplabv3_plus import get_deeplabv3_plus
import numpy as np
import cv2
import time
from utils import landmark_processor

def label_to_mask(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

label_map = [
    [0, 0, 0],  #
    [128, 128, 128],
    [255, 255, 255],
]

class Evaluator(object):
    def __init__(self, gpu_id, output_img_size, nclass, seg_model_path=None):
        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
        # print("gpu_id: ", gpu_id)

        # create network
        self.model = get_deeplabv3_plus(backbone='xception', nclass=nclass)
        model_path = os.path.join(seg_model_path)
        self.model.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage))
        # print("seg device: ", self.model.device)
        self.model.to(self.device)
        self.model.eval()

        # images = torch.randn((1, 3, 512, 512)).to(self.device)
        # torch.onnx.export(self.model, images,
        #                   "deeplabv3_hair512_360_0520_wl.onnx",
        #                   verbose=True,
        #                   opset_version=11,
        #                   input_names=['data'],
        #                   do_constant_folding=True,
        #                   output_names=['output'])

        # exit()

        self.output_img_size = output_img_size
        self.nclass = nclass


    def process_data(self, img):
        img = (img.astype(np.float32) / 255).transpose((2, 0, 1))
        img = torch.from_numpy(img).unsqueeze(0)

        return img

    def eval(self, img, pts1k):
        orig_h, orig_w, _ = img.shape

        M1 = landmark_processor.get_transform_mat_hair(pts1k, self.output_img_size, ratio=0.28, h_ratio=0.3)
        crop_img = cv2.warpAffine(img, M1, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)
        crop_img = self.process_data(crop_img)
        crop_img = crop_img.to(self.device)

        with torch.no_grad():
            # torch.cuda.synchronize()
            outputs = self.model(crop_img)
            pred = torch.argmax(outputs[0], 1)

            pred = pred[0].detach().cpu().numpy()
            predict = pred.astype(np.float32)

            pred_mask = label_to_mask(predict)

            M1_invert = cv2.invertAffineTransform(M1)
            img_pred = cv2.warpAffine(pred_mask, M1_invert, (orig_w, orig_h), flags=cv2.INTER_CUBIC)  #flags=cv2.INTER_NEAREST
            orig_mask = img_pred.copy()

        # show_concat = np.concatenate((img, orig_mask), axis=1)
        # cv2.imshow("show_concat", show_concat)
        # cv2.waitKey()
        return orig_mask