# Copyright (c) Alibaba, Inc. and its affiliates.
import math
import os.path as osp
from typing import Any, Dict

import cv2
import numpy as np
import PIL
import torch

from modelscope.metainfo import Pipelines
from modelscope.outputs import OutputKeys
from modelscope.pipelines.base import Input, Pipeline
from modelscope.pipelines.builder import PIPELINES
from modelscope.pipelines.cv.ocr_utils.model_resnet18_half import \
    CardDetectionCorrectionModel
from modelscope.pipelines.cv.ocr_utils.table_process import (
    bbox_decode, bbox_post_process, decode_by_ind, get_affine_transform, nms)
from modelscope.preprocessors import load_image
from modelscope.preprocessors.image import LoadImage
from modelscope.utils.config import Config
from modelscope.utils.constant import ModelFile, Tasks
from modelscope.utils.device import (create_device, device_placement,
                                     verify_device)
from modelscope.utils.logger import get_logger

logger = get_logger()


@PIPELINES.register_module(
    Tasks.card_detection_correction,
    module_name=Pipelines.card_detection_correction)
class CardDetectionCorrection(Pipeline):
    r""" Card Detection Pipeline.

    Examples:

    >>> from modelscope.pipelines import pipeline

    >>> detector = pipeline(Tasks.card_detection_correction, model='damo/cv_resnet18_card_correction')
    >>> detector("https://modelscope.oss-cn-beijing.aliyuncs.com/test/images/card_detection_correction.jpg")
    >>>   {
    >>>    "polygons": array([[ 60.562023, 110.682144, 688.57715, 77.34028, 720.2409,
    >>>                         480.33508,  70.20054, 504.9171 ]], dtype=float32),
    >>>    "output_imgs": [array([
    >>>    [[[168, 176, 192],
    >>>     [165, 173, 188],
    >>>     [163, 172, 187],
    >>>     ...,
    >>>     [153, 153, 165],
    >>>     [153, 153, 165],
    >>>     [153, 153, 165]],
    >>>    [[187, 194, 210],
    >>>     [184, 192, 203],
    >>>     [183, 191, 199],
    >>>     ...,
    >>>     [168, 166, 186],
    >>>     [169, 166, 185],
    >>>     [169, 165, 184]],
    >>>    [[186, 193, 211],
    >>>     [183, 191, 205],
    >>>     [183, 192, 203],
    >>>     ...,
    >>>     [170, 167, 187],
    >>>     [171, 165, 186],
    >>>     [170, 164, 184]]]], dtype=uint8)}
    """

    def __init__(self,
                 model: str,
                 device: str = 'gpu',
                 device_map=None,
                 **kwargs):
        """
        Args:
            model: model id on modelscope hub.
        """
        super().__init__(model=model, **kwargs)
        model_path = osp.join(self.model, ModelFile.TORCH_MODEL_FILE)
        config_path = osp.join(self.model, ModelFile.CONFIGURATION)
        logger.info(f'loading model from {model_path}')

        self.cfg = Config.from_file(config_path)
        self.K = self.cfg.K

        if device_map is not None:
            assert device == 'gpu', '`device` and `device_map` cannot be input at the same time!'
        self.device_map = device_map
        verify_device(device)
        self.device_name = device
        self.device = create_device(self.device_name)

        self.infer_model = CardDetectionCorrectionModel()
        checkpoint = torch.load(
            model_path, map_location=self.device, weights_only=True)
        if 'state_dict' in checkpoint:
            self.infer_model.load_state_dict(checkpoint['state_dict'])
        else:
            self.infer_model.load_state_dict(checkpoint)
        self.infer_model = self.infer_model.to(self.device)
        self.infer_model.to(self.device).eval()

    def preprocess(self, input: Input) -> Dict[str, Any]:
        img = LoadImage.convert_to_ndarray(input)[:, :, ::-1]
        self.image = np.array(img)

        mean = np.array([0.408, 0.447, 0.470],
                        dtype=np.float32).reshape(1, 1, 3)
        std = np.array([0.289, 0.274, 0.278],
                       dtype=np.float32).reshape(1, 1, 3)
        height, width = img.shape[0:2]
        inp_height, inp_width = self.cfg.input_h, self.cfg.input_w
        c = np.array([width / 2., height / 2.], dtype=np.float32)
        s = max(height, width) * 1.0

        trans_input = get_affine_transform(c, s, 0, [inp_width, inp_height])
        resized_image = cv2.resize(img, (width, height))
        inp_image = cv2.warpAffine(
            resized_image,
            trans_input, (inp_width, inp_height),
            flags=cv2.INTER_LINEAR)
        inp_image = ((inp_image / 255. - mean) / std).astype(np.float32)

        images = inp_image.transpose(2, 0, 1).reshape(1, 3, inp_height,
                                                      inp_width)
        images = torch.from_numpy(images).to(self.device)
        meta = {
            'c': c,
            's': s,
            'input_height': inp_height,
            'input_width': inp_width,
            'out_height': inp_height // 4,
            'out_width': inp_width // 4
        }

        result = {'img': images, 'meta': meta}

        return result

    def distance(self, x1, y1, x2, y2):
        return math.sqrt(pow(x1 - x2, 2) + pow(y1 - y2, 2))

    def crop_image(self, img, position):
        x0, y0 = position[0][0], position[0][1]
        x1, y1 = position[1][0], position[1][1]
        x2, y2 = position[2][0], position[2][1]
        x3, y3 = position[3][0], position[3][1]

        img_width = self.distance((x0 + x3) / 2, (y0 + y3) / 2, (x1 + x2) / 2,
                                  (y1 + y2) / 2)
        img_height = self.distance((x0 + x1) / 2, (y0 + y1) / 2, (x2 + x3) / 2,
                                   (y2 + y3) / 2)

        corners_trans = np.zeros((4, 2), np.float32)
        corners_trans[0] = [0, 0]
        corners_trans[1] = [img_width, 0]
        corners_trans[2] = [img_width, img_height]
        corners_trans[3] = [0, img_height]

        transform = cv2.getPerspectiveTransform(position, corners_trans)
        dst = cv2.warpPerspective(img, transform,
                                  (int(img_width), int(img_height)))
        return dst

    def forward(self, input: Dict[str, Any]) -> Dict[str, Any]:
        pred = self.infer_model(input['img'])
        return {'results': pred, 'meta': input['meta']}

    def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        output = inputs['results'][0]
        meta = inputs['meta']
        hm = output['hm'].sigmoid_()
        wh = output['wh']
        reg = output['reg']
        angle_cls = output['cls'].sigmoid_()
        ftype_cls = output['ftype'].sigmoid_()

        bbox, inds = bbox_decode(hm, wh, reg=reg, K=self.K)
        angle_cls = decode_by_ind(
            angle_cls, inds, K=self.K).detach().cpu().numpy()
        ftype_cls = decode_by_ind(
            ftype_cls, inds,
            K=self.K).detach().cpu().numpy().astype(np.float32)
        bbox = bbox.detach().cpu().numpy()
        for i in range(bbox.shape[1]):
            bbox[0][i][9] = angle_cls[0][i]
        bbox = np.concatenate((bbox, np.expand_dims(ftype_cls, axis=-1)),
                              axis=-1)
        bbox = nms(bbox, 0.3)
        bbox = bbox_post_process(bbox.copy(), [meta['c'].cpu().numpy()],
                                 [meta['s']], meta['out_height'],
                                 meta['out_width'])

        res = []
        angle = []
        sub_imgs = []
        ftype = []
        score = []
        for idx, box in enumerate(bbox[0]):
            if box[8] > 0.3:
                angle.append(int(box[9]))
                res.append(box[0:8])
                sub_img = self.crop_image(self.image,
                                          res[-1].copy().reshape(4, 2))
                if angle[-1] == 1:
                    sub_img = cv2.rotate(sub_img, 2)
                if angle[-1] == 2:
                    sub_img = cv2.rotate(sub_img, 1)
                if angle[-1] == 3:
                    sub_img = cv2.rotate(sub_img, 0)
                sub_imgs.append(sub_img)
                ftype.append(int(box[10]))
                score.append(box[8])

        result = {
            OutputKeys.POLYGONS: res,
            OutputKeys.SCORES: score,
            OutputKeys.OUTPUT_IMGS: sub_imgs,
            OutputKeys.LABELS: angle,
            OutputKeys.LAYOUT: np.array(ftype)
        }
        return result
