# -*- coding: utf-8 -*-
# @Time    : 2019/8/24 12:06
# @Author  : zhoujun

import os
import sys
import pathlib

__dir__ = pathlib.Path(os.path.abspath(__file__))
sys.path.append(str(__dir__))
sys.path.append(str(__dir__.parent.parent))

import time
import cv2
import paddle

from data_loader import get_transforms
from models import build_model
from post_processing import get_post_processing


def resize_image(img, short_size):
    height, width, _ = img.shape
    if height < width:
        new_height = short_size
        new_width = new_height / height * width
    else:
        new_width = short_size
        new_height = new_width / width * height
    new_height = int(round(new_height / 32) * 32)
    new_width = int(round(new_width / 32) * 32)
    resized_img = cv2.resize(img, (new_width, new_height))
    return resized_img


class PaddleModel:
    def __init__(self, model_path, post_p_thre=0.7, gpu_id=None):
        """
        初始化模型
        :param model_path: 模型地址(可以是模型的参数或者参数和计算图一起保存的文件)
        :param gpu_id: 在哪一块gpu上运行
        """
        self.gpu_id = gpu_id

        if (
            self.gpu_id is not None
            and isinstance(self.gpu_id, int)
            and paddle.device.is_compiled_with_cuda()
        ):
            paddle.device.set_device("gpu:{}".format(self.gpu_id))
        else:
            paddle.device.set_device("cpu")
        checkpoint = paddle.load(model_path)

        config = checkpoint["config"]
        config["arch"]["backbone"]["pretrained"] = False
        self.model = build_model(config["arch"])
        self.post_process = get_post_processing(config["post_processing"])
        self.post_process.box_thresh = post_p_thre
        self.img_mode = config["dataset"]["train"]["dataset"]["args"]["img_mode"]
        self.model.set_state_dict(checkpoint["state_dict"])
        self.model.eval()

        self.transform = []
        for t in config["dataset"]["train"]["dataset"]["args"]["transforms"]:
            if t["type"] in ["ToTensor", "Normalize"]:
                self.transform.append(t)
        self.transform = get_transforms(self.transform)

    def predict(self, img_path: str, is_output_polygon=False, short_size: int = 1024):
        """
        对传入的图像进行预测，支持图像地址,opecv 读取图片，偏慢
        :param img_path: 图像地址
        :param is_numpy:
        :return:
        """
        assert os.path.exists(img_path), "file is not exists"
        img = cv2.imread(img_path, 1 if self.img_mode != "GRAY" else 0)
        if self.img_mode == "RGB":
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        h, w = img.shape[:2]
        img = resize_image(img, short_size)
        # 将图片由(w,h)变为(1,img_channel,h,w)
        tensor = self.transform(img)
        tensor = tensor.unsqueeze_(0)

        batch = {"shape": [(h, w)]}
        with paddle.no_grad():
            start = time.time()
            preds = self.model(tensor)
            box_list, score_list = self.post_process(
                batch, preds, is_output_polygon=is_output_polygon
            )
            box_list, score_list = box_list[0], score_list[0]
            if len(box_list) > 0:
                if is_output_polygon:
                    idx = [x.sum() > 0 for x in box_list]
                    box_list = [box_list[i] for i, v in enumerate(idx) if v]
                    score_list = [score_list[i] for i, v in enumerate(idx) if v]
                else:
                    idx = (
                        box_list.reshape(box_list.shape[0], -1).sum(axis=1) > 0
                    )  # 去掉全为0的框
                    box_list, score_list = box_list[idx], score_list[idx]
            else:
                box_list, score_list = [], []
            t = time.time() - start
        return preds[0, 0, :, :].detach().cpu().numpy(), box_list, score_list, t


def save_depoly(net, input, save_path):
    input_spec = [paddle.static.InputSpec(shape=[None, 3, None, None], dtype="float32")]
    net = paddle.jit.to_static(net, input_spec=input_spec)

    # save static model for inference directly
    paddle.jit.save(net, save_path)


def init_args():
    import argparse

    parser = argparse.ArgumentParser(description="DBNet.paddle")
    parser.add_argument("--model_path", default=r"model_best.pth", type=str)
    parser.add_argument(
        "--input_folder", default="./test/input", type=str, help="img path for predict"
    )
    parser.add_argument(
        "--output_folder", default="./test/output", type=str, help="img path for output"
    )
    parser.add_argument("--gpu", default=0, type=int, help="gpu for inference")
    parser.add_argument(
        "--thre", default=0.3, type=float, help="the thresh of post_processing"
    )
    parser.add_argument("--polygon", action="store_true", help="output polygon or box")
    parser.add_argument("--show", action="store_true", help="show result")
    parser.add_argument(
        "--save_result", action="store_true", help="save box and score to txt file"
    )
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    import pathlib
    from tqdm import tqdm
    import matplotlib.pyplot as plt
    from utils.util import show_img, draw_bbox, save_result, get_image_file_list

    args = init_args()
    print(args)
    # 初始化网络
    model = PaddleModel(args.model_path, post_p_thre=args.thre, gpu_id=args.gpu)
    img_folder = pathlib.Path(args.input_folder)
    for img_path in tqdm(get_image_file_list(args.input_folder)):
        preds, boxes_list, score_list, t = model.predict(
            img_path, is_output_polygon=args.polygon
        )
        img = draw_bbox(cv2.imread(img_path)[:, :, ::-1], boxes_list)
        if args.show:
            show_img(preds)
            show_img(img, title=os.path.basename(img_path))
            plt.show()
        # 保存结果到路径
        os.makedirs(args.output_folder, exist_ok=True)
        img_path = pathlib.Path(img_path)
        output_path = os.path.join(args.output_folder, img_path.stem + "_result.jpg")
        pred_path = os.path.join(args.output_folder, img_path.stem + "_pred.jpg")
        cv2.imwrite(output_path, img[:, :, ::-1])
        cv2.imwrite(pred_path, preds * 255)
        save_result(
            output_path.replace("_result.jpg", ".txt"),
            boxes_list,
            score_list,
            args.polygon,
        )