#!/usr/bin/env python3
"""Faithful reproduction of Preprocess.ipynb crop pipeline on samples/board.png,
then classify the 64 squares and compare to the known FEN."""
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"}


# ---- EXACT functions from Preprocess.ipynb ----
def crop_image(imagem):
    cx, cy = imagem.shape[1] // 2, imagem.shape[0] // 2
    return imagem[cy - 520:cy + 450, cx - 550:cx + 550]

def detect_edges(im):       return cv.Canny(im, 100, 150, apertureSize=3)
def detect_lines(b):        return cv.HoughLines(b, 1, np.pi / 180, 160)

def group_lines(linhas):
    g = []
    for linha in linhas:
        rho, theta = linha[0]
        if not any(abs(rho - l[0][0]) < 30 and abs(theta - l[0][1]) < np.pi / 18 for l in g):
            g.append(linha)
    return g

def filter_board_lines(g):
    min_rho = None
    for linha in g:
        rho, theta = linha[0]
        if 1 < theta < 3 and (min_rho is None or rho < min_rho):
            min_rho = rho
    return [l for l in g if l[0][0] != min_rho]

def compute_intersections(linhas_casas):
    lp = []
    for linha in linhas_casas:
        rho, theta = linha[0]
        a, b = np.cos(theta), np.sin(theta)
        x0, y0 = a * rho, b * rho
        lp.append([(int(x0 + 10000 * -b), int(y0 + 10000 * a)),
                   (int(x0 - 10000 * -b), int(y0 - 10000 * a)), theta])
    rows = []
    for i in range(len(lp)):
        pts = []
        for j in range(len(lp)):
            x1, y1 = lp[i][0]; x2, y2 = lp[i][1]
            x3, y3 = lp[j][0]; x4, y4 = lp[j][1]
            det = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4)
            if det != 0:
                ix = int(((x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)) / det)
                iy = int(((x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)) / det)
                if 0 <= ix <= 1920 and 0 <= iy <= 1080 and 1 < lp[i][2] < 3:
                    pts.append((ix, iy))
        if pts:
            pts.sort(key=lambda p: p[0]); rows.append(pts)
    rows.sort(key=lambda p: p[0][0])
    return rows

def build_squares(rows):
    casas = []
    for i in range(len(rows) - 1):
        for j in range(len(rows[i]) - 1):
            casas.append([rows[i][j], rows[i][j + 1], rows[i + 1][j], rows[i + 1][j + 1]])
    return casas


def crop_square(imagem, p):
    xs = [pt[0] for pt in p]; ys = [pt[1] for pt in p]
    x_min, x_max = min(xs), max(xs); y_min, y_max = min(ys), max(ys)
    if y_min - 80 < 0:
        y_min = 80
    return imagem[y_min - 90:y_max, x_min:x_max]


def main():
    img = cv.imread("samples/board.png")
    im = crop_image(img)
    print("cropped:", im.shape)
    lines = detect_lines(detect_edges(im))
    print("hough lines:", 0 if lines is None else len(lines))
    g = group_lines(lines)
    casas = build_squares(compute_intersections(filter_board_lines(g)))
    print("squares detected:", len(casas))

    # montage of crops in detected order
    tiles = []
    for p in casas[:64]:
        c = crop_square(im, p)
        tiles.append(cv.resize(c, (90, 90)) if c.size else np.zeros((90, 90, 3), np.uint8))
    while len(tiles) < 64:
        tiles.append(np.zeros((90, 90, 3), np.uint8))
    mont = np.vstack([np.hstack(tiles[r * 8:r * 8 + 8]) for r in range(8)])
    cv.imwrite("/tmp/sq_montage.png", mont)
    print("saved /tmp/sq_montage.png")

    if len(casas) != 64:
        print("!! not 64 squares — pipeline needs tuning for this image"); return

    sess = ort.InferenceSession("models_onnx/digitizer.fp32.onnx", providers=["CPUExecutionProvider"])
    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])]:
        import chess
        board = chess.Board(None)
        for i, p in enumerate(casas[:64]):
            c = crop_square(im, p)
            if c.size == 0:
                continue
            rgb = cv.cvtColor(cv.resize(c, (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()
            board.set_piece_at(i, chess.Piece.from_symbol(IDX2SYM[int(out.argmax())]))
        print(f"\n[{nm}] FEN: {board.board_fen()}")
        print(f"      TRUE: {TRUE}")


if __name__ == "__main__":
    main()