import torch_mlu

import torch_mlu.core.mlu_model as ct

import torch_mlu.core.mlu_quantize as mlu_quantize
from maskrcnn_benchmark.config import cfg
#from utils.predictor import COCODemo
#from utils.datasets import *
#from utils.parse_config import *
from maskrcnn_benchmark.utils.checkpoint import DetectronCheckpointer
from maskrcnn_benchmark.modeling.roi_heads.mask_head.inference import Masker
from maskrcnn_benchmark.modeling.detector import build_detection_model
from maskrcnn_benchmark.structures.image_list import to_image_list
from maskrcnn_benchmark.structures.bounding_box import BoxList
from maskrcnn_benchmark.structures.segmentation_mask import SegmentationMask
from maskrcnn_benchmark.utils import cv2_util
import torch_mlu.core.mlu_quantize as mlu_quantize
import torch_mlu.core.mlu_model as ct
from torchvision.models.quantization.utils import quantize_model
from PIL import Image
import cv2
import torch
import numpy as np
from skimage.transform import resize
CATEGORIES = [
   "__background",
   "draft",
]
make_mask = Masker(threshold = 0.3, padding = 1)

def map_box_to_ori_img(predictions, ori_img, mode = 'xyxy'):
    img_h, img_w, _ = ori_img.shape
    assert isinstance(predictions, BoxList), "Detections should be packed into BoxList!"
    bbox = predictions.bbox.clone()
    scaling_factors = min(float(800)/ img_w, float(800)/img_h)

    bbox[:, 0] = ((bbox[:, 0] - (800 - scaling_factors * img_w) / 2.0) / scaling_factors)
    bbox[:, 1] = ((bbox[:, 1] - (800 - scaling_factors * img_h) / 2.0) / scaling_factors)
    bbox[:, 2] = ((bbox[:, 2] - (800 - scaling_factors * img_w) / 2.0) / scaling_factors)
    bbox[:, 3] = ((bbox[:, 3] - (800 - scaling_factors * img_h) / 2.0) / scaling_factors)

    if mode == 'xyxy':
       Box = BoxList(bbox, (img_w, img_h), mode="xyxy")
    else:
       Box = BoxList(bbox, (img_w, img_h), mode="xywh")
    Box.add_field("labels", predictions.get_field("labels"))
    Box.add_field("scores", predictions.get_field("scores"))
    Box.add_field("mask", predictions.get_field("mask"))

    return Box

def select_top_predictions(prediction):
    """
    Select only predictions which have a `score` > confidence_threshold,
    and returns the predictions in descending order of score

    Arguments:
        predictions (BoxList): the result of the computation by the model.
            It should contain the field `scores`.

    Returns:
        prediction (BoxList): the detected objects. Additional information
            of the detection properties can be found in the fields of
            the BoxList via `prediction.fields()`
    """
    scores = prediction.get_field("scores")
    keep = torch.nonzero(scores > 0.3).squeeze(1)
    prediction = prediction[keep]
    scores = prediction.get_field("scores")
    _, idx = scores.sort(0, descending=True)
    return prediction[idx]

def overlay_boxes(image, predictions):
   """
   Adds the predicted boxes on top of the image

   Arguments:
       image (np.ndarray): an image as returned by OpenCV
       predictions (BoxList): the result of the computation by the model.
           It should contain the field `labels`.
   """
   labels = predictions.get_field("labels")
   boxes = predictions.bbox
   colors = compute_colors_for_labels(labels).tolist()
   for box, color in zip(boxes, colors):
       box = box.to(torch.int64)
       top_left, bottom_right = box[:2].tolist(), box[2:].tolist()
       image = cv2.rectangle(
           image, tuple(top_left), tuple(bottom_right), tuple(color), 1
       )
   return image

def overlay_mask(image, predictions):
    """
    Adds the instances contours for each predicted object.
    Each label has a different color.

    Arguments:
        image (np.ndarray): an image as returned by OpenCV
        predictions (BoxList): the result of the computation by the model.
            It should contain the field `mask` and `labels`.
    """
    masks = predictions.get_field("mask").numpy()
    labels = predictions.get_field("labels")

    colors = compute_colors_for_labels(labels).tolist()
    for mask, color in zip(masks, colors):
        thresh = mask[0, :, :, None].astype(np.uint8)
        contours, hierarchy = cv2_util.findContours(
            thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
        )
        image = cv2.drawContours(image, contours, -1, color, 3)

    composite = image

    return composite

