"""Разбор карт детектора в рамки строк и порядок чтения.

Модель выдаёт карту по сжатым рамкам: соседние строки на ней разделены. Здесь
эти пятна снова разворачиваются до полной строки (unclip) и раскладываются в
порядок чтения рекурсивным XY-разрезом — уже по чистым рамкам, а не по
морфологическим пятнам.
"""
import json
import os

import numpy as np
from scipy import ndimage

UNCLIP = 1.6          # во сколько раз площадь пятна возвращается к полной строке
PROB_TH = 0.3         # порог по карте вероятностей (как в DB)


def unclip_box(x0, y0, x1, y1, ratio=UNCLIP):
    """Развернуть сжатое пятно обратно: смещение наружу по формуле Vatti."""
    w, h = x1 - x0, y1 - y0
    if w <= 0 or h <= 0:
        return None
    d = (w * h) * ratio / (2.0 * (w + h))
    return (x0 - d, y0 - d, x1 + d, y1 + d)


def boxes_from_maps(prob, stride=2, prob_th=PROB_TH, min_h=3, min_w=6, unclip=UNCLIP):
    """Связные компоненты карты -> рамки строк в координатах входа."""
    mask = prob > prob_th
    lbl, n = ndimage.label(mask)
    if n == 0:
        return []
    out = []
    for sl in ndimage.find_objects(lbl):
        y0, y1 = sl[0].start, sl[0].stop
        x0, x1 = sl[1].start, sl[1].stop
        if y1 - y0 < 1 or x1 - x0 < 1:
            continue
        b = unclip_box(x0, y0, x1, y1, unclip)
        if b is None:
            continue
        bx0, by0, bx1, by1 = (v * stride for v in b)
        if bx1 - bx0 < min_w or by1 - by0 < min_h:
            continue
        out.append((int(round(bx0)), int(round(by0)),
                    int(round(bx1)), int(round(by1))))
    return out


def clip_boxes(boxes, w, h):
    res = []
    for x0, y0, x1, y1 in boxes:
        x0, y0 = max(0, x0), max(0, y0)
        x1, y1 = min(w, x1), min(h, y1)
        if x1 > x0 and y1 > y0:
            res.append((x0, y0, x1, y1))
    return res


