document.addEventListener('DOMContentLoaded', () => {
    const GRID_SIZE = 50;
    const CELL_SIZE = 10; // Adjust for desired visual size
    const CANVAS_WIDTH = GRID_SIZE * CELL_SIZE;
    const CANVAS_HEIGHT = GRID_SIZE * CELL_SIZE;

    const canvas = document.getElementById('gridCanvas');
    const ctx = canvas.getContext('2d');
    const input = document.getElementById('kInput');
    const runButton = document.getElementById('runButton');
    const output = document.getElementById('output');

    canvas.width = CANVAS_WIDTH;
    canvas.height = CANVAS_HEIGHT;

    // --- Grid State ---
    let G = Array(GRID_SIZE).fill().map(() => Array(GRID_SIZE).fill(0)); // 2D matrix grid state

    // --- Drawing ---
    function drawGridLines() {
        ctx.strokeStyle = '#eee'; // Light gray grid lines
        ctx.lineWidth = 1;

        for (let i = 0; i <= GRID_SIZE; i++) {
            // Vertical lines
            ctx.beginPath();
            ctx.moveTo(i * CELL_SIZE + 0.5, 0);
            ctx.lineTo(i * CELL_SIZE + 0.5, CANVAS_HEIGHT);
            ctx.stroke();

            // Horizontal lines
            ctx.beginPath();
            ctx.moveTo(0, i * CELL_SIZE + 0.5);
            ctx.lineTo(CANVAS_WIDTH, i * CELL_SIZE + 0.5);
            ctx.stroke();
        }
    }

    function drawCells() {
        ctx.fillStyle = '#333'; // Color for 'on' cells
        for (let row = 0; row < GRID_SIZE; row++) {
            for (let col = 0; col < GRID_SIZE; col++) {
                if (G[row][col] === 1) {
                    ctx.fillRect(col * CELL_SIZE, row * CELL_SIZE, CELL_SIZE, CELL_SIZE);
                }
            }
        }
    }

    function redraw() {
        // Clear canvas
        ctx.fillStyle = '#fff'; // Background color
        ctx.fillRect(0, 0, CANVAS_WIDTH, CANVAS_HEIGHT);

        drawGridLines();
        drawCells();
    }

    // Helper functions for array operations
    const range = (n) => Array.from({length: n}, (_, i) => i);
    const reshape = (arr, rows, cols) => {
        const result = [];
        for (let i = 0; i < rows; i++) {
            result.push(arr.slice(i * cols, (i + 1) * cols));
        }
        return result;
    };
    const ravel = (arr) => arr.flat();
    const zip = (a, b) => a.map((x, i) => [x, b[i]]);

    // --- K-like Interpreter ---
    function tokenize(code) {
        // First, normalize whitespace
        code = code.replace(/\s+/g, ' ').trim();
        
        // Define all operators and special characters
        const operators = ['+', '-', '*', '/', '%', '(', ')', '!', '@', ':', '=', '<', '>', "'", '|', '$', '#', '~'];
        
        // Add spaces around all operators
        operators.forEach(op => {
            // Use a regex that ensures we don't double-space
            const regex = new RegExp(`\\${op}`, 'g');
            code = code.replace(regex, ` ${op} `);
        });
        
        // Normalize spaces again
        code = code.replace(/\s+/g, ' ').trim();
        
        // Split into tokens
        const tokens = code.split(' ');
        
        // Filter out empty tokens and log for debugging
        const filteredTokens = tokens.filter(t => t.length > 0);
        console.log('Tokenized:', filteredTokens);
        
        return filteredTokens;
    }

    function evaluateExpression(tokens) {
        if (!tokens || tokens.length === 0) throw new Error("Empty expression");

        function parseAtom() {
            let token = tokens.shift();
            if (!token) throw new Error("Unexpected end of expression");

            if (token === '(') {
                const value = parseAddSub();
                if (tokens.length === 0) {
                    throw new Error("Missing closing parenthesis");
                }
                const nextToken = tokens.shift();
                if (nextToken !== ')') {
                    throw new Error(`Expected closing parenthesis, got: ${nextToken}`);
                }
                return value;
            } else if (token === '!') { // Iota (prefix)
                const operand = parseAtom();
                if (typeof operand !== 'number' || !Number.isInteger(operand) || operand < 0) {
                    throw new Error("Operand for ! (iota) must be a non-negative integer");
                }
                const maxIndex = GRID_SIZE * GRID_SIZE;
                const result = Array.from({length: Math.min(operand, maxIndex)}, (_, i) => i);
                console.log(`Iota generated array of length ${result.length}, first few values:`, result.slice(0, 5));
                return result;
            } else if (token === '~') { // Not operator
                const operand = parseAtom();
                if (Array.isArray(operand)) {
                    const result = operand.map(x => {
                        const val = x === 0 ? 1 : 0;
                        console.log(`Not operation: ${x} -> ${val}`);
                        return val;
                    });
                    console.log(`Not operation on array, first few results:`, result.slice(0, 5));
                    console.log(`Input array first few values:`, operand.slice(0, 5));
                    return result;
                } else {
                    const result = operand === 0 ? 1 : 0;
                    console.log(`Not operation (scalar): ${operand} -> ${result}`);
                    return result;
                }
            } else if (/^-?\d+$/.test(token)) {
                return parseInt(token, 10);
            } else {
                throw new Error(`Unrecognized token: ${token}`);
            }
        }

        function parseUnary() {
            let token = tokens[0];
            if (token === '|' || token === '$' || token === '#') {
                tokens.shift();
                const operand = parseUnary();
                let result;
                switch (token) {
                    case '|':
                        result = [...operand].reverse();
                        break;
                    case '$':
                        if (!Array.isArray(operand)) {
                            throw new Error("Rotate operator ($) requires an array operand");
                        }
                        const size = Math.sqrt(operand.length);
                        if (size * size !== operand.length) {
                            throw new Error("Rotate operator ($) requires a square array");
                        }
                        result = [];
                        for (let col = 0; col < size; col++) {
                            for (let row = size - 1; row >= 0; row--) {
                                result.push(operand[row * size + col]);
                            }
                        }
                        break;
                    case '#':
                        if (!Array.isArray(operand)) {
                            throw new Error("Reshape operator (#) requires an array operand");
                        }
                        result = [];
                        for (let i = 0; i < GRID_SIZE; i++) {
                            result.push(operand.slice(i * GRID_SIZE, (i + 1) * GRID_SIZE));
                        }
                        result = result.flat();
                        break;
                }
                return result;
            }
            return parseAtom();
        }

        function applyOperation(a, b, op) {
            const isAList = Array.isArray(a);
            const isBList = Array.isArray(b);

            const scalarOp = (x, y) => {
                let result;
                switch (op) {
                    case '+': result = x + y; break;
                    case '-': result = x - y; break;
                    case '*': result = x * y; break;
                    case '%': result = y === 0 ? 0 : x % y; break;
                    case '/': result = y === 0 ? 0 : Math.floor(x / y); break;
                    case '<': result = x < y ? 1 : 0; break;
                    case '>': result = x > y ? 1 : 0; break;
                    case '=': result = x === y ? 1 : 0; break;
                    default: throw new Error(`Unknown operator: ${op}`);
                }
                return result;
            };

            // Handle scalar operations
            if (!isAList && !isBList) {
                return scalarOp(a, b);
            }

            // Handle array operations
            const arrayOp = (arr, val) => {
                if (Array.isArray(arr)) {
                    const result = arr.map(x => arrayOp(x, val));
                    if (op === '/' || op === '%') {
                        console.log(`Array operation ${op} with ${val}, first few results:`, result.slice(0, 5));
                    }
                    return result;
                }
                return scalarOp(arr, val);
            };

            if (isAList && !isBList) {
                const result = arrayOp(a, b);
                if (op === '+') {
                    console.log(`Array + scalar operation, first few values:`, {
                        array: a.slice(0, 5),
                        scalar: b,
                        result: result.slice(0, 5)
                    });
                }
                return result;
            } else if (!isAList && isBList) {
                const result = arrayOp(b, a);
                if (op === '+') {
                    console.log(`Scalar + array operation, first few values:`, {
                        scalar: a,
                        array: b.slice(0, 5),
                        result: result.slice(0, 5)
                    });
                }
                return result;
            } else {
                // Both are arrays
                if (a.length !== b.length) {
                    throw new Error(`List length mismatch for operator ${op}: ${a.length} vs ${b.length}`);
                }
                const result = a.map((x, i) => {
                    const val = arrayOp(x, b[i]);
                    if (op === '+') {
                        console.log(`Adding values at index ${i}: ${x} + ${b[i]} = ${val}`);
                    }
                    return val;
                });
                if (op === '+') {
                    console.log(`Array addition, first few results:`, result.slice(0, 5));
                }
                return result;
            }
        }

        function parseMulDivMod() {
            let left = parseUnary();
            while (tokens.length > 0 && (tokens[0] === '*' || tokens[0] === '%' || tokens[0] === '/')) {
                const op = tokens.shift();
                const right = parseUnary();
                left = applyOperation(left, right, op);
            }
            return left;
        }

        function parseAddSub() {
            let left = parseMulDivMod();
            while (tokens.length > 0 && (tokens[0] === '+' || tokens[0] === '-')) {
                const op = tokens.shift();
                const right = parseMulDivMod();
                left = applyOperation(left, right, op);
            }
            return left;
        }

        function parseComparison() {
            let left = parseAddSub();
            while (tokens.length > 0 && (tokens[0] === '<' || tokens[0] === '>' || tokens[0] === '=')) {
                const op = tokens.shift();
                const right = parseAddSub();
                left = applyOperation(left, right, op);
            }
            return left;
        }

        function parseNot() {
            let left = parseComparison();
            while (tokens.length > 0 && tokens[0] === '~') {
                tokens.shift();
                left = Array.isArray(left) ? left.map(x => x === 0 ? 1 : 0) : (left === 0 ? 1 : 0);
            }
            return left;
        }

        return parseNot();
    }

    // Main execution function
    function executeK(code) {
        code = code.trim();
        if (!code) return;

        try {
            if (code === 'G : 0') {
                G = Array(GRID_SIZE).fill().map(() => Array(GRID_SIZE).fill(0));
                setOutput("Grid cleared.", "success");
                return;
            }
            if (code === 'G : 1') {
                G = Array(GRID_SIZE).fill().map(() => Array(GRID_SIZE).fill(1));
                setOutput("Grid filled.", "success");
                return;
            }

            const assignMatch = code.match(/^G\s*@\s*(.+?)\s*:\s*(.+)$/);
            if (assignMatch) {
                const indexExpr = assignMatch[1].trim();
                const valueExpr = assignMatch[2].trim();
                const steps = [];

                // Parse and evaluate indices
                steps.push(`1. Evaluating indices expression: ${indexExpr}`);
                const indices = evaluateExpression(tokenize(indexExpr));
                steps.push(`   → Generated ${indices.length} indices`);

                // Parse and evaluate values
                steps.push(`2. Evaluating values expression: ${valueExpr}`);
                const values = evaluateExpression(tokenize(valueExpr));
                steps.push(`   → Generated ${Array.isArray(values) ? values.length : 1} values`);

                const indicesArray = Array.isArray(indices) ? indices : [indices];
                const valuesArray = Array.isArray(values) ? values : [values];

                if (indicesArray.length === 0) {
                    setOutput("Warning: Assignment applied to empty index list.", "info", steps);
                    return;
                }

                // Vectorized assignment
                const assignments = indicesArray.reduce((count, idx, i) => {
                    const row = Math.floor(idx / GRID_SIZE);
                    const col = idx % GRID_SIZE;
                    
                    if (row >= 0 && row < GRID_SIZE && col >= 0 && col < GRID_SIZE) {
                        const valueToAssign = valuesArray[i % valuesArray.length];
                        if (count < 5) { // Only log first few assignments
                            console.log(`Assignment [${row},${col}]: ${valueToAssign} (from index ${idx})`);
                        }
                        G[row][col] = valueToAssign % 2;
                        return count + 1;
                    }
                    return count;
                }, 0);

                steps.push(`3. Assignment complete:`);
                steps.push(`   → Applied ${assignments} assignments to the grid`);
                steps.push(`   → Each value was taken modulo 2 to ensure binary (0/1) values`);

                setOutput(`OK. Performed ${assignments} assignments.`, "success", steps);
            } else {
                const result = evaluateExpression(tokenize(code));
                setOutput(`Evaluated: ${JSON.stringify(result)}`, "info", [
                    `1. Evaluated expression: ${code}`,
                    `2. Result: ${JSON.stringify(result)}`
                ]);
            }
        } catch (error) {
            setOutput(`Error: ${error.message}`, "error", [
                `1. Error occurred while executing: ${code}`,
                `2. Error details: ${error.message}`
            ]);
            console.error("K execution error:", error);
        }
    }

    // --- Output Helper ---
    function setOutput(message, type = "info", steps = []) {
        const outputDiv = document.getElementById('output');
        
        // Create a container for the message and steps
        const container = document.createElement('div');
        container.className = type;
        
        // Add the main message
        const messageDiv = document.createElement('div');
        messageDiv.textContent = message;
        container.appendChild(messageDiv);
        
        // Add steps if provided
        if (steps.length > 0) {
            const stepsDiv = document.createElement('div');
            stepsDiv.className = 'steps';
            steps.forEach(step => {
                const stepDiv = document.createElement('div');
                stepDiv.className = 'step';
                stepDiv.textContent = step;
                stepsDiv.appendChild(stepDiv);
            });
            container.appendChild(stepsDiv);
        }
        
        // Clear previous output and add new content
        outputDiv.innerHTML = '';
        outputDiv.appendChild(container);
    }

    // --- Event Listeners ---
    function handleRun() {
         const code = input.value;
         executeK(code);
         redraw();
         // Optional: Clear input after running
         // input.value = '';
     }

    input.addEventListener('keydown', (event) => {
        if (event.key === 'Enter') {
            handleRun();
        }
    });

    runButton.addEventListener('click', handleRun);


    // --- Initial Draw ---
    setOutput("Grid initialized. Enter commands below.", "info");
    redraw();
});