// rl_script.js

const grid = document.getElementById('grid');
const startButton = document.getElementById('startButton');
const stepButton = document.getElementById('stepButton');
const resetButton = document.getElementById('resetButton');
const learningGraphCanvas = document.getElementById('learningGraph');

const gridRows = 5;
const gridCols = 5;

const initialState = { row: 0, col: 0 };
const goalState = { row: 4, col: 4 };
const obstacles = [{ row: 1, col: 1 }, { row: 2, col: 2 }, { row: 3, col: 3 }];

let agentState = { ...initialState };
let gridCells = [];
let stepCount = 0;
let stepHistory = [];
let myChart = null; // Global variable to hold the chart instance

function createGrid() {
    for (let row = 0; row < gridRows; row++) {
        for (let col = 0; col < gridCols; col++) {
            const cell = document.createElement('div');
            cell.classList.add('cell');
            if (row === agentState.row && col === agentState.col) {
                cell.classList.add('agent');
            } else if (row === goalState.row && col === goalState.col) {
                cell.classList.add('goal');
            } else if (obstacles.some(o => o.row === row && o.col === col)) {
                cell.classList.add('obstacle');
            }
            grid.appendChild(cell);
            gridCells.push(cell);
        }
    }
}

function moveAgent(action) {
    const { row, col } = agentState;
    let newRow = row;
    let newCol = col;

    switch (action) {
        case 'up':
            newRow = Math.max(0, row - 1);
            break;
        case 'down':
            newRow = Math.min(gridRows - 1, row + 1);
            break;
        case 'left':
            newCol = Math.max(0, col - 1);
            break;
        case 'right':
            newCol = Math.min(gridCols - 1, col + 1);
            break;
    }

    if (!obstacles.some(o => o.row === newRow && o.col === newCol)) {
        agentState = { row: newRow, col: newCol };
        stepCount++;
        stepHistory.push(stepCount);
        updateLearningGraph();
    }

    updateGrid();
}

function updateGrid() {
    gridCells.forEach((cell, index) => {
        const row = Math.floor(index / gridCols);
        const col = index % gridCols;
        cell.classList.remove('agent', 'goal', 'obstacle');
        if (row === agentState.row && col === agentState.col) {
            cell.classList.add('agent');
        } else if (row === goalState.row && col === goalState.col) {
            cell.classList.add('goal');
        } else if (obstacles.some(o => o.row === row && o.col === col)) {
            cell.classList.add('obstacle');
        }
    });
}

function startRL() {
    agentState = { ...initialState };
    stepCount = 0;
    stepHistory = [];
    updateGrid();
    updateLearningGraph();
}

function stepRL() {
    const actions = ['up', 'down', 'left', 'right'];
    const action = actions[Math.floor(Math.random() * actions.length)];
    moveAgent(action);
}

function resetRL() {
    startRL();
}

function updateLearningGraph() {
    const ctx = learningGraphCanvas.getContext('2d');

    // Destroy the existing chart if it exists
    if (myChart) {
        myChart.destroy();
    }

    // Create a new chart
    myChart = new Chart(ctx, {
        type: 'line',
        data: {
            labels: Array.from({ length: stepHistory.length }, (_, i) => i + 1),
            datasets: [{
                label: 'Steps Taken',
                data: stepHistory,
                borderColor: 'rgba(75, 192, 192, 1)',
                fill: false
            }]
        },
        options: {
            scales: {
                y: {
                    beginAtZero: true
                }
            }
        }
    });
}

createGrid();
updateLearningGraph();

startButton.addEventListener('click', startRL);
stepButton.addEventListener('click', stepRL);
resetButton.addEventListener('click', resetRL);