#!/usr/bin/env python3
"""Move prediction inference (CNN part) replicating Predict_Human_Move_Train.test_model.
Sanity-checks the 12x8x8 encoding + class order on the start position."""
import sys
import numpy as np
import onnxruntime as ort
import chess

PIECE_ORDER = [chess.PAWN, chess.KNIGHT, chess.BISHOP, chess.ROOK, chess.QUEEN, chess.KING]
SQUARE_MODELS = ["pawn", "knight", "bishop", "rook", "queen", "king"]


def encode(board, row0_rank8=True):
    """12x8x8: channels 0-5 white P,N,B,R,Q,K; 6-11 black. (white-to-move board)"""
    enc = np.zeros((12, 8, 8), np.float32)
    for sq in chess.SQUARES:
        p = board.piece_at(sq)
        if p:
            c = PIECE_ORDER.index(p.piece_type) + (0 if p.color == chess.WHITE else 6)
            r = 7 - (sq // 8) if row0_rank8 else (sq // 8)
            enc[c, r, sq % 8] = 1.0
    return enc[None]


def softmax(x):
    e = np.exp(x - x.max()); return e / e.sum()


def cnn_moves(board, piece_sess, square_sesss, row0_rank8=True):
    enc = encode(board, row0_rank8)
    pieces = softmax(piece_sess.run(None, {"input": enc})[0].flatten())
    legal = list(board.legal_moves)
    move_prob = {}
    for i, ptype in enumerate(PIECE_ORDER):
        squares = softmax(square_sesss[i].run(None, {"input": enc})[0].flatten())
        squares = (squares + pieces[i]) / 2
        for from_sq in board.pieces(ptype, chess.WHITE):
            fs = chess.square_name(from_sq)
            for j in range(64):
                try:
                    mv = chess.Move.from_uci(fs + chess.square_name(j))
                    if mv in legal:
                        move_prob[mv.uci()] = squares[j]
                except Exception:
                    pass
    return [m for m, _ in sorted(move_prob.items(), key=lambda x: x[1], reverse=True)]


def main():
    piece_sess = ort.InferenceSession("models_onnx/piece.int8.onnx", providers=["CPUExecutionProvider"])
    square_sesss = [ort.InferenceSession(f"models_onnx/square_{s}.int8.onnx",
                                         providers=["CPUExecutionProvider"]) for s in SQUARE_MODELS]
    board = chess.Board()  # start position (white to move)
    for orient in (True, False):
        moves = cnn_moves(board, piece_sess, square_sesss, row0_rank8=orient)
        print(f"row0_rank8={orient}: top CNN moves -> {moves[:8]}")


if __name__ == "__main__":
    main()