clc
clear all


nCheckPoints = 7;
nOptions = 3;
alpha = 0.1;
gamma = 0.9;
HRL_QTable = zeros(nCheckPoints,nOptions);
maxEpisodes = 1000;
iteration = 0;
rewardSum = 0;
goalCP=7;
array_it = zeros(1,maxEpisodes);
array_reward = zeros(1,maxEpisodes);

for ii=1:maxEpisodes
    currentCP = randi(6); %% all CP except goal
    iteration = 0;
    rewardSum = 0;
    while (currentCP ~= goalCP)
        if(mod(iteration,5)==0)
            OptionTaken = randi(nOptions);
        else
            [QValue,OptionTaken]=max(HRL_QTable(currentCP,:));
        end
        [reward,nextCP] = MovementProperty(currentCP,OptionTaken);
        HRL_QTable(currentCP,OptionTaken) = HRL_QTable(currentCP,OptionTaken) + alpha*(reward + gamma*max(HRL_QTable(nextCP,:)) - HRL_QTable(currentCP,OptionTaken));
        currentCP = nextCP;
        iteration = iteration + 1;
        rewardSum = rewardSum + reward;
    end
    if (currentCP == goalCP)
        HRL_QTable(currentCP,OptionTaken) = HRL_QTable(currentCP,OptionTaken) + alpha*(reward - HRL_QTable(currentCP,OptionTaken));
        iteration = iteration + 1;
        array_it(ii) = iteration;
        rewardSum = rewardSum + reward;
        array_reward(ii) = rewardSum;
    end
end 

%%

IniParam        % Script 
%##########################################################################
% GENERATE THE ANGULAR VELOCITIES @WHEELS
%##########################################################################
% Virtual time
Ts = 3e-3;
Texe = 100;
t = 0:Ts:Texe;
ANG_vel = 10;
DLevels = [0.055 0.12 0.18 0.25]
%MOTION_type = 0;    % 0: left-rotation about itself
                    % 1: right-rotation about itself
                    % 2: moving forward along an oscillatory path
%[t,w_R,w_L] = WheelMotionGenerator(Ts,Texe,ANG_vel,MOTION_type);  
DIM = length(t);
disp('Simulating the robot motion and the ID sensor data acquisition...')
%##########################################################################
% SIMULATE THE ROBOT'S MOTION
%##########################################################################
% Robot localization (x,y,fi), "ground-truth" --> pre-allocation for speed
% Sensor distance & intersection point --> pre-allocation for speed
NUM_IR_sensors = length(IR_sensor_ori);
Dist_IR_sensor = ones(NUM_IR_sensors,DIM);
IPoint_IR_sensor = zeros(NUM_IR_sensors,2,DIM);
sensor_Levels=zeros(NUM_IR_sensors,DIM);
nLevels = 4;
state = zeros(1,DIM);
nStates = nLevels^size(sensor_Levels,1);
nActions = 9;
nCheckPoints = 7;
nOptions = 3;
%HierarchicalQTable = load('HRL_QTable.mat');
HierarchicalQTable = zeros(nCheckPoints,nOptions);
reward = zeros(1,DIM);
alpha = 0.1;
gamma = 0.9;
currentState = 0;
newState = 0;
Action_Taken= zeros(1,DIM);
Action_Taken(1)=1;
x = zeros(1,DIM);
y = zeros(1,DIM);
fi = zeros(1,DIM);
x(1) = .4; y(1) = 1.05; fi(1) = pi; % Initial robot's configuration

