#!/usr/bin/env python3
"""Diagnostic: classify the 64 squares of samples/board.png and compare to the
known FEN. Helps calibrate crop/normalization and understand the empty-class issue."""
import sys
import cv2 as cv
import numpy as np
import onnxruntime as ort

SIZE = 800
CORNERS = np.array([[520, 140], [1360, 150], [1430, 915], [452, 918]], dtype=np.float32)
TRUE_FEN = "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"}

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, (96, 96)), cv.COLOR_BGR2RGB).astype(np.float32) / 255.0
    rgb = (rgb - MEAN) / STD
    return rgb.transpose(2, 0, 1)[None]


def fen_grid(fen):
    grid = []
    for row in fen.split("/"):
        r = []
        for ch in row:
            if ch.isdigit():
                r += ["."] * int(ch)
            else:
                r.append(ch)
        grid.append(r)
    return grid


def main():
    margin = float(sys.argv[1]) if len(sys.argv) > 1 else 0.6  # top margin (cell fractions)
    img = cv.imread("samples/board.png")
    dst = np.array([[0, 0], [SIZE, 0], [SIZE, SIZE], [0, SIZE]], np.float32)
    Minv = cv.getPerspectiveTransform(dst, CORNERS)  # warped -> original

    sess = ort.InferenceSession("models_onnx/digitizer.int8.onnx",
                                providers=["CPUExecutionProvider"])
    cell = SIZE / 8
    pred = [["." for _ in range(8)] for _ in range(8)]
    prob = [[0.0 for _ in range(8)] for _ in range(8)]

    for r in range(8):
        for c in range(8):
            # cell quad in warped space, extended upward by `margin` cells for the piece
            y0 = (r - margin) * cell
            quad = np.array([[[c * cell, y0], [(c + 1) * cell, y0],
                              [(c + 1) * cell, (r + 1) * cell], [c * cell, (r + 1) * cell]]],
                            np.float32)
            orig = cv.perspectiveTransform(quad, Minv)[0]
            xs, ys = orig[:, 0], orig[:, 1]
            x1, x2 = max(0, int(xs.min())), min(img.shape[1], int(xs.max()))
            y1, y2 = max(0, int(ys.min())), min(img.shape[0], int(ys.max()))
            crop = img[y1:y2, x1:x2]
            if crop.size == 0:
                continue
            out = sess.run(None, {"input": preprocess(crop)})[0].flatten()
            p = np.exp(out) / np.exp(out).sum()
            k = int(p.argmax())
            pred[r][c] = IDX2SYM[k]
            prob[r][c] = float(p[k])

    true = fen_grid(TRUE_FEN)
    print(f"margin={margin}\n")
    print("PRED (top row = rank8 assumed):           TRUE:")
    hits = 0
    for r in range(8):
        pr = " ".join(pred[r])
        tr = " ".join(true[r])
        hits += sum(1 for a, b in zip(pred[r], true[r]) if a == b)
        print(f"  {pr}     {tr}")
    print(f"\nexact-cell match (incl. empties): {hits}/64")
    print("avg max-prob per row (low on empties?):")
    for r in range(8):
        print("  " + " ".join(f"{prob[r][c]:.2f}" for c in range(8)))


if __name__ == "__main__":
    main()