def overlay_class_names(image, predictions):
    """
    Adds detected class names and scores in the positions defined by the
    top-left corner of the predicted bounding box

    Arguments:
        image (np.ndarray): an image as returned by OpenCV
        predictions (BoxList): the result of the computation by the model.
            It should contain the field `scores` and `labels`.
    """
    scores = predictions.get_field("scores").tolist()
    labels = predictions.get_field("labels").tolist()
    labels = [CATEGORIES[i] for i in labels]
    boxes = predictions.bbox

    template = "{}: {:.2f}"
    for box, score, label in zip(boxes, scores, labels):
        x, y = box[:2]
        s = template.format(label, score)
        cv2.putText(
            image, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, .5, (255, 255, 255), 1
        )

    return image

def compute_colors_for_labels(labels):
    """
    Simple function that adds fixed colors depending on the class
    """
    palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
    colors = labels[:, None] * palette
    colors = (colors % 255).numpy().astype("uint8")
    return colors

torch.set_grad_enabled(False)
ct.set_core_version("MLU220")
ct.set_device(-1)
ct.save_as_cambricon("mask_rcnn_8.10")
ct.set_input_format(0)

#若量化的时候加入了mean和std,则数据推理的时候normalize则不需要利用mean和std归一化

config_file = './weight/config.yml'
#model_path = './weight/model_final.pth'

cfg.merge_from_file(config_file)
# cfg.merge_from_list(["MODEL.DEVICE", "mlu"])
cfg.merge_from_list(["MODEL.DEVICE", "cpu"])
#cfg.MODEL.WEIGHT = model_path

cfg.freeze()
config = cfg.clone()
save_dir = config.OUTPUT_DIR
model = build_detection_model(config)
print(model)

torch.set_grad_enabled(False)

# # model = quantize_model(model)
model = mlu_quantize.quantize_dynamic_mlu(model,dtype='int16')
model.load_state_dict(torch.load('./quantize/mask_rcnn.pth', map_location='cpu'), strict=False)
model = model.eval().float()
#
# # parpare input
# example_mlu = torch.randn(1, 3, 800, 800, dtype=torch.float)
randn_mlu = torch.randn(1, 3, 800, 800)#, dtype=torch.float)
#
#net_traced = torch.jit.trace(model.to(ct.mlu_device()),
#                              randn_mlu.to(ct.mlu_device()),
#                              check_trace=False)
#
# # run inference and save cambricon
#net_traced(example_mlu.to(ct.mlu_device()))
#model.float().eval().mlu().set_core_number(16)
quantized_net_mlu = model.to(ct.mlu_device())

quantized_net_mlu(randn_mlu.to(ct.mlu_device()))
#
# # unset offline flag
# ct.save_as_cambricon("")
print("Offline model has been generated!")
#

#生成量化模型
model.eval().float()
device = torch.device(cfg.MODEL.DEVICE)
model = model.to(device)
checkpointer = DetectronCheckpointer(cfg, model, save_dir=save_dir)
_ = checkpointer.load(config.MODEL.WEIGHT)
# mean = [102.9801, 115.9465,122.7717]
# std = [1.0, 1.0, 1.0]
mean = [0.4038, 0.455, 0.481]
std = [1./255., 1./255., 1./255.]
# model = mlu_quantize.quantize_dynamic_mlu(model,dtype='int8')
model_quantized = mlu_quantize.quantize_dynamic_mlu(model,
              {'iteration':100, 'use_ave': False, 'data_scale':1.0, \
              'mean': mean, 'std': std,'per_channel': True}, \
              dtype='int16', gen_quant=True)

# image = Image.open("1981.jpg") #图片发在了build文件夹下
# image = image.resize((800, 800),Image.ANTIALIAS)
# image = np.asarray(image)
# image = torch.Tensor(image).unsqueeze_(dim=0)
# image = image.permute((0, 3, 1, 2)).float()
# print('image:',image.size())
# predict = model(image)
# print("predict:",predict)

img_shape = (800, 800)
img = np.array(cv2.imread('1981.jpg'))
h,w, _ = img.shape
dim_diff = np.abs(h-w)
pad1, pad2 = dim_diff // 2, dim_diff-dim_diff//2
pad = ((pad1, pad2), (0, 0), (0, 0)) if h <= w else ((0, 0), (pad1, pad2), (0, 0))

