function [Q_sarsa,Q_qlearn,rpt_sarsa,rpt_qlearn,n_sarsa,n_qlearn] = learn_cw(alpha,CF,s_start,s_end,MAX_N_EPISODES)
% LEARN_CW - Performs on-policy sarsa and Q-learning to learn the policy for the 
% cliff walking problem example.
% 
% Written by:
% -- 
% John L. Weatherwax                2007-12-03
% 
% email: wax@alum.mit.edu
% 
% Please send comments and especially bug reports to the
% above email address.
% 
%-----

PLOT_STEPS = 0; 

gamma = 1;    % <- take this is an undiscounted task 
  
epsilon = 0.1;  % for our epsilon greedy policy 

% the number of states: 
[sideII,sideJJ] = size(CF); 
nStates         = sideII*sideJJ; 

% on each grid we can choose from among this many actions (except on edges where this action is reduced): 
nActions = 4; 

% An array to hold the values of the action-value function 
Q_sarsa    = zeros(nStates,nActions);
Q_qlearn   = zeros(nStates,nActions);
rpt_sarsa  = zeros(1,MAX_N_EPISODES); 
rpt_qlearn = zeros(1,MAX_N_EPISODES); 
n_sarsa    = zeros(nStates,nActions); % <- lets store the number of times we are in this state and take this action
n_qlearn   = zeros(nStates,nActions); 

if( PLOT_STEPS )
  figure; imagesc( CF ); colorbar; hold on; 
  plot( s_start(2), s_start(1), 'x', 'MarkerSize', 10, 'MarkerFaceColor', 'k' ); 
  plot( s_end(2), s_end(1), 'o', 'MarkerSize', 10, 'MarkerFaceColor', 'k' ); 
end