% Simulate the robot's motion and the ID sensor data acquisition
Ts=3e-3;
Texe=3e-3;
currentCP = 0;
episodes = 30;
count = 0;
array_it=zeros(1,episodes);
rewardSum = zeros(1,episodes);
initialx=0.4;
initialy = [0.15 1.05 0.75];
for jj = 1:episodes,
    QTable = corridor.QTable;
    %ans = randi(3);
        x(1)=0.45;
        y(1)=1.05;
        fi(1)=pi;

    for kk = 1:DIM,
        array_it(jj)=array_it(jj)+1;
        if kk ~=1
        if kk > 15

           if ((x(kk-1)>=0.9 && x(kk-1)<=1.2) && (round(y(kk-1),2)==0.59 && round(y(kk-10),2)==0.60)) %% CheckPoint 1
               disp('Entered CheckPoint 1');
               z = rand();
               %if z <=0.1
               if z <= (1/log((jj+1)^3))
                   ind = randi(3);
               else
                   [dum,ind]=max(HierarchicalQTable(1,:));
               end
               if(ind == 1)
                   QTable = door_front.QTable;
               elseif(ind == 2)
                   QTable = door_right.QTable;
               else
                   QTable = door_left.QTable;
               end
               %reward = -1 -(count*0.001);
               reward = -1;
               rewardSum(jj) = rewardSum(jj)+reward;
               count = 0;
               if(currentCP ~= 0 && currentCP ~=1)
                   HierarchicalQTable(currentCP,OptionTaken) = HierarchicalQTable(currentCP,OptionTaken) + alpha*(reward + gamma*max(HierarchicalQTable(1,:)) - HierarchicalQTable(currentCP,OptionTaken));
               end 
               OptionTaken = ind;    
               currentCP = 1;
           elseif (round(x(kk-1),2)==1.19 && round(x(kk-10),2)==1.20) && (y(kk-1)>=0.3 && y(kk-1)<=0.6) %% CheckPoint 2
               disp('Entered CheckPoint 2');
               z = rand();
               if z <= (1/log((jj+1)^3))
               %if z <=0.1
                    ind = randi(3);
               else
                   [dum,ind]=max(HierarchicalQTable(2,:));
               end
               if(ind == 1)
                   QTable = door_front.QTable;
               elseif(ind == 2)
                   QTable = door_right.QTable;
               else
                   QTable = door_left.QTable;
               end
               reward = -1 -(count*0.001);
               %reward = -1;
               rewardSum(jj) = rewardSum(jj)+reward;
               count = 0;
               if(currentCP ~= 0 && currentCP ~=2)
                   HierarchicalQTable(currentCP,OptionTaken) = HierarchicalQTable(currentCP,OptionTaken) + alpha*(reward + gamma*max(HierarchicalQTable(2,:)) - HierarchicalQTable(currentCP,OptionTaken));
               end 
               OptionTaken = ind;    
               currentCP = 2;
           elseif ((x(kk-1)>=0.9 && x(kk-1)<=1.2) && (round(y(kk-10),2)==0.29 && round(y(kk-1),2)==0.30)) %% CheckPoint 3 
               disp('Entered CheckPoint 3');
               z = rand();
               %if z <=0.1
               if z <= (1/log((jj+1)^3))
                   ind = randi(3);
               else
                   [dum,ind]=max(HierarchicalQTable(3,:));
               end
               if(ind == 1)
                   QTable = door_front.QTable;
               elseif(ind == 2)
                   QTable = door_right.QTable;
               else
                   QTable = door_left.QTable;
               end
               %reward = -1 -(count*0.001);
               reward = -1;
               rewardSum(jj) = rewardSum(jj)+reward;
               count = 0;
               if(currentCP ~= 0 && currentCP ~=3)
                   HierarchicalQTable(currentCP,OptionTaken) = HierarchicalQTable(currentCP,OptionTaken) + alpha*(reward + gamma*max(HierarchicalQTable(3,:)) - HierarchicalQTable(currentCP,OptionTaken));
               end 
               OptionTaken = ind;    
               currentCP = 3;
           elseif ((x(kk-1)>0 && x(kk-1)<=0.3) && (round(y(kk-2),2)==0.59 && round(y(kk-1),2)==0.60)) %% CheckPoint 4
               disp('Entered CheckPoint 4');
               %x(kk-1)
               %y(kk-1)
               %fi(kk-1)
               z = rand();
               
               %if z <=0.3
               if z <= (1/log((jj+1)^3))
                   ind = randi(3);
               else
                   [dum,ind]=max(HierarchicalQTable(4,:));
               end
               if(ind == 1)
                   disp('front')
                   QTable = door_front.QTable;
               elseif(ind == 2)
                   disp('right')
                   QTable = door_right.QTable;
               else
                   
                   reward = -10;
                   rewardSum(jj) = rewardSum(jj)+reward;
                   HierarchicalQTable(4,ind) = HierarchicalQTable(4,ind) + alpha*(reward - HierarchicalQTable(4,ind));
                   %break;
                   
                   x(kk-1)=x(1);
                   y(kk-1)=y(1);
                   fi(kk-1)=fi(1);
               end
               %reward = -1 -(count*0.001);
               reward = -1;
               rewardSum(jj) = rewardSum(jj)+reward;
               count = 0;
               if(currentCP ~= 0 && currentCP ~=4)
                   HierarchicalQTable(currentCP,OptionTaken) = HierarchicalQTable(currentCP,OptionTaken) + alpha*(reward + gamma*max(HierarchicalQTable(4,:)) - HierarchicalQTable(currentCP,OptionTaken));
               end 
               OptionTaken = ind;    
               currentCP = 4;
            elseif ((x(kk-1)>=0 && x(kk-1)<=0.3)) && ((round(y(kk-1),2)==0.89) && round(y(kk-2),2)==0.90) %% CheckPoint 5
               disp('Entered CheckPoint 5')
               z = rand();
               %if z <=0.1
               if z <= (1/log((jj+1)^3))
                   ind = randi(3);
               else
                   [dum,ind]=max(HierarchicalQTable(5,:));
               end
               if(ind == 1)
                   QTable = door_front.QTable;
               elseif(ind == 2)
                   reward = -10;
                   rewardSum(jj) = rewardSum(jj)+reward;
                   HierarchicalQTable(5,ind) = HierarchicalQTable(5,ind) + alpha*(reward - HierarchicalQTable(5,ind));
                   %break;
                   x(kk-1)=x(1);
                   y(kk-1)=y(1);
                   fi(kk-1)=fi(1);
               else
                   QTable = door_left.QTable;
               end
               %reward = -1 -(count*0.001);
               reward = -1;
               rewardSum(jj) = rewardSum(jj)+reward;
               count = 0;
               if(currentCP ~= 0 && currentCP ~=5)
                   HierarchicalQTable(currentCP,OptionTaken) = HierarchicalQTable(currentCP,OptionTaken) + alpha*(reward + gamma*max(HierarchicalQTable(5,:)) - HierarchicalQTable(currentCP,OptionTaken));
               end 
               OptionTaken = ind;    
               currentCP = 5;
           elseif (round(x(kk-1),2) ==0.29 && round(x(kk-10),2)==0.31 && (y(kk-1)>=0.6 && y(kk-1)<=0.9)) %% CheckPoint 6 
               %x(kk-1)
               %x(kk-10)
               disp('Entered CheckPoint 6');
               z = rand();
               if z <= (1/log((jj+1)^3))
               %if z <=0.1
                   ind = randi(3);
               else
                   [dum,ind]=max(HierarchicalQTable(6,:));
               end
               if(ind == 1)
                   reward = -10;
                   rewardSum(jj) = rewardSum(jj)+reward;
                   HierarchicalQTable(6,ind) = HierarchicalQTable(6,ind) + alpha*(reward - HierarchicalQTable(6,ind));
                   %break;
                   
                   x(kk-1)=x(1);
                   y(kk-1)=y(1);
                   fi(kk-1)=fi(1);
               elseif(ind == 2)
                   disp('6-right')
                   QTable = door_right.QTable;
               else
                   disp('6-left')
                   QTable = door_left.QTable;
               end
               %reward = -1 -(count*0.001);
               reward = -1;
               rewardSum(jj) = rewardSum(jj)+reward;
               count = 0;
               if(currentCP ~= 0 && currentCP ~=6)
                   HierarchicalQTable(currentCP,OptionTaken) = HierarchicalQTable(currentCP,OptionTaken) + alpha*(reward + gamma*max(HierarchicalQTable(6,:)) - HierarchicalQTable(currentCP,OptionTaken));
               end 
               OptionTaken = ind;    
               currentCP = 6;
           elseif ((x(kk-1)>=0.91 && x(kk-1)<=1.2) && (y(kk-10)>=0.59 && y(kk-1)<=0.61 && y(kk-10)< y(kk-1))) || ((x(kk-10)>=1.19 && x(kk-1)<=1.21 && x(kk-10)<x(kk-1)) && (y(kk-1)>=0.3 && y(kk-1)<=0.6)) || ((x(kk-1)>=0.9 && x(kk-1)<=1.2) && (y(kk-1)>=0.29 && y(kk-10)<=0.31) && y(kk-10)>y(kk-1)) || ((x(kk-1)>=0 && x(kk-1)<=0.3) && (y(kk-1)>=0.59 && y(kk-10)<=0.61 && y(kk-10)>y(kk-1))) || ((x(kk-1)>=0 && x(kk-1)<=0.3) && (y(kk-10)>=0.89 && y(kk-1)<=0.91 && y(kk-10)<y(kk-1))) || ((x(kk-10)>=0.29 && x(kk-1)<=0.31 && x(kk-10)<x(kk-1)) && (y(kk-1)>=0.6 && y(kk-1)<=0.9)) || ((x(kk-1)>=0.88 && x(kk-10)<=0.9 && x(kk-10)>x(kk-1)) && (y(kk-1)>=0.3 && y(kk-1)<=0.6))%% CheckPoint 1 ao contrário
               %disp('Corridor')
               QTable = corridor.QTable;

           elseif ((x(kk-1)>=0.3 && x(kk-1)<=0.35) && (y(kk-1)>=0.3 && y(kk-1)<=0.6)) %% GoalCell
               disp('Entered the goal');
               reward = 10;
               rewardSum(jj) = rewardSum(jj)+reward;
               HierarchicalQTable(currentCP,OptionTaken) = HierarchicalQTable(currentCP,OptionTaken) + alpha*(reward + gamma*max(HierarchicalQTable(7,:)) - HierarchicalQTable(currentCP,OptionTaken));
