% Deep Q-Network (DQN) on Grid World
%
% Same grid world as rl_gridworld.m, but instead of tabular policy
% iteration, we approximate Q(s,a) with a neural network trained via
% the Bellman equation using Adam optimizer and experience replay.
%
% This demonstrates why neural network function approximation is useful:
% the agent learns from state features (row, col) rather than a lookup
% table, so in principle it could generalize to unseen grid positions.
%
% Comparison: at the end, we compare the DQN-learned policy to the
% exact tabular solution from policy iteration.

function rl_gridworld_dqn()
    %% =====================================================================
    %  1. Define the grid world (same as rl_gridworld.m)
    %  =====================================================================
    % 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;
    nActions = 4;  % 1=up, 2=down, 3=left, 4=right
    dRow = [-1, 1, 0, 0];
    dCol = [0, 0, -1, 1];
    arrowDx = [0, 0, -0.3, 0.3];
    arrowDy = [0.3, -0.3, 0, 0];

    R_step = -0.1;
    R_goal = 10.0;
    R_trap = -5.0;
    R_wall = -0.5;
    gamma  = 0.95;

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

    goalState = find(grid(:) == 3);
    trapStates = find(grid(:) == 2);
    obstacleStates = find(grid(:) == 1);
    terminalStates = [goalState; trapStates];
    freeStates = find(grid(:) == 0);

    % Build transition model (same as tabular version)
    nextStateTbl = zeros(nStates, nActions);
    rewardTbl    = zeros(nStates, nActions);
    for s = 1:nStates
        r = s2row(s); c = s2col(s);
        for a = 1:nActions
            if grid(r,c) == 1
                nextStateTbl(s,a) = s; rewardTbl(s,a) = 0; continue;
            end
            if ismember(s, terminalStates)
                nextStateTbl(s,a) = s; rewardTbl(s,a) = 0; continue;
            end
            nr = r + dRow(a); nc = c + dCol(a);
            if nr<1 || nr>nRows || nc<1 || nc>nCols || grid(nr,nc)==1
                nextStateTbl(s,a) = s; rewardTbl(s,a) = R_wall;
            else
                ns = sub2s(nr, nc);
                nextStateTbl(s,a) = ns;
                if grid(nr,nc)==3, rewardTbl(s,a) = R_goal;
                elseif grid(nr,nc)==2, rewardTbl(s,a) = R_trap;
                else, rewardTbl(s,a) = R_step;
                end
            end
        end
    end

    %% =====================================================================
    %  2. State encoding: (row, col) -> feature vector
    %  =====================================================================
    % Encode state as normalized [row/nRows, col/nCols] plus distance
    % features to give the network useful spatial information
    encodeState = @(s) stateFeatures(s, nRows, nCols, s2row, s2col, goalState);
    stateDim = length(encodeState(1));  % feature dimension

    %% =====================================================================
    %  3. Build Q-network
    %  =====================================================================
    rng(2);
    hiddenSizes = [stateDim, 64, 64, nActions];
    qParams = initNetParams(hiddenSizes);
    tParams = qParams;  % target network

    % Adam state
    adamM = structfun(@(x) dlarray(zeros(size(x),'single')), qParams, 'UniformOutput',false);
    adamV = structfun(@(x) dlarray(zeros(size(x),'single')), qParams, 'UniformOutput',false);
    adamT = 0;
    lr = 5e-4; beta1 = 0.9; beta2 = 0.999; eps_adam = 1e-8;

    %% =====================================================================
    %  4. Experience replay buffer
    %  =====================================================================
    bufSize = 5000;
    buf = struct;
    buf.s  = zeros(bufSize, stateDim, 'single');
    buf.a  = zeros(bufSize, 1);
    buf.r  = zeros(bufSize, 1, 'single');
    buf.ns = zeros(bufSize, stateDim, 'single');
    buf.d  = zeros(bufSize, 1, 'single');  % done flag
    buf.count = 0; buf.idx = 0;

    %% =====================================================================
    %  5. Training parameters
    %  =====================================================================
    nEpisodes    = 800;
    maxStepsEp   = 100;    % max steps per episode
    miniBatch    = 32;
    targetUpdate = 20;     % copy Q -> target every N episodes
    epsilonStart = 1.0;
    epsilonEnd   = 0.05;
    epsilonDecay = 300;

    %% =====================================================================
    %  6. Compute tabular optimal solution for comparison
    %  =====================================================================
    [Vtab, policyTab] = tabularPolicyIteration(grid, nextStateTbl, rewardTbl, ...
                        gamma, nStates, nActions, terminalStates, s2row, s2col);

    %% =====================================================================
    %  7. Set up figure
    %  =====================================================================
    fig = figure('Name','DQN Grid World','Color','w',...
                 'Position',[40 30 1400 850]);

    % Top-left: DQN value function
    ax_vdqn = subplot(2,3,1,'Parent',fig);
    title(ax_vdqn, 'DQN V(s) = max_a Q(s,a)', 'FontSize', 12);

    % Top-center: DQN policy
    ax_pdqn = subplot(2,3,2,'Parent',fig);
    title(ax_pdqn, 'DQN Policy', 'FontSize', 12);

    % Top-right: Tabular optimal (reference)
    ax_tab = subplot(2,3,3,'Parent',fig);
    drawGridV(ax_tab, grid, Vtab, policyTab, nRows, nCols, arrowDx, arrowDy, s2row, s2col);
    title(ax_tab, 'Tabular Optimal (reference)', 'FontSize', 12);

    % Bottom-left: episode reward
    ax_rew = subplot(2,3,4,'Parent',fig);
    hold(ax_rew,'on'); set(ax_rew,'XGrid','on','YGrid','on');
    hRew    = animatedline(ax_rew,'Color',[0.7 0.7 0.7],'LineWidth',0.5);
    hRewAvg = animatedline(ax_rew,'Color',[0 0.45 0.74],'LineWidth',2.5);
    xlabel(ax_rew,'Episode'); ylabel(ax_rew,'Reward');
    title(ax_rew,'Episode Reward','FontSize',12);
    legend(ax_rew,{'Per episode','Avg (20)'},'Location','southeast');

    % Bottom-center: Bellman loss
    ax_loss = subplot(2,3,5,'Parent',fig);
    hold(ax_loss,'on'); set(ax_loss,'XGrid','on','YGrid','on');
    hLoss = animatedline(ax_loss,'Color',[0.47 0.67 0.19],'LineWidth',1.5);
    xlabel(ax_loss,'Episode'); ylabel(ax_loss,'Loss');
    title(ax_loss,'Bellman MSE Loss','FontSize',12);
    set(ax_loss,'YScale','log');

    % Bottom-right: policy accuracy vs tabular
    ax_acc = subplot(2,3,6,'Parent',fig);
    hold(ax_acc,'on'); set(ax_acc,'XGrid','on','YGrid','on');
    hAcc = animatedline(ax_acc,'Color',[0.85 0.33 0.1],'LineWidth',2,...
                        'Marker','o','MarkerSize',3,'MarkerFaceColor',[0.85 0.33 0.1]);
    plot(ax_acc, [0 nEpisodes], [100 100], 'k--', 'LineWidth', 0.8);
    xlabel(ax_acc,'Episode'); ylabel(ax_acc,'Accuracy (%)');
    title(ax_acc,'Policy Match vs Tabular Optimal','FontSize',12);
    ylim(ax_acc, [0 105]);

    iterText = annotation(fig,'textbox',[0.25 0.005 0.5 0.03],...
        'String','Episode: 0','FontSize',12,'EdgeColor','none',...
        'HorizontalAlignment','center','FontWeight','bold');

    drawnow;

    %% =====================================================================
    %  8. Training loop: DQN with experience replay
    %  =====================================================================
    episodeRewards = zeros(nEpisodes, 1);
    episodeLosses  = zeros(nEpisodes, 1);

    for ep = 1:nEpisodes
        % Epsilon decay
        epsil = epsilonEnd + (epsilonStart - epsilonEnd) * exp(-ep / epsilonDecay);

        % Random start from a free cell
        s = freeStates(randi(length(freeStates)));
        totalReward = 0;
        totalLoss = 0;
        nLoss = 0;

        for step = 1:maxStepsEp
            feat = encodeState(s);

            % Epsilon-greedy
            if rand < epsil
                action = randi(nActions);
            else
                dlF = dlarray(single(feat), 'CB');
                qVals = extractdata(qForward(qParams, dlF, hiddenSizes));
                [~, action] = max(qVals);
            end

            % Environment step
            ns = nextStateTbl(s, action);
            rew = rewardTbl(s, action);
            done = ismember(ns, terminalStates);
            totalReward = totalReward + rew;

            % Store experience
            featNS = encodeState(ns);
            buf.idx = mod(buf.idx, bufSize) + 1;
            buf.s(buf.idx,:)  = feat;
            buf.a(buf.idx)    = action;
            buf.r(buf.idx)    = rew;
            buf.ns(buf.idx,:) = featNS;
            buf.d(buf.idx)    = done;
            buf.count = min(buf.count + 1, bufSize);

            s = ns;

            % Train from replay
            if buf.count >= miniBatch
                bIdx = randperm(buf.count, miniBatch);
                dlS  = dlarray(single(buf.s(bIdx,:)'), 'CB');
                bA   = buf.a(bIdx);
                dlR  = dlarray(single(buf.r(bIdx)'));
                dlNS = dlarray(single(buf.ns(bIdx,:)'), 'CB');
                dlD  = dlarray(single(buf.d(bIdx)'));

                % Target: r + gamma * max_a' Q_target(s',a') * (1-done)
                qNext = qForward(tParams, dlNS, hiddenSizes);
                qMaxNext = max(qNext, [], 1);
                targets = dlR + gamma * qMaxNext .* (1 - dlD);

                % Gradient step
                [loss, grads] = dlfeval(@qGradients, qParams, dlS, bA, ...
                                        targets, hiddenSizes, miniBatch);

                % Adam update
                adamT = adamT + 1;
                fnames = fieldnames(grads);
                for fi = 1:length(fnames)
                    fn = fnames{fi};
                    adamM.(fn) = beta1*adamM.(fn) + (1-beta1)*grads.(fn);
                    adamV.(fn) = beta2*adamV.(fn) + (1-beta2)*grads.(fn).^2;
                    mHat = adamM.(fn) / (1 - beta1^adamT);
                    vHat = adamV.(fn) / (1 - beta2^adamT);
                    qParams.(fn) = qParams.(fn) - lr * mHat ./ (sqrt(vHat) + eps_adam);
                end

                totalLoss = totalLoss + extractdata(loss);
                nLoss = nLoss + 1;
            end

            if done, break; end
        end

        episodeRewards(ep) = totalReward;
        episodeLosses(ep)  = totalLoss / max(nLoss, 1);

        % Update target network
        if mod(ep, targetUpdate) == 0
            tParams = qParams;
        end

        % --- Update plots every 10 episodes ---
        if mod(ep, 10) == 0 || ep == 1 || ep == nEpisodes
            % Compute V(s) and policy from Q-network
            Vdqn = zeros(nStates, 1);
            policyDqn = ones(nStates, 1);
            for si = 1:nStates
                if grid(s2row(si),s2col(si)) == 1 || ismember(si, terminalStates)
                    continue;
                end
                feat = encodeState(si);
                dlF = dlarray(single(feat), 'CB');
                qVals = extractdata(qForward(qParams, dlF, hiddenSizes));
                [Vdqn(si), policyDqn(si)] = max(qVals);
            end

            % Draw DQN value + policy
            drawGridV(ax_vdqn, grid, Vdqn, policyDqn, nRows, nCols, ...
                      arrowDx, arrowDy, s2row, s2col);
            title(ax_vdqn, 'DQN V(s) = max_a Q(s,a)', 'FontSize', 12);

            drawGridPolicy(ax_pdqn, grid, policyDqn, nRows, nCols, ...
                           arrowDx, arrowDy, s2row, s2col);
            title(ax_pdqn, 'DQN Policy', 'FontSize', 12);

            % Reward
            addpoints(hRew, ep, totalReward);
            if ep >= 20
                addpoints(hRewAvg, ep, mean(episodeRewards(max(1,ep-19):ep)));
            end

            % Loss
            if episodeLosses(ep) > 0
                addpoints(hLoss, ep, episodeLosses(ep));
            end

            % Policy accuracy: % of free states matching tabular optimal
            nMatch = sum(policyDqn(freeStates) == policyTab(freeStates));
            acc = 100 * nMatch / length(freeStates);
            addpoints(hAcc, ep, acc);

            set(iterText, 'String', sprintf(...
                'Episode: %d / %d  |  Reward: %.1f  |  \\epsilon: %.3f  |  Policy match: %.0f%%', ...
                ep, nEpisodes, totalReward, epsil, acc));

            drawnow;
        end
    end

    %% =====================================================================
    %  9. Animate agent following DQN policy
    %  =====================================================================
    % Extract final DQN policy
    policyDqn = ones(nStates, 1);
    for si = 1:nStates
        if grid(s2row(si),s2col(si))==1 || ismember(si,terminalStates), continue; end
        feat = encodeState(si);
        dlF = dlarray(single(feat),'CB');
        qVals = extractdata(qForward(qParams, dlF, hiddenSizes));
        [~, policyDqn(si)] = max(qVals);
    end

    startState = sub2s(1,1);
    path = startState;
    s = startState;
    for step = 1:nStates
        if ismember(s, terminalStates), break; end
        a = policyDqn(s);
        s = nextStateTbl(s, a);
        path(end+1) = s; %#ok<AGROW>
    end

    % Animate on DQN policy panel
    hold(ax_pdqn,'on');
    agentDot = plot(ax_pdqn, NaN, NaN, 'o', 'MarkerSize', 18, ...
                    'MarkerFaceColor', [1 0.84 0], 'MarkerEdgeColor', 'k', 'LineWidth', 2);
    trailLine = animatedline(ax_pdqn, 'Color', [1 0.84 0], 'LineWidth', 3);

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

        if ismember(si, goalState)
            set(iterText, 'String', sprintf('DQN agent reached GOAL in %d steps!', step-1));
        elseif ismember(si, trapStates)
            set(iterText, 'String', sprintf('DQN agent fell into TRAP at step %d!', step-1));
            set(agentDot, 'MarkerFaceColor', [1 0 0]);
        end
        drawnow;
        pause(0.3);
    end

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

    %% =====================================================================
    %  10. Summary
    %  =====================================================================
    nMatch = sum(policyDqn(freeStates) == policyTab(freeStates));
    acc = 100 * nMatch / length(freeStates);
    fprintf('\n=============================================\n');
    fprintf('  DQN Grid World Summary\n');
    fprintf('=============================================\n');
    fprintf('Grid: %d x %d, States: %d, Free: %d\n', nRows, nCols, nStates, length(freeStates));
    fprintf('Q-network: %d -> 64 -> 64 -> %d\n', stateDim, nActions);
    fprintf('Optimizer: Adam (lr=%.4f)\n', lr);
    fprintf('Episodes: %d, Replay buffer: %d\n', nEpisodes, bufSize);
    fprintf('Gamma: %.2f, Target update: every %d eps\n', gamma, targetUpdate);
    fprintf('---------------------------------------------\n');
    fprintf('Policy match vs tabular optimal: %.1f%%\n', acc);
    fprintf('DQN path length: %d steps\n', length(path)-1);
    fprintf('Tabular path length: %d steps\n', tabularPathLength(sub2s(1,1), policyTab, nextStateTbl, terminalStates, nStates));
    fprintf('Last 20 avg reward: %.2f\n', mean(episodeRewards(max(1,nEpisodes-19):nEpisodes)));
    fprintf('=============================================\n\n');
end

%% =========================================================================
%  State feature encoding
%  =========================================================================
function feat = stateFeatures(s, nRows, nCols, s2row, s2col, goalState)
    r = s2row(s); c = s2col(s);
    gr = s2row(goalState); gc = s2col(goalState);
    % Features: normalized position, distance to goal, one-hot quadrant
    feat = [r / nRows; ...
            c / nCols; ...
            (r - gr) / nRows; ...
            (c - gc) / nCols; ...
            sqrt((r-gr)^2 + (c-gc)^2) / sqrt(nRows^2+nCols^2); ...
            sin(pi * r / nRows); ...
            cos(pi * c / nCols)];
end

%% =========================================================================
%  Q-Network forward / gradients / init
%  =========================================================================
function qVals = qForward(params, x, sizes)
    h = stripdims(x);
    nLayers = length(sizes) - 1;
    for i = 1:nLayers
        h = params.(['W' num2str(i)]) * h + params.(['b' num2str(i)]);
        if i < nLayers, h = relu(h); end
    end
    qVals = h;
end

function [loss, grads] = qGradients(params, states, actions, targets, sizes, B)
    qAll = qForward(params, states, sizes);
    idx = sub2ind([size(qAll,1), B], actions(:)', 1:B);
    qSel = qAll(idx);
    loss = mean((qSel - stripdims(targets)).^2);
    grads = dlgradient(loss, params);
end

function params = initNetParams(sizes)
    params = struct;
    for i = 1:length(sizes)-1
        fi = sizes(i); fo = sizes(i+1);
        params.(['W' num2str(i)]) = dlarray(single(randn(fo,fi)*sqrt(2/fi)));
        params.(['b' num2str(i)]) = dlarray(single(zeros(fo,1)));
    end
end

%% =========================================================================
%  Tabular policy iteration (for ground-truth comparison)
%  =========================================================================
function [V, policy] = tabularPolicyIteration(grid, nextStateTbl, rewardTbl, ...
                       gamma, nStates, nActions, terminalStates, s2row, s2col)
    policy = randi(nActions, nStates, 1);
    V = zeros(nStates, 1);
    for iter = 1:50
        % Policy evaluation
        for ev = 1:200
            Vnew = V;
            for s = 1:nStates
                if ismember(s, terminalStates) || grid(s2row(s),s2col(s))==1, continue; end
                a = policy(s);
                Vnew(s) = rewardTbl(s,a) + gamma * V(nextStateTbl(s,a));
            end
            if max(abs(Vnew-V)) < 1e-6, V = Vnew; break; end
            V = Vnew;
        end
        % Policy improvement
        stable = true;
        for s = 1:nStates
            if ismember(s, terminalStates) || grid(s2row(s),s2col(s))==1, continue; end
            qV = zeros(nActions,1);
            for a = 1:nActions
                qV(a) = rewardTbl(s,a) + gamma * V(nextStateTbl(s,a));
            end
            [~, best] = max(qV);
            if best ~= policy(s), stable = false; policy(s) = best; end
        end
        if stable, break; end
    end
end

function pLen = tabularPathLength(start, policy, nextStateTbl, terminalStates, maxLen)
    s = start; pLen = 0;
    for i = 1:maxLen
        if ismember(s, terminalStates), break; end
        s = nextStateTbl(s, policy(s));
        pLen = pLen + 1;
    end
end

%% =========================================================================
%  Grid drawing helpers
%  =========================================================================
function drawGridV(ax, grid, V, policy, nRows, nCols, arrowDx, arrowDy, s2row, s2col)
    cla(ax); hold(ax,'on');
    Vgrid = flipud(reshape(V, nRows, nCols));
    imagesc(ax, Vgrid); colormap(ax, parula);
    if max(V) > min(V), caxis(ax, [min(V)-0.1, max(V)+0.1]); end

    for s = 1:nRows*nCols
        r = s2row(s); c = s2col(s); yr = nRows-r+1;
        if grid(r,c)==1
            fill(ax, 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
            fill(ax, 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(ax, c, yr, 'G','HorizontalAlignment','center','FontWeight','bold','FontSize',9,'Color','w');
        elseif grid(r,c)==2
            fill(ax, 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(ax, c, yr, 'T','HorizontalAlignment','center','FontWeight','bold','FontSize',9,'Color','w');
        else
            text(ax, c, yr, sprintf('%.1f',V(s)),'HorizontalAlignment','center','FontSize',6,'Color','w');
            a = policy(s);
            quiver(ax, c, yr, arrowDx(a)*0.6, arrowDy(a)*0.6, 0, 'Color',[1 1 1 0.5],'LineWidth',1,'MaxHeadSize',1);
        end
    end
    axis(ax,'equal','tight');
    set(ax,'XTick',1:nCols,'YTick',1:nRows,'YTickLabel',nRows:-1:1);
    hold(ax,'off');
end

function drawGridPolicy(ax, grid, policy, nRows, nCols, arrowDx, arrowDy, s2row, s2col)
    cla(ax); hold(ax,'on');
    for s = 1:nRows*nCols
        r = s2row(s); c = s2col(s); yr = nRows-r+1;
        if grid(r,c)==1
            fill(ax, 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(ax, 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(ax, c, yr, 'GOAL','HorizontalAlignment','center','FontWeight','bold','FontSize',8,'Color','w');
        elseif grid(r,c)==2
            fill(ax, 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(ax, c, yr, 'TRAP','HorizontalAlignment','center','FontWeight','bold','FontSize',8,'Color','w');
        else
            fill(ax, 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]);
            a = policy(s);
            quiver(ax, c, yr, arrowDx(a), arrowDy(a), 0,'Color',[0 0.3 0.7],'LineWidth',2,'MaxHeadSize',0.8);
        end
    end
    % Start marker
    fill(ax, 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(ax, 1, nRows, 'S','HorizontalAlignment','center','FontWeight','bold','FontSize',10);
    axis(ax,'equal','tight');
    set(ax,'XTick',1:nCols,'YTick',1:nRows,'YTickLabel',nRows:-1:1);
    hold(ax,'off');
end
