% Reinforcement Learning: Policy Iteration on a Grid World
%
% Demonstrates the Bellman equation, policy evaluation, and policy
% improvement on a grid world with obstacles, a goal, and traps.
%
% The agent starts at the top-left and must reach the goal (bottom-right)
% while avoiding traps. Actions: up, down, left, right.
%
% Animation shows:
%   1. Value function V(s) as a heatmap evolving during policy evaluation
%   2. Policy pi(s) as arrows updating during policy improvement
%   3. The agent following the learned policy after convergence

function rl_gridworld()
    %% =====================================================================
    %  1. Define the grid world
    %  =====================================================================
    % Grid: 0 = free, 1 = obstacle, 2 = trap, 3 = goal
    grid = [
        0  0  0  0  0  0  0  0
        0  1  1  0  0  1  0  0
        0  0  0  0  1  0  0  2
        0  1  0  0  0  0  1  0
        0  1  0  1  1  0  0  0
        0  0  0  0  1  0  1  0
        2  0  1  0  0  0  0  0
        0  0  0  0  1  0  0  3
    ];

    [nRows, nCols] = size(grid);
    nStates = nRows * nCols;

    % Actions: 1=up, 2=down, 3=left, 4=right
    nActions = 4;
    actionNames = {'Up','Down','Left','Right'};
    % Row, col deltas for each action
    dRow = [-1, 1, 0, 0];
    dCol = [0, 0, -1, 1];
    % Arrow symbols for plotting
    arrowDx = [0, 0, -0.3, 0.3];
    arrowDy = [0.3, -0.3, 0, 0];

    % Rewards
    R_step = -0.1;    % small penalty per step (encourages short paths)
    R_goal = 10.0;    % reaching the goal
    R_trap = -5.0;    % falling into a trap
    R_wall = -0.5;    % bumping into wall/obstacle

    % Discount factor
    gamma = 0.95;

    % State indexing helpers
    sub2s = @(r,c) (r-1)*nCols + c;
    s2row = @(s) ceil(s / nCols);
    s2col = @(s) s - (s2row(s)-1)*nCols;

    % Identify terminal states
    goalState = find(grid(:) == 3);
    trapStates = find(grid(:) == 2);
    obstacleStates = find(grid(:) == 1);
    terminalStates = [goalState; trapStates];

    %% =====================================================================
    %  2. Build transition model P(s'|s,a) and R(s,a)
    %  =====================================================================
    % Deterministic transitions (can add stochasticity later)
    % nextState(s, a) and reward(s, a)
    nextState = zeros(nStates, nActions);
    reward    = zeros(nStates, nActions);

    for s = 1:nStates
        r = s2row(s); c = s2col(s);

        for a = 1:nActions
            if grid(r,c) == 1  % obstacle: can't be here
                nextState(s,a) = s;
                reward(s,a) = 0;
                continue;
            end
            if ismember(s, terminalStates)  % terminal: stay
                nextState(s,a) = s;
                reward(s,a) = 0;
                continue;
            end

            nr = r + dRow(a);
            nc = c + dCol(a);

            % Check bounds
            if nr < 1 || nr > nRows || nc < 1 || nc > nCols || grid(nr,nc) == 1
                % Hit wall or obstacle: stay in place
                nextState(s,a) = s;
                reward(s,a) = R_wall;
            else
                ns = sub2s(nr, nc);
                nextState(s,a) = ns;
                if grid(nr,nc) == 3
                    reward(s,a) = R_goal;
                elseif grid(nr,nc) == 2
                    reward(s,a) = R_trap;
                else
                    reward(s,a) = R_step;
                end
            end
        end
    end

    %% =====================================================================
    %  3. Initialize policy and value function
    %  =====================================================================
    % Random initial policy (uniform over valid actions)
    policy = randi(nActions, nStates, 1);
    V = zeros(nStates, 1);

    %% =====================================================================
    %  4. Set up figure
    %  =====================================================================
    fig = figure('Name','RL Grid World: Policy Iteration','Color','w',...
                 'Position',[80 60 1300 600]);

    % --- Left panel: value function heatmap ---
    ax1 = subplot(1,2,1,'Parent',fig);
    title(ax1, 'Value Function V(s)', 'FontSize', 14);

    % --- Right panel: policy arrows on grid ---
    ax2 = subplot(1,2,2,'Parent',fig);
    title(ax2, 'Policy \pi(s)', 'FontSize', 14);

    iterText = annotation(fig, 'textbox', [0.30 0.01 0.4 0.04], ...
        'String', 'Iteration: 0', 'FontSize', 13, 'EdgeColor', 'none', ...
        'HorizontalAlignment', 'center', 'FontWeight', 'bold');

    %% =====================================================================
    %  5. Policy Iteration with animation
    %  =====================================================================
    maxPolicyIter = 20;
    maxEvalIter   = 100;   % max sweeps per policy evaluation
    evalTol       = 1e-4;  % convergence tolerance for evaluation

    for pIter = 1:maxPolicyIter
        %% --- Policy Evaluation ---
        for evalStep = 1:maxEvalIter
            Vnew = V;
            for s = 1:nStates
                if ismember(s, terminalStates) || grid(s2row(s), s2col(s)) == 1
                    continue;
                end
                a = policy(s);
                ns = nextState(s, a);
                Vnew(s) = reward(s, a) + gamma * V(ns);
            end
            delta = max(abs(Vnew - V));
            V = Vnew;
            if delta < evalTol
                break;
            end
        end

        %% --- Animate value function and policy ---
        drawGrid(ax1, ax2, grid, V, policy, nRows, nCols, ...
                 arrowDx, arrowDy, obstacleStates, goalState, trapStates);
        set(iterText, 'String', sprintf(...
            'Policy Iteration: %d  |  Eval sweeps: %d  |  V range: [%.2f, %.2f]', ...
            pIter, evalStep, min(V(V~=0)), max(V)));
        drawnow;
        pause(0.5);

        %% --- Policy Improvement ---
        policyStable = true;
        for s = 1:nStates
            if ismember(s, terminalStates) || grid(s2row(s), s2col(s)) == 1
                continue;
            end
            % Greedy improvement: find best action
            qValues = zeros(nActions, 1);
            for a = 1:nActions
                ns = nextState(s, a);
                qValues(a) = reward(s, a) + gamma * V(ns);
            end
            [~, bestAction] = max(qValues);
            if bestAction ~= policy(s)
                policyStable = false;
                policy(s) = bestAction;
            end
        end

        % Redraw after improvement
        drawGrid(ax1, ax2, grid, V, policy, nRows, nCols, ...
                 arrowDx, arrowDy, obstacleStates, goalState, trapStates);
        drawnow;
        pause(0.3);

        if policyStable
            set(iterText, 'String', sprintf(...
                'CONVERGED at iteration %d  |  Policy is optimal', pIter));
            drawnow;
            break;
        end
    end

    %% =====================================================================
    %  6. Animate the agent following the optimal policy
    %  =====================================================================
    pause(1.0);

    % Start from top-left free cell
    startState = sub2s(1, 1);

    % Trace the path
    path = startState;
    s = startState;
    maxPathLen = nStates;  % safety limit
    for step = 1:maxPathLen
        if ismember(s, terminalStates)
            break;
        end
        a = policy(s);
        s = nextState(s, a);
        path(end+1) = s; %#ok<AGROW>
    end

    % Animate agent moving along the path
    axes(ax2); hold(ax2, 'on');
    agentDot = plot(ax2, NaN, NaN, 'o', 'MarkerSize', 20, ...
                    'MarkerFaceColor', [1 0.84 0], 'MarkerEdgeColor', 'k', ...
                    'LineWidth', 2);
    trailLine = animatedline(ax2, 'Color', [1 0.84 0], 'LineWidth', 3, ...
                             'LineStyle', '-');

    for step = 1:length(path)
        s = path(step);
        r = s2row(s); c = s2col(s);
        set(agentDot, 'XData', c, 'YData', nRows - r + 1);
        addpoints(trailLine, c, nRows - r + 1);

        if ismember(s, goalState)
            set(iterText, 'String', sprintf('GOAL reached in %d steps!', step-1));
        elseif ismember(s, trapStates)
            set(iterText, 'String', sprintf('Fell into TRAP at step %d!', step-1));
            set(agentDot, 'MarkerFaceColor', [1 0 0]);
        end

        drawnow;
        pause(0.4);
    end

    saveas(fig, fullfile(fileparts(mfilename('fullpath')), 'rl_gridworld.png'));

    %% =====================================================================
    %  7. Summary
    %  =====================================================================
    fprintf('\n=============================================\n');
    fprintf('  Grid World Policy Iteration Summary\n');
    fprintf('=============================================\n');
    fprintf('Grid size: %d x %d (%d states)\n', nRows, nCols, nStates);
    fprintf('Obstacles: %d,  Traps: %d,  Goal: 1\n', ...
            length(obstacleStates), length(trapStates));
    fprintf('Discount gamma: %.2f\n', gamma);
    fprintf('Policy converged at iteration: %d\n', pIter);
    fprintf('Optimal path length: %d steps\n', length(path)-1);
    fprintf('Path value V(start): %.4f\n', V(startState));
    fprintf('=============================================\n\n');
end

%% =========================================================================
%  Draw grid world: value heatmap + policy arrows
%  =========================================================================
function drawGrid(ax1, ax2, grid, V, policy, nRows, nCols, ...
                  arrowDx, arrowDy, obstacleStates, goalState, trapStates)

    s2row = @(s) ceil(s / nCols);
    s2col = @(s) s - (s2row(s)-1)*nCols;

    %% --- Value function heatmap ---
    cla(ax1);
    hold(ax1, 'on');

    Vgrid = reshape(V, nRows, nCols);
    % Flip for display (row 1 at top)
    VgridFlip = flipud(Vgrid);

    imagesc(ax1, VgridFlip);
    colormap(ax1, parula);
    colorbar(ax1);
    caxis(ax1, [min(V)-0.1, max(V)+0.1]);

    % Overlay cell values and markers
    for s = 1:nRows*nCols
        r = s2row(s); c = s2col(s);
        yr = nRows - r + 1;  % flipped row for display

        if grid(r,c) == 1
            % Obstacle: dark square
            fill(ax1, c + [-0.5 0.5 0.5 -0.5], yr + [-0.5 -0.5 0.5 0.5], ...
                 [0.2 0.2 0.2], 'EdgeColor', 'none');
        elseif grid(r,c) == 3
            % Goal: gold
            fill(ax1, c + [-0.5 0.5 0.5 -0.5], yr + [-0.5 -0.5 0.5 0.5], ...
                 [0.2 0.8 0.2], 'EdgeColor', 'none');
            text(ax1, c, yr, 'GOAL', 'HorizontalAlignment','center', ...
                 'FontWeight','bold','FontSize',8,'Color','w');
        elseif grid(r,c) == 2
            % Trap: red
            fill(ax1, c + [-0.5 0.5 0.5 -0.5], yr + [-0.5 -0.5 0.5 0.5], ...
                 [0.8 0.1 0.1], 'EdgeColor', 'none');
            text(ax1, c, yr, 'TRAP', 'HorizontalAlignment','center', ...
                 'FontWeight','bold','FontSize',8,'Color','w');
        else
            text(ax1, c, yr, sprintf('%.1f', V(s)), ...
                 'HorizontalAlignment','center','FontSize',7,'Color','w');
        end
    end

    % Grid lines
    for i = 0.5:(nCols+0.5)
        plot(ax1, [i i], [0.5 nRows+0.5], 'k-', 'LineWidth', 0.5);
    end
    for i = 0.5:(nRows+0.5)
        plot(ax1, [0.5 nCols+0.5], [i i], 'k-', 'LineWidth', 0.5);
    end

    axis(ax1, 'equal', 'tight');
    set(ax1, 'XTick', 1:nCols, 'YTick', 1:nRows, 'YTickLabel', nRows:-1:1);
    xlabel(ax1, 'Column'); ylabel(ax1, 'Row');
    title(ax1, 'Value Function V(s)', 'FontSize', 14);
    hold(ax1, 'off');

    %% --- Policy arrows ---
    cla(ax2);
    hold(ax2, 'on');

    % Background coloring
    for s = 1:nRows*nCols
        r = s2row(s); c = s2col(s);
        yr = nRows - r + 1;

        if grid(r,c) == 1
            fill(ax2, c + [-0.5 0.5 0.5 -0.5], yr + [-0.5 -0.5 0.5 0.5], ...
                 [0.2 0.2 0.2], 'EdgeColor', [0.3 0.3 0.3]);
        elseif grid(r,c) == 3
            fill(ax2, c + [-0.5 0.5 0.5 -0.5], yr + [-0.5 -0.5 0.5 0.5], ...
                 [0.2 0.8 0.2], 'EdgeColor', [0.3 0.3 0.3]);
            text(ax2, c, yr, 'GOAL', 'HorizontalAlignment','center', ...
                 'FontWeight','bold','FontSize',9,'Color','w');
        elseif grid(r,c) == 2
            fill(ax2, c + [-0.5 0.5 0.5 -0.5], yr + [-0.5 -0.5 0.5 0.5], ...
                 [0.8 0.1 0.1], 'EdgeColor', [0.3 0.3 0.3]);
            text(ax2, c, yr, 'TRAP', 'HorizontalAlignment','center', ...
                 'FontWeight','bold','FontSize',9,'Color','w');
        else
            fill(ax2, c + [-0.5 0.5 0.5 -0.5], yr + [-0.5 -0.5 0.5 0.5], ...
                 [0.95 0.95 0.95], 'EdgeColor', [0.7 0.7 0.7]);
            % Draw policy arrow
            a = policy(s);
            quiver(ax2, c, yr, arrowDx(a), arrowDy(a), 0, ...
                   'Color', [0 0.3 0.7], 'LineWidth', 2, ...
                   'MaxHeadSize', 0.8);
        end
    end

    % Start marker
    fill(ax2, 1 + [-0.3 0.3 0.3 -0.3], nRows + [-0.3 -0.3 0.3 0.3], ...
         [0 0.6 1], 'EdgeColor', 'k', 'FaceAlpha', 0.4);
    text(ax2, 1, nRows, 'S', 'HorizontalAlignment','center', ...
         'FontWeight','bold','FontSize',10,'Color','k');

    axis(ax2, 'equal', 'tight');
    set(ax2, 'XTick', 1:nCols, 'YTick', 1:nRows, 'YTickLabel', nRows:-1:1);
    xlabel(ax2, 'Column'); ylabel(ax2, 'Row');
    title(ax2, 'Policy \pi(s)', 'FontSize', 14);
    hold(ax2, 'off');
end
