#!/usr/bin/env python3
"""Full image -> FEN: exact grid + crop + transform + empty detection (binarize
center) + flipUD orientation. Compares to the known FEN for i10."""
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"
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, STD = np.array([0.485,0.456,0.406],np.float32), 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 is_empty(im, quad):
    cx = int(np.mean([p[0] for p in quad])); cy = int(np.mean([p[1] for p in quad]))
    w = max(6, int(0.22*(max(p[0] for p in quad)-min(p[0] for p in quad))))
    patch = im[max(0,cy-w):cy+w, max(0,cx-w):cx+w]
    if patch.size == 0:
        return True
    g = cv.cvtColor(patch, cv.COLOR_BGR2GRAY)
    return g.std() < 18      # flat colour -> empty


def grid_to_fen(grid):
    out=[]
    for row in grid:
        s=""; e=0
        for cell in row:
            if cell==".": e+=1
            else:
                if e: s+=str(e); e=0
                s+=cell
        if e: s+=str(e)
        out.append(s)
    return "/".join(out)


def main():
    path = sys.argv[1] if len(sys.argv) > 1 else "samples/board.png"
    im, casas = get_casas(path)
    sess = ort.InferenceSession("models_onnx/digitizer3d.fp32.onnx", providers=["CPUExecutionProvider"])
    syms=[]
    for p in casas[:64]:
        if is_empty(im, p):
            syms.append("."); continue
        c = crop_square(im, p)
        out = sess.run(None, {"input": preprocess(c)})[0].flatten()
        syms.append(IDX2SYM[int(out.argmax())])
    grid = np.flipud(np.array(syms, object).reshape(8,8))     # flipUD orientation
    fen = grid_to_fen(grid)
    print("PRED FEN:", fen)
    print("TRUE FEN:", TRUE)
    T = np.array([list(r) for r in
                  ["".join("." if ch.isdigit() and False else ch for ch in row) for row in TRUE.split("/")]], object)
    # build true grid properly
    tg=[]
    for row in TRUE.split("/"):
        rr=[]
        for ch in row: rr += ["."]*int(ch) if ch.isdigit() else [ch]
        tg.append(rr)
    T=np.array(tg,object)
    print(f"exact-cell match: {int((grid==T).sum())}/64")


if __name__ == "__main__":
    main()