"""Оценка детектора строк на отложенных выпусках.

Меряем то, что важно продукту: сколько строк найдено (полнота), сколько
выдумано (точность) и на скольких страницах порядок чтения совпал с вёрсткой.
"""
import argparse, json, os, sys, time

import numpy as np
import torch
from PIL import Image, ImageFilter

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from linedet import LineDet
from detect import lines_from_prob, iou_match


@torch.no_grad()
def run_page(model, im, dev, stride, max_side=2600):
    """Карта вероятностей для целой страницы; длинная сторона ограничена."""
    s = min(1.0, max_side / max(im.width, im.height))
    if s < 1.0:
        im = im.resize((int(im.width * s), int(im.height * s)), Image.BILINEAR)
    w, h = im.width, im.height
    pw, ph = (w + 31) // 32 * 32, (h + 31) // 32 * 32
    a = np.zeros((ph, pw), np.float32)
    a[:h, :w] = 1.0 - np.asarray(im, np.float32) / 255.0
    x = torch.from_numpy(a)[None, None].to(dev)
    with torch.amp.autocast("cuda", enabled=(dev == "cuda")):
        prob, _ = model(x)
    prob = torch.sigmoid(prob.float())[0, 0].cpu().numpy()
    prob = prob[:h // stride, :w // stride]
    return prob, s


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--data", required=True)
    ap.add_argument("--pages", type=int, default=0)
    ap.add_argument("--skip", type=int, default=0,
                    help="пропустить первые N страниц: по ним отбирался чекпоинт")
    ap.add_argument("--iou", type=float, default=0.5)
    ap.add_argument("--max-side", type=int, default=2000)
    ap.add_argument("--unclip", type=float, default=1.6)
    ap.add_argument("--book-scale", type=float, default=4.0)
    ap.add_argument("--book", action="store_true",
                    help="книжно-экранное условие: та же разметка, крупнее и бледнее")
    ap.add_argument("--max-gap", type=float, default=2.0,
                    help="порог склейки кусков строки в долях высоты строки")
    a = ap.parse_args()

    dev = "cuda" if torch.cuda.is_available() else "cpu"
    ck = torch.load(a.ckpt, map_location="cpu", weights_only=False)
    model = LineDet(ck.get("width", 32)).to(dev).eval()
    model.load_state_dict(ck["model"])
    stride = ck.get("stride", 2)

    rows = [json.loads(l) for l in
            open(os.path.join(a.data, "index.jsonl"), encoding="utf-8")]
    rows = rows[a.skip:]
    if a.pages:
        rows = rows[:a.pages]

    tot_rec = tot_prec = 0.0
    n_gt = n_pred = n_hit = 0
    ok_pairs = all_pairs = 0
    t0 = time.time()
    for r in rows:
        im = Image.open(os.path.join(a.data, r["f"])).convert("L")
        pre, off = 1.0, (0, 0)
        gt_src = r["b"]
        if a.book:
            # Книжно-экранное условие. Просто увеличить страницу мало:
            # run_page тут же ужмёт её обратно под max_side, и высота строки
            # не изменится. Поэтому после увеличения вырезаем окно размером
            # с обычную страницу — получается маленький лист с крупным
            # бледным текстом, как на скриншоте книги.
            pre = a.book_scale
            big = im.resize((int(im.width * pre), int(im.height * pre)), Image.BILINEAR)
            cw = min(big.width, a.max_side)
            ch = min(big.height, a.max_side)
            # окно центрируем по массе текста: середина увеличенной полосы
            # часто приходится на фотографию, и в окне не оказывается строк
            cx = float(np.mean([(b[0] + b[2]) / 2 for b in r["b"]])) * pre
            cy = float(np.mean([(b[1] + b[3]) / 2 for b in r["b"]])) * pre
            ox = int(min(max(0.0, cx - cw / 2), big.width - cw))
            oy = int(min(max(0.0, cy - ch / 2), big.height - ch))
            im = big.crop((ox, oy, ox + cw, oy + ch))
            arr = np.asarray(im, np.float32) / 255.0
            im = Image.fromarray(((0.30 + arr * 0.60) * 255).astype(np.uint8)).filter(
                ImageFilter.GaussianBlur(0.6))
            off = (ox, oy)
            gt_src = [[b[0] * pre - ox, b[1] * pre - oy, b[2] * pre - ox, b[3] * pre - oy]
                      for b in r["b"]]
            # оставляем строки, попавшие в окно хотя бы наполовину, и
            # обрезаем их по окну: иначе точность падает на строках,
            # которые видны у края лишь частично
            keep = []
            for b in gt_src:
                x0, y0 = max(0.0, b[0]), max(0.0, b[1])
                x1, y1 = min(float(cw), b[2]), min(float(ch), b[3])
                if x1 <= x0 or y1 <= y0:
                    continue
                if (x1 - x0) * (y1 - y0) < 0.5 * (b[2] - b[0]) * (b[3] - b[1]):
                    continue
                keep.append([x0, y0, x1, y1])
            gt_src = keep
            pre = 1.0
        prob, s = run_page(model, im, dev, stride, a.max_side)
        w_s, h_s = int(im.width * s), int(im.height * s)
        pred, lh = lines_from_prob(prob, stride, w_s, h_s, a.max_gap, a.unclip)
        gt = [tuple(int(v * pre * s) for v in b) for b in gt_src]
        rec, prec, hit = iou_match(pred, gt, a.iou)
        n_gt += len(gt); n_pred += len(pred); n_hit += hit

        # порядок чтения: сравниваем с порядком строк в вёрстке исходника
        ordered = pred
        idx = []
        for b in ordered:
            best, bi = 0.0, -1
            ba = max(1, (b[2]-b[0]) * (b[3]-b[1]))
            for i, g in enumerate(gt):
                ox = min(b[2], g[2]) - max(b[0], g[0])
                oy = min(b[3], g[3]) - max(b[1], g[1])
                if ox <= 0 or oy <= 0:
                    continue
                v = ox * oy / ba
                if v > best:
                    best, bi = v, i
            if best > 0.25:
                idx.append(bi)
        ok_pairs += sum(1 for p, q in zip(idx, idx[1:]) if q >= p)
        all_pairs += max(1, len(idx) - 1)

    print(f"страниц {len(rows)}, строк в эталоне {n_gt}, найдено рамок {n_pred}")
    print(f"IoU {a.iou}: полнота {n_hit/max(1,n_gt)*100:.1f}%, "
          f"точность {n_hit/max(1,n_pred)*100:.1f}%")
    print(f"порядок чтения: {ok_pairs/all_pairs*100:.1f}% соседних пар верно")
    print(f"{(time.time()-t0)/len(rows)*1000:.0f} мс на страницу ({dev})")


if __name__ == "__main__":
    main()