#!/usr/bin/env python3
"""Decisive test: 4-corner grid (robust) + EXACT Preprocess crop (bbox + 90px top),
classify the 64 squares, compare to FEN under all 8 orientations and 3 norms."""
import cv2 as cv
import numpy as np
import onnxruntime as ort

TRUE = "rn1qkb1r/pb2pppp/5n2/1ppp4/8/3P1NP1/PPP1PPBP/RNBQ1RK1"
IDX2SYM = {0:"K",1:"Q",2:"R",3:"B",4:"N",5:"P",6:"k",7:"q",8:"r",9:"b",10:"n",11:"p"}
# corners on the FULL photo (TL,TR,BR,BL of the 8x8 area)
CORNERS_FULL = np.array([[520,140],[1360,150],[1430,915],[452,918]], np.float32)


def crop_image(im):
    cx, cy = im.shape[1]//2, im.shape[0]//2
    return im[cy-520:cy+450, cx-550:cx+550], (cx-550, cy-520)


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():
    img = cv.imread("samples/board.png")
    im, (ox, oy) = crop_image(img)
    corners = CORNERS_FULL - [ox, oy]                       # into crop space
    unit = np.array([[0,0],[8,0],[8,8],[0,8]], np.float32)
    Hmat = cv.getPerspectiveTransform(unit, corners)        # unit grid -> image

    def node(c, r):
        p = cv.perspectiveTransform(np.array([[[c, r]]], np.float32), Hmat)[0][0]
        return p

    crops = []
    for r in range(8):
        for c in range(8):
            pts = [node(c, r), node(c+1, r), node(c+1, r+1), node(c, r+1)]
            xs = [p[0] for p in pts]; ys = [p[1] for p in pts]
            x1, x2 = int(min(xs)), int(max(xs))
            y1 = max(0, int(min(ys)) - 90); y2 = int(max(ys))
            crops.append(im[y1:y2, x1:x2])

    sess = ort.InferenceSession("models_onnx/digitizer.fp32.onnx", providers=["CPUExecutionProvider"])
    T = true_grid()
    occ = (T != ".")
    n_occ = int(occ.sum())

    for nm, fn in [("/255", lambda r: r/255.0),
                   ("imagenet", lambda r: (r/255.0-[0.485,0.456,0.406])/[0.229,0.224,0.225]),
                   ("[-1,1]", lambda r: r/127.5-1)]:
        P = np.empty((8,8), object)
        for i, crop in enumerate(crops):
            if crop.size == 0:
                P[i//8, i%8] = "?"; continue
            rgb = cv.cvtColor(cv.resize(crop,(96,96)), cv.COLOR_BGR2RGB).astype(np.float32)
            x = np.asarray(fn(rgb), np.float32).transpose(2,0,1)[None]
            out = sess.run(None, {"input": x})[0].flatten()
            P[i//8, i%8] = IDX2SYM[int(out.argmax())]
        # best over 8 dihedral orientations, counting only TRUE-occupied squares
        best = 0; bestname = ""
        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)),
                     ("transp",P.T),("anti",np.fliplr(np.rot90(P)))]:
            if Q.shape!=T.shape: continue
            m = int(((Q==T)&occ).sum())
            if m>best: best, bestname = m, k
        print(f"[{nm:9s}] best occupied-match {best:2d}/{n_occ}  ({bestname})")


if __name__ == "__main__":
    main()