import threading
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor
from queue import Queue
from typing import List, Optional
from PIL import Image
import numpy as np

from surya.detection import batch_detection
from surya.postprocessing.heatmap import keep_largest_boxes, get_and_clean_boxes, get_detected_boxes
from surya.schema import LayoutResult, LayoutBox, TextDetectionResult
from surya.settings import settings


def get_regions_from_detection_result(detection_result: TextDetectionResult, heatmaps: List[np.ndarray], orig_size, id2label, segment_assignment, vertical_line_width=20) -> List[LayoutBox]:
    logits = np.stack(heatmaps, axis=0)
    vertical_line_bboxes = detection_result.vertical_lines
    line_bboxes = detection_result.bboxes

    # Scale back to processor size
    for line in vertical_line_bboxes:
        line.rescale_bbox(orig_size, list(reversed(heatmaps[0].shape)))

    for line in line_bboxes:
        line.rescale(orig_size, list(reversed(heatmaps[0].shape)))

    for bbox in vertical_line_bboxes:
        # Give some width to the vertical lines
        vert_bbox = list(bbox.bbox)
        vert_bbox[2] = min(heatmaps[0].shape[0], vert_bbox[2] + vertical_line_width)

        logits[:, vert_bbox[1]:vert_bbox[3], vert_bbox[0]:vert_bbox[2]] = 0  # zero out where the column lines are

    logits[:, logits[0] >= .5] = 0 # zero out where blanks are

    # Zero out where other segments are
    for i in range(logits.shape[0]):
        logits[i, segment_assignment != i] = 0

    detected_boxes = []
    for heatmap_idx in range(1, len(id2label)):  # Skip the blank class
        heatmap = logits[heatmap_idx]
        if np.max(heatmap) < settings.DETECTOR_BLANK_THRESHOLD:
            continue
        bboxes = get_detected_boxes(heatmap)
        bboxes = [bbox for bbox in bboxes if bbox.area > 25]
        for bb in bboxes:
            bb.fit_to_bounds([0, 0, heatmap.shape[1] - 1, heatmap.shape[0] - 1])

        for bbox in bboxes:
            detected_boxes.append(LayoutBox(polygon=bbox.polygon, label=id2label[heatmap_idx], confidence=1))

    detected_boxes = sorted(detected_boxes, key=lambda x: x.confidence, reverse=True)
    # Expand bbox to cover intersecting lines
    box_lines = defaultdict(list)
    used_lines = set()

    # We try 2 rounds of identifying the correct lines to snap to
    # First round is majority intersection, second lowers the threshold
    for thresh in [.5, .4]:
        for bbox_idx, bbox in enumerate(detected_boxes):
            for line_idx, line_bbox in enumerate(line_bboxes):
                if line_bbox.intersection_pct(bbox) > thresh and line_idx not in used_lines:
                    box_lines[bbox_idx].append(line_bbox.bbox)
                    used_lines.add(line_idx)

    new_boxes = []
    for bbox_idx, bbox in enumerate(detected_boxes):
        if bbox.label == "Picture" and bbox.area < 200: # Remove very small figures
            continue

        # Skip if we didn't find any lines to snap to, except for Pictures and Formulas
        if bbox_idx not in box_lines and bbox.label not in ["Picture", "Formula"]:
            continue

        covered_lines = box_lines[bbox_idx]
        # Snap non-picture layout boxes to correct text boundaries
        if len(covered_lines) > 0 and bbox.label not in ["Picture"]:
            min_x = min([line[0] for line in covered_lines])
            min_y = min([line[1] for line in covered_lines])
            max_x = max([line[2] for line in covered_lines])
            max_y = max([line[3] for line in covered_lines])

            # Tables and formulas can contain text, but text isn't the whole area
            if bbox.label in ["Table", "Formula"]:
                min_x_box = min([b[0] for b in bbox.polygon])
                min_y_box = min([b[1] for b in bbox.polygon])
                max_x_box = max([b[0] for b in bbox.polygon])
                max_y_box = max([b[1] for b in bbox.polygon])

                min_x = min(min_x, min_x_box)
                min_y = min(min_y, min_y_box)
                max_x = max(max_x, max_x_box)
                max_y = max(max_y, max_y_box)

            bbox.polygon = [
                [min_x, min_y],
                [max_x, min_y],
                [max_x, max_y],
                [min_x, max_y]
            ]

        if bbox_idx in box_lines and bbox.label in ["Picture"]:
            bbox.label = "Figure"

        new_boxes.append(bbox)

    # Merge tables together (sometimes one column is detected as a separate table)
    mergeable_types = ["Table", "Picture", "Figure"]
    for ftype in mergeable_types:
        to_remove = set()
        for bbox_idx, bbox in enumerate(new_boxes):
            if bbox.label != ftype or bbox_idx in to_remove:
                continue

            for bbox_idx2, bbox2 in enumerate(new_boxes):
                if bbox2.label != ftype or bbox_idx2 in to_remove or bbox_idx == bbox_idx2:
                    continue

                if bbox.intersection_pct(bbox2, x_margin=.25) > .1:
                    bbox.merge(bbox2)
                    to_remove.add(bbox_idx2)

        new_boxes = [bbox for idx, bbox in enumerate(new_boxes) if idx not in to_remove]

    # Ensure we account for all text lines in the layout
    unused_lines = [line for idx, line in enumerate(line_bboxes) if idx not in used_lines]
    for bbox in unused_lines:
        new_boxes.append(LayoutBox(polygon=bbox.polygon, label="Text", confidence=.5))

    for bbox in new_boxes:
        bbox.rescale(list(reversed(heatmaps[0].shape)), orig_size)

    detected_boxes = [bbox for bbox in new_boxes if bbox.area > 16]

    # Remove bboxes contained inside others, unless they're captions
    contained_bbox = []
    for i, bbox in enumerate(detected_boxes):
        for j, bbox2 in enumerate(detected_boxes):
            if i == j:
                continue

            if bbox2.intersection_pct(bbox) >= .95 and bbox2.label not in ["Caption"]:
                contained_bbox.append(j)

    detected_boxes = [bbox for idx, bbox in enumerate(detected_boxes) if idx not in contained_bbox]

    return detected_boxes


