trigo-web/tests/mctsTerminalPropagation.test.ts

/**
 * Unit Tests for MCTS Terminal State Propagation
 *
 * Tests the minimax propagation logic where nodes with all terminal children
 * are themselves marked as terminal.
 *
 * Based on GPT-5.1 review recommendations.
 */

import { describe, it, expect, beforeEach } from "vitest";
import { TrigoGame } from "../inc/trigo/game.js";
import { MCTSAgent } from "../inc/mctsAgent.js";
import { TrigoTreeAgent } from "../inc/trigoTreeAgent.js";
import { TrigoEvaluationAgent } from "../inc/trigoEvaluationAgent.js";
import { ModelInferencer } from "../inc/modelInferencer.js";
import type { Move } from "../inc/trigo/types.js";


/**
 * Mock Tree Agent that returns uniform priors
 */
class MockTreeAgent extends TrigoTreeAgent {
	constructor() {
		super(null as any);
	}

	async scoreMoves(game: TrigoGame, moves: Move[]): Promise<Array<{ move: Move; score: number }>> {
		// Return uniform log probabilities
		const logProb = Math.log(1.0 / moves.length);
		return moves.map(move => ({ move, score: logProb }));
	}
}


/**
 * Mock Evaluation Agent that returns fixed values
 */
class MockEvaluationAgent extends TrigoEvaluationAgent {
	private mockValue: number;

	constructor(value: number = 0.0) {
		super(null as any);
		this.mockValue = value;
	}

	setMockValue(value: number): void {
		this.mockValue = value;
	}

	async evaluatePosition(game: TrigoGame): Promise<{ value: number; interpretation: string }> {
		return {
			value: this.mockValue,
			interpretation: this.mockValue > 0 ? "White advantage" : "Black advantage"
		};
	}
}


/**
 * Helper to access private MCTSNode properties via type assertion
 */
interface MCTSNode {
	state: TrigoGame | null;
	parent: MCTSNode | null;
	action: Move | null;
	N: Map<string, number>;
	W: Map<string, number>;
	Q: Map<string, number>;
	P: Map<string, number>;
	children: Map<string, MCTSNode>;
	expanded: boolean;
	terminalValue: number | null;
	depth: number;
	playerToMove: number;
}


