#!/usr/bin/env python3
"""Digitization with the EXACT transform from Utils/read_imgs.py:
Resize(224,224) + ToTensor + ImageNet Normalize, and the correct class order."""
import sys
import cv2 as cv
import numpy as np
import onnxruntime as ort
from digitize_full import get_casas, crop_square

TRUE = "rn1qkb1r/pb2pppp/5n2/1ppp4/8/3P1NP1/PPP1PPBP/RNBQ1RK1"
# from read_imgs.py: {'P':0,'N':1,'B':2,'R':3,'Q':4,'K':5,'p':6,'n':7,'b':8,'r':9,'q':10,'k':11}
IDX2SYM = {0:"P",1:"N",2:"B",3:"R",4:"Q",5:"K",6:"p",7:"n",8:"b",9:"r",10:"q",11:"k"}
MEAN = np.array([0.485,0.456,0.406], np.float32)
STD = np.array([0.229,0.224,0.225], np.float32)


def preprocess(bgr):
    rgb = cv.cvtColor(cv.resize(bgr,(224,224)), cv.COLOR_BGR2RGB).astype(np.float32)/255.0
    return ((rgb-MEAN)/STD).transpose(2,0,1)[None]


def true_grid():
    g=[]
    for row in TRUE.split("/"):
        r=[]
        for ch in row: r += ["."]*int(ch) if ch.isdigit() else [ch]
        g.append(r)
    return np.array(g)


def main():
    path = sys.argv[1] if len(sys.argv) > 1 else "samples/board.png"
    im, casas = get_casas(path)
    print("squares:", len(casas))
    crops = [crop_square(im, p) for p in casas[:64]]
    T = true_grid(); occ = (T!="."); n=int(occ.sum())

    for model in ["digitizer3d.fp32.onnx", "digitizer.fp32.onnx"]:
        sess = ort.InferenceSession(f"models_onnx/{model}", providers=["CPUExecutionProvider"])
        syms=[]
        for c in crops:
            if c.size==0: syms.append("?"); continue
            out = sess.run(None, {"input": preprocess(c)})[0].flatten()
            syms.append(IDX2SYM[int(out.argmax())])
        P = np.array(syms, object).reshape(8,8)
        best = max((int(((Q==T)&occ).sum()), k) for k,Q in
                   [("id",P),("rot90",np.rot90(P)),("rot180",np.rot90(P,2)),("rot270",np.rot90(P,3)),
                    ("flipUD",np.flipud(P)),("flipLR",np.fliplr(P)),("T",P.T),("anti",np.fliplr(np.rot90(P)))])
        tag = "3D" if "3d" in model else "REAL"
        print(f"  {tag:4s}: best occupied-match {best[0]:2d}/{n}  ({best[1]})")


if __name__ == "__main__":
    main()