%                for jj = kk-30:10:kk-10 ,
%                     GraphicalAnimation(A,B,...                          % Maze layout
%                                        [x(jj) y(jj)],fi(jj),...         % Robot location
%                                        IPoint_IR_sensor(:,:,jj),t)  % IR sensor data
%                     if jj == 1, pause, else pause(0.001), end    
%                end
%                 H1=plot(x(1:16500),y(1:16500),'r.')
%                 set(H1,'LineWidth',2),pause
               
               break;
           else
    %            if (Dist_IR_sensor(1,kk-1)<=0.09)||((Dist_IR_sensor(3,kk-1)<=0.09)&&(Dist_IR_sensor(2,kk-1)<=0.09))||((Dist_IR_sensor(5,kk-1)<=0.09)&&(Dist_IR_sensor(6,kk-1)<=0.09))||(Dist_IR_sensor(3,kk-1)<=0.09)||(Dist_IR_sensor(5,kk-1)<=0.09)||(Dist_IR_sensor(6,kk-1)<=0.09)||(Dist_IR_sensor(2,kk-1)<=0.09)
    %                disp("Backwards")
    %                
    %            elseif(Dist_IR_sensor(1,kk-1)<=0.06)||(Dist_IR_sensor(2,kk-1)<=0.06)||(Dist_IR_sensor(3,kk-1)<=0.06)||(Dist_IR_sensor(4,kk-1)<=0.06)||(Dist_IR_sensor(5,kk-1)<=0.06)||(Dist_IR_sensor(6,kk-1)<=0.06)
    %                disp("xxx");
    %                
    %            elseif ((Dist_IR_sensor(4,kk-1)<=0.09)&&(Dist_IR_sensor(3,kk-1)<=0.09))||((Dist_IR_sensor(4,kk-1)<=0.09)&&(Dist_IR_sensor(5,kk-1)<=0.09))
    %                disp("Frontwards")
    %                
    %            elseif (Dist_IR_sensor(1,kk-1)>=1.1)||(Dist_IR_sensor(2,kk-1)>=1.1)||(Dist_IR_sensor(3,kk-1)>=1.1)||(Dist_IR_sensor(4,kk-1)>=1.1)||(Dist_IR_sensor(5,kk-1)>=1.1)||(Dist_IR_sensor(6,kk-1)>=1.1)
    %     
                %disp('else')
           end
        end
                   currentState = newState;
                   [aux,ind]=max(QTable(currentState,:));
                   MOTION_type = ind;
                   Action_Taken(1,kk)=MOTION_type;
                   [t,w_R,w_L] = WheelMotionGenerator(Ts,Texe,ANG_vel,MOTION_type); 
                   % Update the robot's localization (numerical integration)
                   x(kk) = x(kk-1)+R_wheels*cos(fi(kk-1))*Ts*(w_R(1)+w_L(1))/2;
                   y(kk) = y(kk-1)+R_wheels*sin(fi(kk-1))*Ts*(w_R(1)+w_L(1))/2;
                   fi(kk) = fi(kk-1)+R_wheels*Ts*(w_R(1)-w_L(1))/L_wheels;
                   if fi(kk) > pi, fi(kk) = fi(kk)-2*pi; end
                   if fi(kk) < -pi, fi(kk) = fi(kk)+2*pi; end

                   [Dist_IR_sensor(:,kk),IPoint_IR_sensor(:,:,kk)] = ...
                   IR_SensorData(A,B,x(kk),y(kk),fi(kk),IR_sensor_ori);
                   sensor_Levels(:,kk)= getSensorLevels(Dist_IR_sensor(:,kk),DLevels);
                   state(1,kk)=get_State(sensor_Levels(:,kk));
                   newState = state(1,kk);
                   count = count +1;
               end

        [Dist_IR_sensor(:,kk),IPoint_IR_sensor(:,:,kk)] = ...
        IR_SensorData(A,B,x(kk),y(kk),fi(kk),IR_sensor_ori);
        sensor_Levels(:,kk)= getSensorLevels(Dist_IR_sensor(:,kk),DLevels);
        state(1,kk)=get_State(sensor_Levels(:,kk));
        newState = state(1,kk);

    end
end

figure 
subplot(2,1,1),plot(1:episodes,array_it),xlabel('Nº of episodes'),ylabel('Nº of Iterations')
subplot(2,1,2),plot(1:episodes,rewardSum),xlabel('Nº of episodes'),ylabel('Reward Sum')

hf=figure
hf=colordef(hf,'white');
hf.Color='w';
subplot(2,1,1),plot(1:episodes,array_it),xlabel('Nº of episodes','FontSize',20),ylabel('Nº of Iterations','FontSize',20)
subplot(2,1,2),plot(1:episodes,rewardSum),xlabel('Nº of episodes','FontSize',20),ylabel('Reward Sum','FontSize',20)
set(gca,'FontSize',20)