describe("MCTS Terminal Propagation", () => {
	let game: TrigoGame;
	let treeAgent: MockTreeAgent;
	let evalAgent: MockEvaluationAgent;
	let mctsAgent: MCTSAgent;

	beforeEach(() => {
		treeAgent = new MockTreeAgent();
		evalAgent = new MockEvaluationAgent(0.0);
		mctsAgent = new MCTSAgent(treeAgent, evalAgent, {
			numSimulations: 10,
			cPuct: 1.0,
			temperature: 1.0,
			dirichletAlpha: 0.03,
			dirichletEpsilon: 0.0  // Disable Dirichlet noise for testing
		});
	});


	describe("Test 1: Single-Step Endgame", () => {
		it("should propagate terminal value when Black chooses between terminal children", async () => {
			// Setup: 5x1x1 board with 2 empty positions (a and b)
			// Black to move, can play at position 0 (index 2) or position 1 (index 1)
			game = new TrigoGame({ x: 5, y: 1, z: 1 });
			game.startGame();

			// Simulate near-terminal game: Black at z(4), y(3); White at a(0), b(1)
			// Remaining: position 2 (center '0')
			game.drop({ x: 4, y: 0, z: 0 });  // Black at z
			game.drop({ x: 0, y: 0, z: 0 });  // White at a
			game.drop({ x: 3, y: 0, z: 0 });  // Black at y
			game.drop({ x: 1, y: 0, z: 0 });  // White at b

			// Now Black to move at position 2 (center)
			// This will end the game
			const result = await mctsAgent.selectMove(game, 5);

			// Verify move was selected
			expect(result.move).toBeDefined();

			// Check that terminal propagation happened
			// (Terminal value should be set on nodes)
			expect(result.rootValue).toBeDefined();
		});
	});


	describe("Test 2: Two-Step Propagation", () => {
		it("should propagate terminal values up two levels", async () => {
			// Setup: Near-terminal position
			game = new TrigoGame({ x: 3, y: 1, z: 1 });
			game.startGame();

			// Black at position 0, White at position 2
			// Remaining: position 1
			game.drop({ x: 0, y: 0, z: 0 });  // Black
			game.drop({ x: 2, y: 0, z: 0 });  // White

			// Now Black to move - only one move possible
			const result = await mctsAgent.selectMove(game, 3);

			// After enough simulations, should mark nodes as terminal
			expect(result.move).toBeDefined();
		});
	});


	describe("Test 3: White vs Black Turn Correctness", () => {
		it("should use max for White's turn and min for Black's turn", async () => {
			// This is a conceptual test - we verify via the implementation
			// White should maximize Q-values (white-positive)
			// Black should minimize Q-values (white-positive)

			// Setup simple 3x1x1 board
			game = new TrigoGame({ x: 3, y: 1, z: 1 });
			game.startGame();

			// Run a few simulations
			const result = await mctsAgent.selectMove(game, 1);

			// Verify basic functionality works
			expect(result.move).toBeDefined();
			expect(result.visitCounts.size).toBeGreaterThan(0);
		});
	});


	describe("Test 4: Terminal Leaf with No Actions", () => {
		it("should handle terminal nodes with empty action sets", async () => {
			// Create a finished game
			game = new TrigoGame({ x: 3, y: 1, z: 1 });
			game.startGame();

			// Fill the board completely
			game.drop({ x: 0, y: 0, z: 0 });  // Black
			game.drop({ x: 1, y: 0, z: 0 });  // White
			game.drop({ x: 2, y: 0, z: 0 });  // Black

			// Game is now terminal (no more moves)
			// Double-pass to finish
			game.pass();
			game.pass();

			// Should return immediately without error
			const result = await mctsAgent.selectMove(game, 4);

			expect(result.move.isPass).toBe(true);
			expect(result.rootValue).toBeDefined();
		});
	});


	describe("Test 5: Selection Skips Terminal Nodes", () => {
		it("should not expand terminal nodes during selection", async () => {
			// Setup near-terminal position
			game = new TrigoGame({ x: 5, y: 1, z: 1 });
			game.startGame();

			// Play several moves to approach terminal state
			game.drop({ x: 0, y: 0, z: 0 });  // Black at a
			game.drop({ x: 4, y: 0, z: 0 });  // White at z
			game.drop({ x: 1, y: 0, z: 0 });  // Black at b
			game.drop({ x: 3, y: 0, z: 0 });  // White at y

			// One position left (center)
			const result = await mctsAgent.selectMove(game, 5);

			expect(result.move).toBeDefined();

			// Should have visit counts for available moves
			expect(result.visitCounts.size).toBeGreaterThan(0);
		});
	});


	describe("Test 6: White-Positive Minimax Verification", () => {
		it("should correctly apply minimax with white-positive values", async () => {
			// Test scenario from verification doc:
			// White to move with children [+5, +1, -2] should choose +5
			// Black to move with children [+5, +1, -2] should choose -2

			game = new TrigoGame({ x: 3, y: 1, z: 1 });
			game.startGame();

			// Run MCTS
			const result = await mctsAgent.selectMove(game, 1);

			// Basic verification
			expect(result.move).toBeDefined();
			expect(result.rootValue).toBeDefined();
		});
	});


	describe("Test 7: 5x1x1 Board Complete Game", () => {
		it("should handle 5x1x1 board terminal propagation correctly", async () => {
			// Test terminal propagation on a 5x1x1 board
			// Focus on verifying the mechanics work, not exact game outcome

			game = new TrigoGame({ x: 5, y: 1, z: 1 });
			game.startGame();

			// Play a few moves
			game.drop({ x: 2, y: 0, z: 0 });  // Black center
			game.drop({ x: 0, y: 0, z: 0 });  // White a
			game.drop({ x: 4, y: 0, z: 0 });  // Black z

			// Now run MCTS to find best move for White
			// Should handle near-terminal position correctly
			const result = await mctsAgent.selectMove(game, 4);

			expect(result.move).toBeDefined();
			expect(result.visitCounts.size).toBeGreaterThan(0);
			expect(result.rootValue).toBeDefined();

			// The key test: terminal propagation should work without errors
			// Even if we don't have terminal nodes yet, the mechanism should be ready
		});
	});
});