# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 Decosa
"""Score a page read by a document reader: P(each table cell or text region is misread) and the likely error.
pip install torch safetensors numpy pillow   (transformers only for a vision variant; this one has none)

    python usage.py page.png page.json      # prints the 20 units most likely misread

The page dict is one page of a decosa.docreader.v1 document: {"elements": [{id, kind, bbox: [x, y, w, h], text?,
table?: {rows, cols, cells: [{r, c, rs, cs, text, tok?}]}, source, conf_bp?, min_bp?, ...}]}, with bboxes in the pixels of
the page image you pass. Cell token confidences (cell["tok"]) are optional but help a lot (see the eval).
"""
import json
import sys

import numpy as np
import torch
from PIL import Image
from safetensors.torch import load_file

import features as F
from model import TYPES, ReadConf, encode_text

D = "."
cfg = json.load(open(f"{D}/readconf_config.json"))
vcfg = None
if cfg.get("vision"):
    from transformers import Siglip2VisionConfig
    vcfg = Siglip2VisionConfig.from_dict(cfg["vision"]["config"])
m = ReadConf(cfg["n_feat"], vision_config=vcfg, vision_layers=(cfg.get("vision") or {}).get("layers", 0))
m.load_state_dict(load_file(f"{D}/model.safetensors"))
m.eval()


@torch.no_grad()
def score(image: Image.Image, page: dict) -> list[dict]:
    us = F.page_units(image.convert("RGB"), page)
    if not us:
        return []
    x = torch.tensor(np.array([u["x"] for u in us], np.float32))
    ch = encode_text([u["text"] for u in us])
    img = None
    if m.vision_layers:
        w = cfg["crop"][1]
        a = np.full((len(us), F.CROP_H, w), 255, np.uint8)
        for i, u in enumerate(us):
            c = F.crop(image.convert("RGB"), u, max_w=w)
            a[i, :, : c.width] = np.asarray(c, np.uint8)
        img = torch.tensor(a)
    lw, lt = m(x, ch, img)
    p, tp = m.p_wrong(lw), lt.softmax(-1)
    out = []
    for i, u in enumerate(us):
        # No thresholds ship with this model: rank by p_wrong and review from the top (see the model card).
        out.append({"uid": u["uid"], "text": u["text"], "p_wrong": round(float(p[i]), 4), "type": TYPES[int(tp[i].argmax())]})
    return out


if __name__ == "__main__":
    im = Image.open(sys.argv[1])
    page = json.load(open(sys.argv[2]))
    if "pages" in page:  # a whole decosa.docreader.v1 document: score its first page
        page = page["pages"][0]
    for r in sorted(score(im, page), key=lambda r: -r["p_wrong"])[:20]:
        print(r)