package app.playmancala.bot;

import java.util.Arrays;

import app.playmancala.engine.Position;

// Scores every move of a position with negamax and alpha-beta pruning. Scores are from the mover's view in
// hundredths of a stone; a move that gives an extra turn keeps the side to move, so its score is not negated.
// Deepens one ply at a time up to the depth; when the node budget runs out, the scores of the last completed depth
// are returned. Nothing depends on time, so the same position, depth and budget always give the same scores.
public final class BotSearch {
    public static final long DEFAULT_NODE_BUDGET = 3_000_000;
    // An ended (or decided) game scores beyond any bank difference.
    public static final int WIN = 100_000;
    // A bank with more than half of the 48 stones has won.
    private static final int DECIDING_BANK = 24;
    private static final int STONE = 100;
    // Leaf score for each stone more in the mover's pits than in the opponent's.
    static final int PIT_STONE_WEIGHT = 0;

    public record Result(int[] pits, int[] scores, int depth, long nodes) {
        public int count() {
            return pits.length;
        }
    }

    private static final class OutOfBudget extends RuntimeException {
        OutOfBudget() {
            super(null, null, false, false);
        }
    }

    private final long nodeBudget;
    private Position[][] children;
    private int[][] order;
    private boolean[][] extraTurns;
    private int[][] orderKeys;
    private long nodes;
    private boolean budgetApplies;

    private BotSearch(long nodeBudget) {
        this.nodeBudget = nodeBudget;
    }

    public static Result scoreMoves(Position root, int maxDepth) {
        return scoreMoves(root, maxDepth, DEFAULT_NODE_BUDGET);
    }

    public static Result scoreMoves(Position root, int maxDepth, long nodeBudget) {
        return new BotSearch(nodeBudget).run(root, Math.max(1, maxDepth));
    }

    private Result run(Position root, int maxDepth) {
        children = new Position[maxDepth + 1][Position.PIT_COUNT];
        order = new int[maxDepth + 1][Position.PIT_COUNT];
        extraTurns = new boolean[maxDepth + 1][Position.PIT_COUNT];
        orderKeys = new int[maxDepth + 1][Position.PIT_COUNT];
        for (Position[] ply : children) {
            for (int i = 0; i < ply.length; i++) ply[i] = root.copy();
        }
        int[] pits = new int[Position.PIT_COUNT];
        int count = root.legalPits(pits);
        if (count == 0) throw new IllegalArgumentException("No moves in " + root);
        pits = Arrays.copyOf(pits, count);
        Position[] rootChildren = new Position[count];
        boolean[] rootExtraTurns = new boolean[count];
        for (int i = 0; i < count; i++) {
            rootChildren[i] = root.copy();
            rootExtraTurns[i] = root.play(pits[i], rootChildren[i]);
        }

        int[] scores = null;
        int completedDepth = 0;
        for (int depth = 1; depth <= maxDepth; depth++) {
            budgetApplies = depth > 1;
            int[] depthScores = new int[count];
            try {
                for (int i = 0; i < count; i++) {
                    nodes++;
                    int value = negamax(rootChildren[i], depth - 1, -Integer.MAX_VALUE, Integer.MAX_VALUE, 1);
                    depthScores[i] = rootExtraTurns[i] ? value : -value;
                }
            } catch (OutOfBudget e) {
                break;
            }
            scores = depthScores;
            completedDepth = depth;
        }
        return new Result(pits, scores, completedDepth, nodes);
    }

    private int negamax(Position position, int depth, int alpha, int beta, int ply) {
        if (budgetApplies && nodes > nodeBudget) throw new OutOfBudget();
        if (position.isOver()) return endScore(position.bankDifference());
        int ownBank = position.stonesAt(position.isPlayer1ToMove() ? Position.PLAYER1_BANK : Position.PLAYER2_BANK);
        int opponentBank = position.stonesAt(position.isPlayer1ToMove() ? Position.PLAYER2_BANK : Position.PLAYER1_BANK);
        if (ownBank > DECIDING_BANK || opponentBank > DECIDING_BANK) return endScore(ownBank - opponentBank);
        if (depth == 0) return STONE * (ownBank - opponentBank) + PIT_STONE_WEIGHT * position.pitStoneDifference();

        // Play every move once to order them: extra turns first, then by the stones they bring to the bank.
        Position[] next = children[ply];
        int[] moveOrder = order[ply];
        boolean[] extra = extraTurns[ply];
        int[] keys = orderKeys[ply];
        int count = 0;
        for (int pit = 0; pit < Position.PIT_COUNT; pit++) {
            if (position.stonesAt(position.boardIndex(pit)) == 0) continue;
            boolean extraTurn = position.play(pit, next[pit]);
            extra[pit] = extraTurn;
            int gain = next[pit].stonesAt(position.isPlayer1ToMove() ? Position.PLAYER1_BANK : Position.PLAYER2_BANK) - ownBank;
            int key = (extraTurn ? 1000 : 0) + gain;
            int at = count++;
            while (at > 0 && keys[at - 1] < key) {
                keys[at] = keys[at - 1];
                moveOrder[at] = moveOrder[at - 1];
                at--;
            }
            keys[at] = key;
            moveOrder[at] = pit;
        }

        int best = -Integer.MAX_VALUE;
        for (int i = 0; i < count; i++) {
            int pit = moveOrder[i];
            nodes++;
            int value = extra[pit]
                    ? negamax(next[pit], depth - 1, alpha, beta, ply + 1)
                    : -negamax(next[pit], depth - 1, -beta, -alpha, ply + 1);
            if (value > best) best = value;
            if (best > alpha) alpha = best;
            if (alpha >= beta) break;
        }
        return best;
    }

    private static int endScore(int bankDifference) {
        int sign = Integer.signum(bankDifference);
        return sign * WIN + STONE * bankDifference;
    }
}