# -*- coding: utf-8 -*-
# @Author  : zhousf
# @Date    : 2023/12/19 
# @Function:
from pathlib import Path
from typing import Optional

import cv2
import yaml
import numpy
from PIL import Image
import fastdeploy as fd
from easydict import EasyDict
from fastdeploy.libs.fastdeploy_main.vision import (
    ClassifyResult,
    DetectionResult,
    SegmentationResult,
    OCRResult
)

from backend import Backend
from zhousflib.image import read
from zhousflib.image import op
from zhousflib.image import pil_util
from zhousflib.util import ocr_vis_util
from zhousflib.font import Font_SimSun
from zhousflib.image import write


class BackendFastDeploy(Backend):
    def __init__(self, *args, **kwargs):
        self.model_dir: Path = kwargs.get("model_dir")
        self.device_id = kwargs.get("device_id")
        self.backend = None
        self.plugin = None
        self.model = None
        self.label_list = []
        super().__init__(*args)
        self.plugins = self.__register_plugin()

    @staticmethod
    def __register_plugin():
        plugin_classification = {
            "fd.vision.classification.PaddleClasModel": fd.vision.classification.PaddleClasModel,
            "fd.vision.classification.PPShiTuV2Recognizer": fd.vision.classification.PPShiTuV2Recognizer,
        }
        plugin_detection = {
            "fd.vision.detection.PaddleDetectionModel": fd.vision.detection.PaddleDetectionModel,
            "fd.vision.detection.PPYOLOE": fd.vision.detection.PPYOLOE,
            "fd.vision.detection.PPYOLO": fd.vision.detection.PPYOLO,
            "fd.vision.detection.PPYOLOER": fd.vision.detection.PPYOLOER,
            "fd.vision.detection.PaddleYOLOX": fd.vision.detection.PaddleYOLOX,
            "fd.vision.detection.YOLOX": fd.vision.detection.YOLOX,
            "fd.vision.detection.YOLOv8": fd.vision.detection.YOLOv8,
            "fd.vision.detection.YOLOv7": fd.vision.detection.YOLOv7,
            "fd.vision.detection.YOLOv6": fd.vision.detection.YOLOv6,
        }
        plugin_segmentation = {
            "fd.vision.segmentation.PaddleSegModel": fd.vision.segmentation.PaddleSegModel,
        }
        plugin_ocr = {
            "fd.vision.ocr.DBDetector": fd.vision.ocr.DBDetector,
            "fd.vision.ocr.Classifier": fd.vision.ocr.Classifier,
            "fd.vision.ocr.Recognizer": fd.vision.ocr.Recognizer,
            "fd.vision.ocr.PPOCRv4": fd.vision.ocr.PPOCRv4,
            "fd.vision.ocr.PPOCRv3": fd.vision.ocr.PPOCRv3,
        }
        return plugin_classification | plugin_detection | plugin_segmentation | plugin_ocr

    def build(self, **kwargs):
        self.plugin = self.plugins.get(kwargs.get("plugin"), None)
        if self.plugin is None:
            raise Exception("Plugin not found: {0}".format(self.plugin))
        runtime_option = kwargs.get("runtime_option", fd.RuntimeOption())
        if self.device_id >= 0:
            runtime_option.use_gpu(device_id=self.device_id)
        else:
            runtime_option.use_cpu()
        fd_backend = kwargs.get("fd_backend", None)
        if fd_backend is None:
            runtime_option.use_paddle_backend()
        model_format = kwargs.get("model_format")
        if self.model_dir is not None:
            union = [
                self.get_file_path_by_suffix(model_dir=self.model_dir, suffix=".pdmodel"),
                self.get_file_path_by_suffix(model_dir=self.model_dir, suffix=".pdiparams"),
                self.get_file_path_by_suffix(model_dir=self.model_dir, suffix=".yaml"),
                self.get_file_path_by_suffix(model_dir=self.model_dir, suffix=".txt"),
                runtime_option,
                model_format
            ]
            union = [sublist for sublist in union if sublist]
            try:
                self.model = self.plugin(*union)
            except Exception:
                raise Exception("Please check model_dir files: {0}".format(self.model_dir))
        else:
            if self.plugin in [fd.vision.ocr.PPOCRv4, fd.vision.ocr.PPOCRv3]:
                det_model = kwargs.get("det_model", None)
                cls_model = kwargs.get("cls_model", None)
                rec_model = kwargs.get("rec_model", None)
                self.model = self.plugin(det_model, cls_model, rec_model)
        if self.model_dir is not None:
            # init label list
            config_file = self.get_file_path_by_suffix(model_dir=self.model_dir, suffix=".yaml")
            if config_file is not None and isinstance(config_file, str):
                with Path(config_file).open("r") as f:
                    yml_conf = yaml.safe_load(f)
                    if "label_list" in yml_conf:
                        self.label_list = yml_conf["label_list"]
            if len(self.label_list) == 0:
                label_file = self.get_file_path_by_suffix(model_dir=self.model_dir, suffix=".label")
                if label_file is not None and isinstance(label_file, str):
                    with Path(label_file).open("r", encoding="utf-8") as f:
                        for label in f.readlines():
                            label = label.replace("\n", "")
                            item = label.split(" ")
                            if len(item) == 2:
                                self.label_list.append(item[1])

    @staticmethod
    def draw_boxes(infer_result, **kwargs):
        if not hasattr(infer_result, "boxes"):
            return
        boxes = infer_result.boxes
        if len(boxes) == 0:
            return
        arr = numpy.asarray(boxes)
        if len(arr.shape) == 2 and arr.shape[1] == 8:
            arr = arr.reshape([arr.shape[0], 4, 2])
            boxes = arr.tolist()
        # config of visualization
        vis_font = kwargs.get("vis_font", None)
        vis_show = kwargs.get("vis_show", False)
        vis_image_size = kwargs.get("vis_image_size", None)
        vis_fill_transparent = kwargs.get("vis_fill_transparent", 60)
        # config of save visualization file
        vis_image_file: Optional[Path] = kwargs.get("vis_image_file", None)

        image_file = kwargs.get("image_path", None)
        if vis_image_file is not None or vis_show:
            draw_img = pil_util.draw_polygon(polygon=boxes,
                                             image_file=image_file,
                                             image_size=vis_image_size,
                                             font=vis_font,
                                             fill_transparent=vis_fill_transparent,
                                             show=vis_show)
            if vis_image_file is not None:
                draw_img.save(vis_image_file)

    def op_classify(self, infer_result, image_arr, **kwargs):
        result = EasyDict()
        label_names = []
        label_ids = infer_result.label_ids
        scores = infer_result.scores
        for i, score in enumerate(scores):
            label_name = self.label_list[label_ids[i]]
            label_names.append(label_name)
        result.feature = infer_result.feature
        result.label_ids = infer_result.label_ids
        result.label_names = label_names
        result.scores = infer_result.scores
        return result

    def op_detection(self, infer_result, image_arr, **kwargs):
        result = EasyDict()
        label_ids = infer_result.label_ids
        scores = infer_result.scores
        bboxes = []
        items = []
        label_names = []
        for i, box in enumerate(infer_result.boxes):
            bboxes.append([int(box[0]), int(box[1]), int(box[2]), int(box[3])])
            class_name = self.label_list[label_ids[i]]
            items.append([class_name, scores[i], int(box[0]), int(box[1]), int(box[2]), int(box[3])])
            label_names.append(class_name)
        result.items = items
        result.label_ids = infer_result.label_ids
        result.label_names = label_names
        result.boxes = bboxes
        result.scores = infer_result.scores
        result.masks = infer_result.masks
        result.contain_masks = infer_result.contain_masks
        self.draw_boxes(infer_result, **kwargs)
        return result

    @staticmethod
    def op_segmentation(infer_result, image_arr, **kwargs):
        result = EasyDict()
        result.label_map = infer_result.label_map
        result.contain_score_map = infer_result.contain_score_map
        result.score_map = infer_result.score_map
        result.shape = infer_result.shape
        if len(result.label_map) > 0:
            # 是否开启掩码最大连通域计算
            max_connectivity_domain = kwargs.get("max_connectivity_domain", False)
            # 掩码空洞填充阈值，面积占整图比例，0则不填充，小于该阈值时则填充
            mask_hole_area_thresh = kwargs.get("mask_hole_area_thresh", 0)
            arr_mask = numpy.array(result.label_map).reshape(result.shape[0], result.shape[1])
            """
            最大连通域计算
            """
            if max_connectivity_domain:
                max_region = op.max_connectivity_domain(mask_arr=arr_mask)
                result.label_map = max_region.flatten().tolist()
            """
            去掉连通域中的空洞
            """
            if mask_hole_area_thresh > 0:
                image_area = image_arr.shape[0] * image_arr.shape[1]
                arr_mask = numpy.array(result.label_map).reshape(result.shape[0], result.shape[1])
                img_max = numpy.where(arr_mask == 1, 0, 1) * 255
                # 掩码类型转换
                img_max = img_max.astype(dtype=numpy.uint8)
                contours, _ = cv2.findContours(img_max, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
                cv_contours = []
                for contour in contours:
                    area = cv2.contourArea(contour)
                    if area / image_area <= mask_hole_area_thresh:
                        cv_contours.append(contour)
                    else:
                        continue
                if len(cv_contours):
                    cv2.fillPoly(img_max, cv_contours, (0, 0, 0))
                    img_max = numpy.where(img_max == 0, 1, 0)
                    img_max = img_max.astype(dtype=numpy.int32)
                    result.label_map = img_max.flatten().tolist()
            vis_show = kwargs.get("vis_show", False)
            vis_transparent_weight = kwargs.get("vis_transparent_weight", 0.6)
            vis_image_file: Optional[Path] = kwargs.get("vis_image_file", None)
            if vis_image_file is not None or vis_show:
                im_vis = fd.vision.visualize.vis_segmentation(image_arr, infer_result, vis_transparent_weight)
                pil_image = Image.fromarray(im_vis)
                if vis_image_file is not None:
                    pil_image.save(vis_image_file, quality=100)
                if vis_show:
                    pil_image.show()
        return result

    def op_ocr(self, infer_result, image_arr, **kwargs):
        result = EasyDict()
        result.boxes = infer_result.boxes
        result.cls_labels = infer_result.cls_labels
        result.cls_scores = infer_result.cls_scores
        result.rec_scores = infer_result.rec_scores
        result.text = infer_result.text
        if len(result.text) == 0 and len(result.boxes) > 0:
            self.draw_boxes(infer_result, **kwargs)
        if len(result.text) > 0 and len(result.boxes) > 0:
            arr = numpy.asarray(result.boxes)
            if len(arr.shape) == 2 and arr.shape[1] == 8:
                arr = arr.reshape([arr.shape[0], 4, 2])
                result.boxes = arr.tolist()
            vis_image_file: Optional[Path] = kwargs.get("vis_image_file", None)
            vis_show = kwargs.get("vis_show", False)
            if vis_image_file is not None or vis_show:
                boxes = []
                texts = []
                scores = []
                for i, bbox in enumerate(result.boxes):
                    bbox_arr = numpy.asarray(bbox)
                    boxes.append(bbox_arr)
                    texts.append(result.text[i])
                    scores.append(result.rec_scores[i])
                image_file: Optional[Path] = kwargs.get("image_path", None)
                image = Image.open(image_file)
                draw_img = ocr_vis_util.draw_ocr_box_txt(
                    image,
                    boxes,
                    texts,
                    scores,
                    drop_score=0.5,
                    font_path=str(Font_SimSun))
                if vis_image_file is not None:
                    write(image=draw_img[:, :, ::-1], img_write_path=vis_image_file)
                if vis_show:
                    draw_img = Image.fromarray(draw_img)
                    draw_img.show()
        return result

    def post_process(self, infer_result, image_arr, **kwargs):
        operations = {
            ClassifyResult: lambda: self.op_classify(infer_result, image_arr, **kwargs),
            DetectionResult: lambda: self.op_detection(infer_result, image_arr, **kwargs),
            SegmentationResult: lambda: self.op_segmentation(infer_result, image_arr, **kwargs),
            OCRResult: lambda: self.op_ocr(infer_result, image_arr, **kwargs),
        }
        return operations.get(type(infer_result), lambda: infer_result)()

    def inference(self, input_data, **kwargs):
        if isinstance(input_data, Path):
            kwargs["image_path"] = input_data
        image_arr = read(input_data)
        infer_res = self.model.predict(image_arr)
        res = self.post_process(infer_res, image_arr, **kwargs)
        return res