# input_img = np.pad(img, pad, 'constant', constant_values=128)
# input_img = cv2.resize(input_img, \
#                      (img_shape[0], img_shape[1]), \
#                       interpolation = cv2.INTER_AREA)
# input_img = np.transpose(input_img, (2, 0, 1))
# # As pytorch tensor
# input_img = torch.from_numpy(input_img).float().unsqueeze_(dim=0)

# #/正确的
# input_img = np.pad(img, pad, 'constant', constant_values=128) / 255.
# input_img = resize(input_img, (img_shape[0], img_shape[1], 3), mode='reflect')
# #\

#测试/
input_img = np.pad(img, pad, 'constant', constant_values=128)
input_img = cv2.resize(input_img,(img_shape[0], img_shape[1]))
input_img = input_img/255.
#\

input_img = np.transpose(input_img, (2, 0, 1))
# As pytorch tensor
input_img = torch.from_numpy(input_img).float()
mean = torch.as_tensor(mean)
std = torch.as_tensor(std)
input_img = (input_img - mean[:, None, None]) / std[:, None, None]
input_img = input_img.unsqueeze_(dim=0)



print('image:',input_img.size())
predict = model_quantized(input_img)
torch.save(model_quantized.state_dict(), "./quantize/mask_rcnn_quan.pth")
print("quan finished")
print("predict:",predict)
bbox__aaa= predict[0].bbox.clone()

quantized_net = torch_mlu.core.mlu_quantize.quantize_dynamic_mlu(model)
quantized_net.load_state_dict(torch.load("./quantize/mask_rcnn_quan.pth",map_location='cpu'),False)
#model.float().eval().mlu().set_core_number(16)
quantized_net_mlu = quantized_net.to(ct.mlu_device())
ct.set_core_number(1)
data_mlu = input_img.to(ct.mlu_device())
torch.set_grad_enabled(False)
output = quantized_net_mlu(data_mlu)
#model_traced = torch.jit.trace(quantized_net_mlu, data_mlu, check_trace=False)  #进行融合

#output = model_traced(data_mlu)

print(output.cpu().shape)

for j, i in enumerate(bbox__aaa):
    print("boxes"+str(j)+":",[float(i[0].numpy()),float(i[1].numpy()),float(i[2].numpy()),float(i[3].numpy())])
#normalize
# input_img = np.pad(img, pad, 'constant', constant_values=128) / 255.
# input_img = resize(input_img, (img_shape[0], img_shape[1], 3), mode='reflect')
# mean = torch.as_tensor([0.4038, 0.455, 0.481])
# std = torch.as_tensor([1./255., 1./255., 1./255.])
# input_img = (input_img - mean[:, None, None]) / std[:, None, None]


# print('image:',image.size())
# predict = model(image)
# print("predict:",predict)

for si, predict1 in enumerate(predict):
    print("si:",si)
    ori_img = np.array(cv2.imread('1981.jpg'))
    predict1 = map_box_to_ori_img(predict1, ori_img, mode='xyxy')
    predict1 = select_top_predictions(predict1)
    if predict1.has_field("mask"):
        masks = predict1.get_field("mask")
    print("mask.mask_befor_shape:",masks.shape)

    masks = make_mask([masks], [predict1])[0]
    predict1.add_field("mask", masks)
    result = ori_img.copy()
    # render detection boxes
    result = overlay_boxes(result, predict1)
    # render detection masks
    result = overlay_mask(result, predict1)
    # render detection labels
    result = overlay_class_names(result, predict1)
    cv2.imwrite("cpu_" + '1981' + ".jpg", result)

for si, output1 in enumerate(output):
    print("si:",si)
    ori_img = np.array(cv2.imread('1981.jpg'))
    output1 = map_box_to_ori_img(output1, ori_img, mode='xyxy')
    output1 = select_top_predictions(output1)
    if output1.has_field("mask"):
        masks = output1.get_field("mask")
    print("mask.mask_befor_shape:",masks.shape)

    masks = make_mask([masks], [output1])[0]
    output1.add_field("mask", masks)
    result = ori_img.copy()
    # render detection boxes
    result = overlay_boxes(result, output1)
    # render detection masks
    result = overlay_mask(result, output1)
    # render detection labels
    result = overlay_class_names(result, output1)
    cv2.imwrite("mlu_" + '1981' + ".jpg", result)