"""Find a scoreboard's fields and corners in a picture.

usage: python example.py PICTURE  (needs onnxruntime, opencv-python, numpy)"""
import sys
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort

HERE = Path(__file__).parent
CLASSES = ["home", "away", "clock", "period", "other", "board"]
CROP, MARGIN = 256, 0.15
MEAN = np.array([0.485, 0.456, 0.406], np.float32)
STD = np.array([0.229, 0.224, 0.225], np.float32)


def detect(sess, bgr):
    h, w = bgr.shape[:2]
    rgb = cv2.cvtColor(cv2.resize(bgr, (640, 640), interpolation=cv2.INTER_LINEAR), cv2.COLOR_BGR2RGB)
    x = (rgb.astype(np.float32) / 255).transpose(2, 0, 1)[None]
    labels, boxes, scores = sess.run(None, {"images": x, "orig_target_sizes": np.array([[w, h]], np.int64)})
    return labels[0], boxes[0], scores[0]


def corners(sess, bgr, box):
    """The board's four corners (top left, top right, bottom right, bottom left) in picture pixels."""
    x1, y1, x2, y2 = box
    mw, mh = MARGIN * (x2 - x1), MARGIN * (y2 - y1)
    src = np.array([[x1 - mw, y1 - mh], [x2 + mw, y1 - mh], [x2 + mw, y2 + mh], [x1 - mw, y2 + mh]], np.float32)
    H = cv2.getPerspectiveTransform(src, np.array([[0, 0], [CROP, 0], [CROP, CROP], [0, CROP]], np.float32))
    crop = cv2.cvtColor(cv2.warpPerspective(bgr, H, (CROP, CROP)), cv2.COLOR_BGR2RGB).astype(np.float32) / 255
    c = sess.run(None, {"crop": ((crop - MEAN) / STD).transpose(2, 0, 1)[None]})[0][0] * CROP
    return cv2.perspectiveTransform(c.reshape(-1, 1, 2), np.linalg.inv(H)).reshape(4, 2)


def main():
    bgr = cv2.imread(sys.argv[1])
    fields = ort.InferenceSession(str(HERE / "fields_board.onnx"), providers=["CPUExecutionProvider"])
    corner_net = ort.InferenceSession(str(HERE / "corners.onnx"), providers=["CPUExecutionProvider"])
    labels, boxes, scores = detect(fields, bgr)
    for i, name in enumerate(CLASSES[:4]):  # the best box of each role
        sel = (labels == i) & (scores >= 0.3)
        if sel.any():
            k = np.argmax(np.where(sel, scores, -1))
            print(f"{name}: box {np.round(boxes[k]).astype(int).tolist()} score {scores[k]:.2f}")
    sel = (labels == CLASSES.index("board")) & (scores >= 0.3)
    if sel.any():
        k = np.argmax(np.where(sel, scores, -1))
        print("board corners:", np.round(corners(corner_net, bgr, boxes[k])).astype(int).tolist())


if __name__ == "__main__":
    main()