/**
 * Gomoku RL Inference Module (Browser)
 * Loads ONNX model and provides inference for RL-based AI
 */

const GomokuRL = (() => {
    let session = null;
    let isLoaded = false;
    let isLoading = false;
    
    const BOARD_SIZE = 15;
    
    /**
     * Load ONNX model
     */
    async function loadModel(modelPath = 'models/gomoku_model.onnx') {
        if (isLoaded) return true;
        if (isLoading) {
            // Wait for current loading to complete
            while (isLoading) {
                await new Promise(resolve => setTimeout(resolve, 100));
            }
            return isLoaded;
        }
        
        isLoading = true;
        
        try {
            console.log('Loading ONNX model...');
            
            // Load ONNX.js if not already loaded
            if (typeof ort === 'undefined') {
                console.error('ONNX Runtime not found. Please include onnxruntime-web.min.js');
                isLoading = false;
                return false;
            }
            
            // Create inference session
            session = await ort.InferenceSession.create(modelPath);
            isLoaded = true;
            isLoading = false;
            
            console.log('✅ RL model loaded successfully');
            return true;
            
        } catch (error) {
            console.error('Failed to load RL model:', error);
            isLoading = false;
            isLoaded = false;
            return false;
        }
    }
    
    /**
     * Convert board to model input format
     * Returns 3-channel tensor: [current_player, opponent, player_indicator]
     */
    function boardToTensor(board, currentPlayer) {
        const WHITE = 1;
        const BLACK = 2;
        
        // Normalize current player to 1 or -1
        const player = (currentPlayer === 'white' || currentPlayer === WHITE) ? 1 : -1;
        const opponent = -player;
        
        // Create 3-channel input
        const channels = 3;
        const tensor = new Float32Array(channels * BOARD_SIZE * BOARD_SIZE);
        
        for (let r = 0; r < BOARD_SIZE; r++) {
            for (let c = 0; c < BOARD_SIZE; c++) {
                const cell = board[r][c];
                let cellValue = 0;
                
                // Convert cell to -1, 0, or 1
                if (cell === 'white' || cell === WHITE) cellValue = 1;
                else if (cell === 'black' || cell === BLACK) cellValue = -1;
                
                const idx = r * BOARD_SIZE + c;
                
                // Channel 0: current player pieces
                tensor[0 * BOARD_SIZE * BOARD_SIZE + idx] = (cellValue === player) ? 1 : 0;
                
                // Channel 1: opponent pieces
                tensor[channels * BOARD_SIZE * BOARD_SIZE + idx] = (cellValue === opponent) ? 1 : 0;
                
                // Channel 2: current player indicator
                tensor[2 * BOARD_SIZE * BOARD_SIZE + idx] = player;
            }
        }
        
        return tensor;
    }
    
    /**
     * Run model inference
     * Returns { policy: Float32Array, value: number }
     */
    async function predict(board, currentPlayer = 'white') {
        if (!isLoaded) {
            console.warn('RL model not loaded');
            return null;
        }
        
        try {
            // Prepare input
            const inputTensor = boardToTensor(board, currentPlayer);
            const dims = [1, 3, BOARD_SIZE, BOARD_SIZE];
            const input = new ort.Tensor('float32', inputTensor, dims);
            
            // Run inference
            const feeds = { input };
            const results = await session.run(feeds);
            
            // Extract outputs
            const policyLogits = results.policy.data;  // Log probabilities
            const value = results.value.data[0];       // Position value
            
            // Convert log probabilities to probabilities
            const policy = new Float32Array(policyLogits.length);
            let maxLogit = -Infinity;
            for (let i = 0; i < policyLogits.length; i++) {
                maxLogit = Math.max(maxLogit, policyLogits[i]);
            }
            let sum = 0;
            for (let i = 0; i < policyLogits.length; i++) {
                policy[i] = Math.exp(policyLogits[i] - maxLogit);
                sum += policy[i];
            }
            for (let i = 0; i < policyLogits.length; i++) {
                policy[i] /= sum;
            }
            
            return { policy, value };
            
        } catch (error) {
            console.error('Inference error:', error);
            return null;
        }
    }
    
    /**
     * Get best move using RL policy
     */
    async function getBestMove(board, currentPlayer = 'white', validMoves = null) {
        const result = await predict(board, currentPlayer);
        if (!result) return null;
        
        const { policy } = result;
        
        // Get valid moves
        if (!validMoves) {
            validMoves = [];
            for (let r = 0; r < BOARD_SIZE; r++) {
                for (let c = 0; c < BOARD_SIZE; c++) {
                    if (!board[r][c] || board[r][c] === 0) {
                        validMoves.push({ row: r, col: c });
                    }
                }
            }
        }
        
        // Find valid move with highest probability
        let bestMove = null;
        let bestProb = -1;
        
        for (const move of validMoves) {
            const idx = move.row * BOARD_SIZE + move.col;
            const prob = policy[idx];
            if (prob > bestProb) {
                bestProb = prob;
                bestMove = move;
            }
        }
        
        return bestMove;
    }
    
    /**
     * Evaluate position using RL value network
     */
    async function evaluatePosition(board, currentPlayer = 'white') {
        const result = await predict(board, currentPlayer);
        return result ? result.value : 0;
    }
    
    /**
     * Get top N moves with probabilities
     */
    async function getTopMoves(board, currentPlayer = 'white', topN = 5) {
        const result = await predict(board, currentPlayer);
        if (!result) return [];
        
        const { policy } = result;
        
        // Get valid moves with probabilities
        const moves = [];
        for (let r = 0; r < BOARD_SIZE; r++) {
            for (let c = 0; c < BOARD_SIZE; c++) {
                if (!board[r][c] || board[r][c] === 0) {
                    const idx = r * BOARD_SIZE + c;
                    moves.push({
                        row: r,
                        col: c,
                        probability: policy[idx]
                    });
                }
            }
        }
        
        // Sort by probability and return top N
        moves.sort((a, b) => b.probability - a.probability);
        return moves.slice(0, topN);
    }
    
    // Public API
    return {
        loadModel,
        isLoaded: () => isLoaded,
        getBestMove,
        evaluatePosition,
        getTopMoves,
        predict
    };
})();

// Export for module systems
if (typeof module !== 'undefined' && module.exports) {
    module.exports = GomokuRL;
}