% keep track of how many timestep we take per episode:
ets = zeros(MAX_N_EPISODES,1); ts=0; 
for ei=1:MAX_N_EPISODES,
  tic; 
  if( ei==1 ) 
    fprintf('working on episode %d...\n',ei);
  else
    fprintf('working on episode %d (ptt=%10.6f secs)...\n',ei, toc); tic; 
  end
  ets(ei,1) = ts+1; 
  % initialize the starting state
  st_sarsa  = s_start; sti_sarsa  = sub2ind( [sideII,sideJJ], st_sarsa(1),  st_sarsa(2) ); 
  st_qlearn = s_start; sti_qlearn = sub2ind( [sideII,sideJJ], st_qlearn(1), st_qlearn(2) ); 

  % pick action using an epsilon greedy policy derived from Q: 
  % 
  [dum,at_sarsa] = max(Q_sarsa(sti_sarsa,:));  % at \in [1,2,3,4]=[up,down,right,left]
  if( rand<epsilon )         % explore ... with a random action 
    tmp=randperm(nActions); at_sarsa=tmp(1); 
  end
  [dum,at_qlearn] = max(Q_qlearn(sti_qlearn,:));  % at \in [1,2,3,4]=[up,down,right,left]
  if( rand<epsilon )         % explore ... with a random action 
    tmp=randperm(nActions); at_qlearn=tmp(1); 
  end
  
  % begin an episode
  R_sarsa=0; R_qlearn=0; 
  while( 1 ) 
    ts=ts+1; 
    %fprintf('st=(%d,%d); act=%3d...\n',st(1),st(2),at);
    % propagate to state stp1 and collect a reward rew
    [rew_sarsa, stp1_sarsa]  = stNac2stp1(st_sarsa, at_sarsa,CF, s_start,s_end); R_sarsa=R_sarsa+rew_sarsa; 
    [rew_qlearn,stp1_qlearn] = stNac2stp1(st_qlearn,at_qlearn,CF,s_start,s_end); R_qlearn=R_qlearn+rew_qlearn; 
    %fprintf('stp1=(%d,%d); rew=%3d...\n',stp1(1),stp1(2),rew);
    % pick the greedy action from state stp1: 
    stp1i_sarsa = sub2ind( [sideII,sideJJ], stp1_sarsa(1), stp1_sarsa(2) ); 
    stp1i_qlearn = sub2ind( [sideII,sideJJ], stp1_qlearn(1), stp1_qlearn(2) ); 
    
    % SARSA: 
    % --make the greedy action selection: 
    [dum,atp1_sarsa] = max(Q_sarsa(stp1i_sarsa,:)); 
    if( rand<epsilon )         % explore ... with a random action 
      tmp=randperm(nActions); atp1_sarsa=tmp(1); 
    end
    n_sarsa(sti_sarsa,at_sarsa) = n_sarsa(sti_sarsa,at_sarsa)+1; 
    if( ~( (stp1_sarsa(1)==s_end(1)) && (stp1_sarsa(2)==s_end(2)) ) ) % stp1 is not the terminal state
      Q_sarsa(sti_sarsa,at_sarsa) = Q_sarsa(sti_sarsa,at_sarsa) + alpha*( rew_sarsa + gamma*Q_sarsa(stp1i_sarsa,atp1_sarsa) - Q_sarsa(sti_sarsa,at_sarsa) ); 
    else                                                  % stp1 IS the terminal state ... no Q_sarsa(s';a') term in the sarsa update
      Q_sarsa(sti_sarsa,at_sarsa) = Q_sarsa(sti_sarsa,at_sarsa) + alpha*( rew_sarsa - Q_sarsa(sti_sarsa,at_sarsa) ); 
      break; 
    end
    %update (st,at) pair: 
    st_sarsa = stp1_sarsa; sti_sarsa = stp1i_sarsa; at_sarsa = atp1_sarsa; 
    
    % Q-learning:
    % --make the greedy action selection: 
    n_qlearn(sti_qlearn,at_qlearn) = n_qlearn(sti_qlearn,at_qlearn)+1; 
    [dum,atp1_qlearn] = max(Q_qlearn(stp1i_qlearn,:)); 
    if( ~( (stp1_qlearn(1)==s_end(1)) && (stp1_qlearn(2)==s_end(2)) ) ) % stp1 is not the terminal state
      Q_qlearn(sti_qlearn,at_qlearn) = Q_qlearn(sti_qlearn,at_qlearn) + alpha*( rew_qlearn + gamma*Q_qlearn(stp1i_qlearn,atp1_qlearn) - Q_qlearn(sti_qlearn,at_qlearn) ); 
    else                                                  % stp1 IS the terminal state ... no Q_qlearn(s';a') term in the qlearn update
      Q_qlearn(sti_qlearn,at_qlearn) = Q_qlearn(sti_qlearn,at_qlearn) + alpha*( rew_qlearn - Q_qlearn(sti_qlearn,at_qlearn) ); 
      break; 
    end
    %update (st,at) pair: 
    st_qlearn = stp1_qlearn; sti_qlearn = stp1i_qlearn; at_qlearn = atp1_qlearn; 
    
    
    %pause; 
  end % end policy while 
  rpt_sarsa(ei) = R_sarsa; rpt_qlearn(ei) = R_qlearn;
end % end episode loop 



function [rew,stp1] = stNac2stp1(st,act,CF,s_start,s_end)
% STNAC2STP1 - state and action to state plus one and reward 
%   

% extract dimensions: 
[sideII,sideJJ] = size(CF); 

% convert to row/column notation: 
ii = st(1); jj = st(2); 

% incorporate any actions and fix our position if we end up outside the grid:
% 
switch act
 case 1, 
  %
  % action = UP 
  %
  stp1 = [ii-1,jj];
 case 2,
  %
  % action = DOWN
  %
  stp1 = [ii+1,jj];
 case 3,
  %
  % action = RIGHT
  %
  stp1 = [ii,jj+1];
 case 4
  %
  % action = LEFT 
  %
  stp1 = [ii,jj-1];
 otherwise
  error(sprintf('unknown value for of action = %d',act)); 
end

% adjust our position of we have fallen outside of the grid:
%
if( stp1(1)<1      ) stp1(1)=1;      end
if( stp1(1)>sideII ) stp1(1)=sideII; end
if( stp1(2)<1      ) stp1(2)=1;      end
if( stp1(2)>sideJJ ) stp1(2)=sideJJ; end

% get the reward for this step: 
% 
if( (ii==s_end(1)) && (jj==s_end(2)) )  % were at the end :)
  %rew = +1;
  rew = 0;
elseif( CF(stp1(1),stp1(2))==0 )        % we fell off the cliff :(
  rew  = -100;
  stp1 = s_start; 
else                                    % normal step 
  rew = -1;
end