%% Gridworld Simulation with GPs for Q learning
clear all;
close all;
clc;
%% Gridworld Parameters
N_grid = 5; % Number of grids on one direction
N_state = N_grid*N_grid; % Square Gridworld
N_act = 5; % 4 directions + null action
%% Create the undirected grid graph
Edge_start = []; %Edge starting point
Edge_end = [];% Edge ending point
% There are 5*5 nodes in the grid. They are numbered by 1~25.  
for k = 1:(N_grid-1)
    temp1 = [];
    temp2 = [];
    for j = 1:(N_grid-1)
        temp1 = [temp1 j j];
        temp2 = [temp2 j+1 j+N_grid];
    end
    temp1 = [temp1 N_grid]+ones(1,N_grid*2-1)*(k-1)*N_grid;
    temp2 = [temp2 N_grid*2]+ones(1,N_grid*2-1)*(k-1)*N_grid;
    Edge_start = [Edge_start temp1];
    Edge_end = [Edge_end temp2];
end
temp1 = [];
temp2 = [];
for j = 1:(N_grid-1)
    temp1 = [temp1 j];
    temp2 = [temp2 j+1];
end
temp1 = temp1+ones(1,N_grid-1)*(N_grid-1)*N_grid;
temp2 = temp2+ones(1,N_grid-1)*(N_grid-1)*N_grid;
Edge_start = [Edge_start temp1];
Edge_end = [Edge_end temp2];
Grid = graph(Edge_start, Edge_end); %Create the grid graph
%% Initialization
n_state_dim=2;
s_init = [1;1]; % Start in the top left corner
a_init=2;
s_goal = [N_grid;N_grid]; % Goal is at the upper right corner [end,end]
N_eps_length =100; % Length of an episode
%% storage variables
STATE_STORE = zeros(N_eps_length,2);
STATE_STORE(1,:) = s_init';
state_number = 1;


%% Algorithm Execution
        for k = 1: N_eps_length            
            %% This is the main part of the code
            % As the code stands, it computes the next state given an
            % action for N_eps_length iterations, 
            % there are 5 actions, they are described in
            % gridworld_trans_HW
            % the purpose of this code is to show you how to use the
            % gridworld_trans_HW function and plot the results
            % in the current code the action has been fixed
            % your job is to write a search algorithm that gets you from
            % s_init to s_goal
            
            % Next State computation
            randomnode = randi([1 25],1);%Random node is picked
            s_new = nodetovector(randomnode,N_grid);
            if attemptconnection(s_new,STATE_STORE,Edge_start, Edge_end, N_grid)==1%Connection is possible and the new node is not currently in the tree.
                state_number = state_number+1;
                STATE_STORE(state_number,:)=s_new';
                if s_new == s_goal%If the goal point is reached, the search loop is terminated.
                    break;
                end
            end
            
        end % end main loop

%% Plots
figure(1)
STATE_STORE(1:state_number,3)=(1:state_number)';
scatter3(STATE_STORE(1:state_number,1),STATE_STORE(1:state_number,2),STATE_STORE(1:state_number,3),25,linspace(1,state_number,state_number),'filled');
grid on
xlabel('x state')
ylabel('y state')
zlabel('Number of nodes discovered')