from typing import List
from .pdf_detection import Pipeline
from utils import non_max_suppression, merge_text_and_title_boxes, LayoutBox, PageDetectionResult
from tqdm import tqdm


"""
    0 - Text
    1 - Title
    2 - Figure
    3 - Figure caption
    4 - Table
    5 - Table caption
    6 - Header
    7 - Footer
    8 - Reference
    9 - Equation
"""
pipeline = Pipeline('./models/PaddleDetection/inference_model/picodet_lcnet_x1_0_fgd_layout_cdla_infer')

effective_labels = [0, 1, 2, 4, 5]
# nms优先级，索引越低优先级越低，box重叠时优先保留表格
label_scores = [1, 5, 0, 2, 4]
expand_pixel = 10


def layout_analysis(image_paths) -> List[PageDetectionResult]:
    layout_analysis_results = []
    for image_path in tqdm(image_paths, '版面分析'):
        page_detecion_outputs = pipeline(image_path)
        layout_boxes = []
        for i in range(len(page_detecion_outputs)):
            clsid, box, confidence = page_detecion_outputs[i]
            if clsid in effective_labels:
                layout_boxes.append(LayoutBox(clsid, box, confidence))
        page_detecion_outputs = PageDetectionResult(layout_boxes, image_path)

        scores = []
        poses = []
        for box in page_detecion_outputs.boxes:
            # 相同的label重叠时，保留面积更大的
            area = (box.pos[3] - box.pos[1]) * (box.pos[2] - box.pos[0])
            area_score = area / 5000000
            scores.append(label_scores.index(box.clsid) + area_score)
            poses.append(box.pos)
        indices = non_max_suppression(poses, scores, 0.2)
        _boxes = []
        for i in indices:
            _boxes.append(page_detecion_outputs.boxes[i])
        page_detecion_outputs.boxes = _boxes

        for i in range(len(page_detecion_outputs.boxes) - 1, -1, -1):
            box = page_detecion_outputs.boxes[i]
            if box.clsid in (0, 5):
                # 移除Table box和Figure box中的Table caption box和Text box (有些扫描件会被识别为Figure)
                for _box in page_detecion_outputs.boxes:
                    if _box.clsid != 2 and _box.clsid != 4:
                        continue
                    if box.pos[0] > _box.pos[0] and box.pos[1] > _box.pos[1] and box.pos[2] < _box.pos[2] and box.pos[3] < _box.pos[3]:
                        page_detecion_outputs.boxes.remove(box)

        # 将text和title合并起来，便于转成markdown格式
        page_detecion_outputs.boxes = merge_text_and_title_boxes(page_detecion_outputs.boxes, (0, 1))
        # 对box进行排序
        page_detecion_outputs.boxes.sort(key=lambda x: (x.pos[1], x.pos[0]))
        layout_analysis_results.append(page_detecion_outputs)

    return layout_analysis_results