def _widest_gap(boxes, axis, thr, skip_wide=0.0):
    lo = min(b[axis] for b in boxes)
    hi = max(b[axis + 2] for b in boxes)
    if hi - lo < 2:
        return None
    src = boxes
    if skip_wide:
        narrow = [b for b in boxes if b[axis + 2] - b[axis] <= skip_wide * (hi - lo)]
        if narrow:
            src = narrow
    occ = np.zeros(hi - lo + 1, dtype=bool)
    for b in src:
        occ[b[axis] - lo:b[axis + 2] - lo] = True
    best, run = None, None
    for i, v in enumerate(occ):
        if not v:
            if run is None:
                run = i
        else:
            if run is not None:
                if i - run >= thr and (best is None or i - run > best[1]):
                    best = (lo + (run + i) // 2, i - run)
                run = None
    return best


def xy_regions(boxes, page_w, lh=None):
    """Разложение страницы на регионы в порядке чтения рекурсивным XY-разрезом.

    Режем регион по самому широкому сквозному просвету: сначала вертикальному
    (колонки), иначе горизонтальному (полосы). Заголовок или снимок поперёк
    колонок не должен запрещать деление — такие объекты исключаются из
    профиля и выдаются перед колонками, как их и читают.

    Возвращаем листья, а не плоский список: внутри листа колонок заведомо нет,
    и куски строки можно склеивать широким порогом, не рискуя сшить соседние
    колонки.
    """
    if not boxes:
        return []
    if lh is None:
        lh = max(6, int(np.median([b[3] - b[1] for b in boxes])))
    vthr = max(int(page_w * 0.008), int(lh * 0.8))
    hthr = max(int(lh * 0.6), 4)

    def cut(reg, depth=0):
        if len(reg) < 2 or depth > 14:
            return [sorted(reg, key=lambda b: (b[1], b[0]))]
        v = _widest_gap(reg, 0, vthr, skip_wide=0.6)
        if v is not None:
            x = v[0]
            span = [b for b in reg if b[0] < x < b[2]]
            left = [b for b in reg if b[2] <= x]
            right = [b for b in reg if b[0] >= x]
            # Разрез пересекает много строк — значит это не межколоночный
            # зазор, а случайный просвет между узкими кусками. Иначе обычные
            # строки во всю ширину принимаются за заголовок поперёк колонок
            # и выдаются раньше него.
            if left and right and len(span) <= 0.2 * len(reg):
                head = [sorted(span, key=lambda b: (b[1], b[0]))] if span else []
                return head + cut(left, depth + 1) + cut(right, depth + 1)
        h = _widest_gap(reg, 1, hthr)
        if h is not None:
            y = h[0]
            top = [b for b in reg if b[3] <= y]
            bot = [b for b in reg if b[1] >= y]
            mid = [b for b in reg if b[1] < y < b[3]]
            if top and bot:
                mid_leaf = [sorted(mid, key=lambda b: (b[1], b[0]))] if mid else []
                return cut(top, depth + 1) + mid_leaf + cut(bot, depth + 1)
        return [sorted(reg, key=lambda b: (b[1], b[0]))]

    return cut(list(boxes))


def dedup(boxes, iou_th=0.4):
    """Убираем перекрывающиеся рамки: детектор иногда даёт на одну строку две."""
    if not boxes:
        return boxes
    order = sorted(boxes, key=lambda b: -( (b[2]-b[0]) * (b[3]-b[1]) ))
    kept = []
    for b in order:
        drop = False
        for i, k in enumerate(kept):
            ox = min(b[2], k[2]) - max(b[0], k[0])
            oy = min(b[3], k[3]) - max(b[1], k[1])
            if ox <= 0 or oy <= 0:
                continue
            inter = ox * oy
            small = min((b[2]-b[0])*(b[3]-b[1]), (k[2]-k[0])*(k[3]-k[1]))
            if inter / max(1, small) > iou_th:
                kept[i] = (min(b[0], k[0]), min(b[1], k[1]),
                           max(b[2], k[2]), max(b[3], k[3]))
                drop = True
                break
        if not drop:
            kept.append(b)
    return kept


def merge_line_fragments(boxes, lh, max_gap=2.0):
    """Склеиваем куски одной строки.

    В выключенном по ширине наборе пробелы между словами широкие, и детектор
    честно выдаёт два пятна вместо одной строки. Склеиваем только то, что
    лежит на одной строке по вертикали и разделено зазором меньше
    межколоночного — иначе сольются соседние колонки.
    """
    if not boxes:
        return boxes
    lim = max_gap * lh
    out = sorted(boxes, key=lambda b: (b[1], b[0]))
    changed = True
    while changed:
        changed = False
        res, used = [], [False] * len(out)
        for i, a in enumerate(out):
            if used[i]:
                continue
            x0, y0, x1, y1 = a
            for j in range(i + 1, len(out)):
                if used[j]:
                    continue
                b = out[j]
                ov = min(y1, b[3]) - max(y0, b[1])
                if ov < 0.6 * min(y1 - y0, b[3] - b[1]):
                    continue
                gap = max(b[0] - x1, x0 - b[2])
                if gap > lim:
                    continue
                x0, y0 = min(x0, b[0]), min(y0, b[1])
                x1, y1 = max(x1, b[2]), max(y1, b[3])
                used[j] = True
                changed = True
            res.append((x0, y0, x1, y1))
        out = sorted(res, key=lambda b: (b[1], b[0]))
    return out


def lines_from_prob(prob, stride, page_w, page_h, max_gap=2.0, unclip=UNCLIP):
    """Полный разбор карты: рамки -> дубли -> регионы -> склейка внутри региона.

    Порядок операций важен. Склейка до раскладки на регионы сшивает соседние
    колонки, а склейка узким порогом оставляет выключенную строку разорванной
    на четыре куска, и куски выходят не по порядку. Внутри региона колонок нет,
    поэтому порог можно держать широким.
    """
    boxes = clip_boxes(boxes_from_maps(prob, stride, unclip=unclip), page_w, page_h)
    if not boxes:
        return [], 12
    boxes = dedup(boxes)
    lh = max(6, int(np.median([b[3] - b[1] for b in boxes])))
    out = []
    for reg in xy_regions(boxes, page_w, lh):
        out += merge_line_fragments(reg, lh, max_gap)
    lh = max(6, int(np.median([b[3] - b[1] for b in out]))) if out else lh
    return out, lh


def iou_match(pred, gt, thr=0.5):
    """Полнота и точность по рамкам при заданном пороге IoU."""
    if not gt:
        return 0.0, 0.0, 0
    used = [False] * len(pred)
    hit = 0
    for g in gt:
        best, bi = 0.0, -1
        for i, p in enumerate(pred):
            if used[i]:
                continue
            ox = min(g[2], p[2]) - max(g[0], p[0])
            oy = min(g[3], p[3]) - max(g[1], p[1])
            if ox <= 0 or oy <= 0:
                continue
            inter = ox * oy
            u = (g[2]-g[0])*(g[3]-g[1]) + (p[2]-p[0])*(p[3]-p[1]) - inter
            v = inter / u if u > 0 else 0.0
            if v > best:
                best, bi = v, i
        if best >= thr:
            used[bi] = True
            hit += 1
    rec = hit / len(gt)
    prec = hit / len(pred) if pred else 0.0
    return rec, prec, hit


class LineDetector:
    """Обученный детектор строк поверх onnxruntime (CPU).

    Заменяет морфологическую сегментацию: у той на 36 отложенных полосах
    полнота 50.1% и точность 37.6% при IoU 0.5 — почти две трети «строк»
    строками не были. Пороги под это подобрать нельзя, поэтому рамки
    предсказывает сеть, обученная на разметке из born-digital PDF.
    """

    # Длинная сторона 2000 — не компромисс ради скорости, а оптимум по
    # качеству: на отложенных полосах 93.5/94.9 против 84.3/87.3 при 2600.
    # Модель обучалась на кропах со случайным масштабом, и в этой полосе
    # масштабов строка ей привычнее.
    VPAD = 0.3          # доля высоты строки, на которую рамка растёт вверх и вниз

    def __init__(self, model_dir=None, max_side=2000, vpad=None):
        import onnxruntime as ort
        d = model_dir or os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                      "model_det")
        self.meta = json.load(open(os.path.join(d, "sakha_det.json"), encoding="utf-8"))
        self.stride = self.meta.get("stride", 2)
        self.pad_to = self.meta.get("pad_to", 32)
        self.max_side = max_side
        self.vpad = self.VPAD if vpad is None else vpad
        so = ort.SessionOptions()
        so.intra_op_num_threads = max(1, (os.cpu_count() or 4) // 2)
        self.sess = ort.InferenceSession(os.path.join(d, "sakha_det.onnx"),
                                         sess_options=so,
                                         providers=["CPUExecutionProvider"])

    def __call__(self, im):
        """im — страница в градациях серого. Возвращает (рамки, высота строки)."""
        from PIL import Image
        w0, h0 = im.width, im.height
        s = min(1.0, self.max_side / max(w0, h0))
        work = im if s >= 1.0 else im.resize((int(w0 * s), int(h0 * s)), Image.BILINEAR)
        w, h = work.width, work.height
        k = self.pad_to
        pw, ph = (w + k - 1) // k * k, (h + k - 1) // k * k
        a = np.zeros((ph, pw), np.float32)
        a[:h, :w] = 1.0 - np.asarray(work.convert("L"), np.float32) / 255.0
        prob = self.sess.run(None, {"image": a[None, None]})[0][0, 0]
        prob = prob[:h // self.stride, :w // self.stride]

        boxes, lh = lines_from_prob(prob, self.stride, w, h)
        if s < 1.0:                       # возвращаемся в координаты исходной страницы
            inv = 1.0 / s
            boxes = [(int(b[0] * inv), int(b[1] * inv),
                      int(b[2] * inv), int(b[3] * inv)) for b in boxes]
            boxes = clip_boxes(boxes, w0, h0)
            lh = int(lh * inv)
        # Рамку расширяем только по вертикали. Модель учится на сжатых рамках,
        # и после разворота нижняя граница садится на базовую линию: выносные
        # элементы букв обрезаются, «у» читается как «ү», строчные как
        # прописные. По горизонтали не расширяем — там соседние колонки.
        if self.vpad:
            dy = max(1, int(lh * self.vpad))
            boxes = clip_boxes([(b[0], b[1] - dy, b[2], b[3] + dy) for b in boxes],
                               w0, h0)
        return boxes, max(6, lh)