import math
import numpy as np
from typing import List
import cv2
import collections
import logging
import imgaug.augmenters as iaa
from imgaug.augmentables.polys import Polygon, PolygonsOnImage

from sdsv_dewarp.models import PaddleTextClassifier, PaddleTextDetector
from sdsv_dewarp.config import Cfg
from .utils import *


MIN_LONG_EDGE = 40**2
NUMBER_BOX_FOR_ALIGNMENT = 200
MAX_ANGLE = 180
MIN_ANGLE = 1
MIN_NUM_BOX_TEXT = 3
CROP_SIZE = 3000

logging.basicConfig(level=logging.INFO)
LOGGER = logging.getLogger(__name__)


class AlignImage:
    """Rotate image to 0 degree
    Args:
        text_detector (deepmodel): Text detection model
        text_cls (deepmodel): Text classification model (0 or 180)

    Return:
        is_blank (bool): Blank image when haven't boxes text
        image_align: Image after alignment
        angle_align: Degree of angle alignment
    """

    def __init__(self, text_detector: dict, text_cls: dict, device: str = 'cpu'):
        self.text_detector = None
        self.text_cls = None
        self.use_gpu = True if device != 'cpu' else False

        self._init_model(text_detector, text_cls)

    def _init_model(self, text_detector, text_cls):
        det_config = Cfg.load_config_from_file(text_detector['config'])
        det_config['model_dir'] = text_detector['weight']
        cls_config = Cfg.load_config_from_file(text_cls['config'])
        cls_config['model_dir'] = text_cls['weight']

        self.text_detector = PaddleTextDetector(config=det_config, use_gpu=self.use_gpu)
        self.text_cls = PaddleTextClassifier(config=cls_config, use_gpu=self.use_gpu)

    def _cal_width(self, poly_box):
        """Calculate width of a polygon [[x1, y1], [x2, y2], [x3, y3], [x4, y4]]"""
        tl, tr, br, bl = poly_box
        edge_s, edge_l = distance(tl, tr), distance(tr, br)

        return max(edge_s, edge_l)

    def _get_most_frequent(self, values):
        values = np.array(values)
        # create the histogram
        hist, bins = np.histogram(values, bins=np.arange(0, 181, 10))

        # get the index of the most frequent angle
        index = np.argmax(hist)

        # get the most frequent angle
        most_frequent_angle = (bins[index] + bins[index + 1]) / 2

        return most_frequent_angle

    def _cal_angle(self, poly_box):
        """Calculate the angle between two point"""
        a = poly_box[0]
        b = poly_box[1]
        c = poly_box[2]

        # Get the longer edge
        if distance(a, b) >= distance(b, c):
            x, y = a, b
        else:
            x, y = b, c

        angle = math.degrees(math.atan2(-(y[1] - x[1]), y[0] - x[0]))

        if angle < 0:
            angle = 180 - abs(angle)

        return angle

    def _reject_outliers(self, data, m=5.0):
        """Remove noise angle"""
        list_index = np.arange(len(data))
        d = np.abs(data - np.median(data))
        mdev = np.median(d)
        s = d / (mdev if mdev else 1.0)

        return list_index[s < m], data[s < m]

    def __call__(self, image):
        """image (np.ndarray): BGR image"""

        # Crop center image to increase speed of text detection

        image_resized = crop_image(image, crop_size=CROP_SIZE).copy() if max(image.shape) > CROP_SIZE else image.copy()
        poly_box_texts = self.text_detector(image_resized)

        # draw_img = vis_ocr(
        #     image_resized,
        #     poly_box_texts,
        # )
        # cv2.imwrite("draw_img.jpg", draw_img)

        is_blank = False

        # Check image is blank
        if len(poly_box_texts) <= MIN_NUM_BOX_TEXT:
            is_blank = True
            return image, is_blank, 0

        # # Crop document
        # poly_np = np.array(poly_box_texts)
        # min_x = poly_box_texts[:, 0].min()
        # max_x = poly_box_texts[:, 2].max()
        # min_y = poly_box_texts[:, 1].min()
        # max_y = poly_box_texts[:, 3].max()

        # Filter small poly
        poly_box_areas = [
            [self._cal_width(poly_box), id]
            for id, poly_box in enumerate(poly_box_texts)
        ]

        poly_box_areas = sorted(poly_box_areas)[-NUMBER_BOX_FOR_ALIGNMENT:]
        poly_box_areas = [poly_box_texts[id[1]] for id in poly_box_areas]

        # Calculate angle
        list_angle = [self._cal_angle(poly_box) for poly_box in poly_box_areas]
        list_angle = [angle if angle >= MIN_ANGLE else 180 for angle in list_angle]

        # LOGGER.info(f"List angle before reject outlier: {list_angle}")
        list_angle = np.array(list_angle)
        list_index, list_angle = self._reject_outliers(list_angle)
        # LOGGER.info(f"List angle after reject outlier: {list_angle}")

        if len(list_angle):

            frequent_angle = self._get_most_frequent(list_angle)
            list_angle = [angle for angle in list_angle if abs(angle - frequent_angle) <= 45]
            # LOGGER.info(f"List angle after reject angle: {list_angle}")
            angle = np.mean(list_angle)
        else:
            angle = 0

        # LOGGER.info(f"Avg angle: {angle}")

        # Reuse poly boxes detected by text detection
        polys_org = PolygonsOnImage(
            [Polygon(poly_box_areas[index]) for index in list_index],
            shape=image_resized.shape,
        )
        seq_augment = iaa.Sequential([iaa.Rotate(angle, fit_output=True, order=3)])

        # Rotate image by degree
        if angle >= MIN_ANGLE and angle <= MAX_ANGLE:
            image_resized, polys_aug = seq_augment(
                image=image_resized, polygons=polys_org
            )
        else:
            angle = 0
            image_resized, polys_aug = image_resized, polys_org

        # cv2.imwrite("image_resized.jpg", image_resized)
        
        # Classify image 0 or 180 degree
        list_poly = [poly.coords for poly in polys_aug]

        image_crop_list = [
            dewarp_by_polygon(image_resized, poly)[0] for poly in list_poly
        ]

        cls_res = self.text_cls(image_crop_list)
        cls_labels = [cls_[0] for cls_ in cls_res[1]]
        # LOGGER.info(f"Angle lines: {cls_labels}")
        counter = collections.Counter(cls_labels)

        angle_align = angle
        if counter["0"] <= counter["180"]:
            aug = iaa.Rotate(angle + 180, fit_output=True, order=3)
            angle_align = angle + 180
        else:
            aug = iaa.Rotate(angle, fit_output=True, order=3)

        # Rotate the image by degree
        image = aug.augment_image(image)

        return image, is_blank, angle_align
        # return image
