function [ ts, Ps, Ks ] = discreteLQR(A, B, Q, R, h, Qf, tf, x0)
% discreteLQR: recursive solution of finite time discrete LQR
%
%   [ ts, Ps, Ks ] = discreteLQR(A, B, Q, R, h, Qf, tf, x0)
%
% Solve the discrete time LQR recursion to find the value function (given by
% V(x, t) = x^T P(t) x where P(t) is an SPD matrix) and optimal input (given
% by u(x, t) = K(t) x where K(t) is a matrix) for an LQR optimal control
% problem (Linear dynamics and Quadratic cost Regulator).
%
% The dynamics are x(t+1) = A x(t) + B u(t) where x \in R^n and u \in R^m.
%
% The cost function is \sum_t x(t)^T Q x(t) + u(t)^T R u(t) + x(tf)^T Qf x(tf)
%
% We take the recursion from Stephen Boyd's EE 363 course notes lecture
% 1 p. 1-23, which can be found at
%
%           http://www.stanford.edu/class/ee363/dlqr.pdf
%
% If n == 2, a plot is generated showing the optimal value function at
% the initial time V(x, 0) for the interval [ -1, +1 ]^2.
%
% If n == 2 & m == 1, plots are generated showing the time evolution of the
% components of P(t) and K(t), and if an initial condition is given, then an
% open and a closed loop trajectory is plotted in phase space.
%
% Note that the input parameters are not checked for appropriate dimensions.
% If dimensions are wrong, a standard Matlab dimension error will occur
% somewhere in this function.
%
% Input Parameters:
%
%   A: n x n matrix.
%
%   B: n x m matrix.
%
%   Q: n x n matrix, Q^T = Q >= 0.
%
%   R: m x m matrix, R^T = R > 0.
%
%   h: Double.  Time step size.  The sample times a which the cost
%   function is evaluated are t_i = i h such that t_i \in [ 0, tf ].
%
%   Qf: n x n matrix.  Qf^T = Qf >= 0.  Optional.  Default Qf = Q.
%
%   tf: Double.  Final time.  Optional.  Default tf = 10.
%
%   x0: m x 1 vector.  Initial condition for the plot.  Optional.  If no
%   x0 is given, then no trajectory plot is generated.
%
% Output Parameters:
%
%   ts: Vector of length N.  Sample times t_i.
%
%   Ps: Array n x n x N.  For each value of the final index, this array
%   gives the matrix defining the quadratic optimal value function at the
%   corresponding time t_i.
%
%   Ks: Array n x m x N.  For each value of the final index, this array
%   gives the gain matrix defining the optimal input at the corresponding
%   time t_i.
%
% Ian Mitchell for CPSC 532M, 5/4/08.
  
  % Check optional parameters.
  if(nargin < 6)
    Qf = Q;
  end
  
  if(nargin < 7)
    tf = 10;
  end
  
  if(nargin < 8)
    x0 = [];
  end

  % Determine our timesteps and allocate space for the results.
  ts = 0 : h : tf;
  N =  length(ts);
  Ps = zeros([ size(A), N ]);
  Ks = zeros([ size(B'), N ]);
  
  % Compute P matrix working backward.  We'll also compute the K matrix,
  % although it could be computed in either direction.

  % Start the iteration.
  Ps(:,:,N) = Qf;
  P = Ps(:,:,N);

  for i = N - 1 : -1 : 1
    % From a numerical standpoint, it would be better to compute this
    % inverse as an operator by factorization; however, we are dealing
    % with small well-conditioned matrices so I'll just invert it.
    inverse = inv(R + B' * P * B);

    % From Boyd's notes, we are computing K_i using P_{i+1}.
    Ks(:,:,i) = -inverse * B' * P * A;

    % From Boyd's notes, we are computing P_i using P_{i+1}.
    P = Q + A' * (P - P * B * inverse * B' * P) * A;

    % Store the result.
    Ps(:,:,i) = P;
  end

  % Plot of the value function at the initial time.
  if(size(A, 1) == 2)
    figure;
    [ xs, ys ] = meshgrid(-1:0.01:+1);
    V = P(1,1,1) * xs.^2 + 2 * P(1,2,1) * xs .* ys + P(2,2,1) * ys.^2;
    surf(xs, ys, V);
    xlabel('x_1');  ylabel('x_2');  zlabel('V(x,0)');
    title('Discrete time value function at the initial time');
    daspect([ 1 1 max(V(:)) ]);
  end
  
  % Now we generate a plot showing how P and K evolve with time.  We
  % should notice a quick convergence to a steady state value.
  if all(size(B) == [ 2, 1 ])
    figure;
  
    % First the components of P (which is symmetric).
    subplot(2,1,1)
    plot(ts, squeeze(Ps(1,1,:)), 'b-+', ...
         ts, squeeze(Ps(1,2,:)), 'k:*', ...
         ts, squeeze(Ps(2,2,:)), 'r--o', ...
         'LineWidth', 2);
    legend('P_{11}', 'P_{12} = P_{21}', 'P_{22}', 'Location', 'SouthWest');
    ylabel('P(t)');
    title('Discrete time evolution of matrices P(t) and K(t)');
    grid on;
    
    % Then the components of K.
    subplot(2,1,2)
    plot(ts, squeeze(Ks(1,1,:)), 'b-+', ...
         ts, squeeze(Ks(1,2,:)), 'r--o', 'LineWidth', 2);
    legend('K_1', 'K_2', 'Location', 'NorthWest');
    ylabel('K(t)');
    xlabel('t');
    grid on;

    % Generate trajectory plots if necessary.
    if ~isempty(x0)
      
      % Allocate space for the trajectories.
      x_open = zeros([ length(x0), N ]);
      x_closed = zeros(size(x_open));
      
      % Initial conditions.
      x_open(:, 1) = x0;
      x_closed(:, 1) = x0;
      
      % March through time following the dynamics.
      for i = 2 : N
        
        % Input is zero for the open loop case.
        x_open(:, i) = A * x_open(:, i-1);
        
        % Input is optimal for the closed loop case.
        u = Ks(:,:,i-1) * x_closed(:,i-1);
        x_closed(:, i) = A * x_closed(:, i-1) + B * u;
        
      end
      
      % Plot the results in phase space.
      figure;
      plot(x_open(1,:), x_open(2,:), 'r--+', ...
           x_closed(1,:), x_closed(2,:), 'b-o', ...
           'LineWidth', 2);
      legend('Open Loop', 'Closed Loop');
      xlabel('x_1');  ylabel('x_2');
      title({ [ 'Discrete time evolution of open ' ...
                'and optimal closed loop trajectories' ], ...
              [ 'from initial condition [ ' num2str(x0') ' ]^T' ] })
      grid on;
      axis equal;
    end
  
  end
  