def get_regions(heatmaps: List[np.ndarray], orig_size, id2label, segment_assignment) -> List[LayoutBox]:
    bboxes = []
    for i in range(1, len(id2label)):  # Skip the blank class
        heatmap = heatmaps[i]
        assert heatmap.shape == segment_assignment.shape
        heatmap[segment_assignment != i] = 0  # zero out where another segment is

        # Skip processing empty labels
        if np.max(heatmap) < settings.DETECTOR_BLANK_THRESHOLD:
            continue

        bbox = get_and_clean_boxes(heatmap, list(reversed(heatmap.shape)), orig_size)
        for bb in bbox:
            bboxes.append(LayoutBox(polygon=bb.polygon, label=id2label[i]))

    bboxes = keep_largest_boxes(bboxes)
    return bboxes


def parallel_get_regions(heatmaps: List[np.ndarray], orig_size, id2label, detection_results=None) -> LayoutResult:
    logits = np.stack(heatmaps, axis=0)
    segment_assignment = logits.argmax(axis=0)
    if detection_results is not None:
        bboxes = get_regions_from_detection_result(detection_results, heatmaps, orig_size, id2label,
                                                   segment_assignment)
    else:
        bboxes = get_regions(heatmaps, orig_size, id2label, segment_assignment)

    segmentation_img = Image.fromarray(segment_assignment.astype(np.uint8))

    result = LayoutResult(
        bboxes=bboxes,
        segmentation_map=segmentation_img,
        heatmaps=heatmaps,
        image_bbox=[0, 0, orig_size[0], orig_size[1]]
    )

    return result


def batch_layout_detection(images: List, model, processor, detection_results: Optional[List[TextDetectionResult]] = None, batch_size=None) -> List[LayoutResult]:
    layout_generator = batch_detection(images, model, processor, batch_size=batch_size)
    id2label = model.config.id2label

    results = []
    result_lock = threading.Lock()
    max_workers = min(settings.DETECTOR_POSTPROCESSING_CPU_WORKERS, len(images))
    parallelize = not settings.IN_STREAMLIT and len(images) >= settings.DETECTOR_MIN_PARALLEL_THRESH
    batch_queue = Queue(maxsize=4)

    def inference_producer():
        for batch in layout_generator:
            batch_queue.put(batch)
        batch_queue.put(None)  # Signal end of batches

    def postprocessing_consumer():
        if parallelize:
            with ProcessPoolExecutor(max_workers=max_workers) as executor:
                img_idx = 0
                while True:
                    batch = batch_queue.get()
                    if batch is None:
                        break

                    preds, orig_sizes = batch
                    img_idxs = [img_idx + i for i in range(len(preds))]
                    batch_results = list(executor.map(
                        parallel_get_regions,
                        preds,
                        orig_sizes,
                        [id2label] * len(preds),
                        [detection_results[idx] for idx in img_idxs] if detection_results else [None] * len(preds)
                    ))

                    with result_lock:
                        results.extend(batch_results)
                        img_idx += len(preds)
        else:
            img_idx = 0
            while True:
                batch = batch_queue.get()
                if batch is None:
                    break

                preds, orig_sizes = batch
                for pred, orig_size in zip(preds, orig_sizes):
                    results.append(parallel_get_regions(
                        pred,
                        orig_size,
                        id2label,
                        detection_results[img_idx] if detection_results else None
                    ))

                    img_idx += 1

    # Start producer and consumer threads
    producer = threading.Thread(target=inference_producer)
    consumer = threading.Thread(target=postprocessing_consumer)

    producer.start()
    consumer.start()

    # Wait for both threads to complete
    producer.join()
    consumer.join()

    return results