function [ ts, Ps, Ks ] = continuousLQR(A, B, Q, R, Qf, tf, x0)
% continuousLQR: recursive solution of finite time continuous LQR
%
%   [ ts, Ps, Ks ] = continuousLQR(A, B, Q, R, Qf, tf, x0)
%
% Solve the continuous 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 \dot x = A x + B u where x \in R^n and u \in R^m.
%
% The cost function is \int x(t)^T Q x(t) + u(t)^T R u(t) dt + x(tf)^T Qf x(tf)
%
% We take the recursion from Stephen Boyd's EE 363 course notes lecture
% 4 p. 4-9, which can be found at
%
%           http://www.stanford.edu/class/ee363/clqr.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.
%
%   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 < 5)
    Qf = Q;
  end
  
  if(nargin < 6)
    tf = 10;
  end
  
  if(nargin < 7)
    x0 = [];
  end

  % 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.
  Rinv = inv(R);

  % The old method of passing parameters to the ODE subroutine does not seem
  % to be supported any more, so we'll use an anonymous function to freeze
  % the parameters for the ODE routine.
  odeLQRfixed = @(t, y) odeLQRgeneral(t, y, A, B, Q, Rinv);
  
  % Compute P matrix working backward by solving the terminal value ODE.
  % Since the ODE solver is designed to work forwards in time, we
  % actually multiply the ODE through by -1 (to reverse time) and then
  % swap the solution (to reverse time again).
  [ ts, ps ] = ode45(odeLQRfixed, [ 0 tf ], Qf(:));
  
  % As we said, reverse time to get things in the right order.
  ts = flipud(tf - ts);
  ps = flipud(ps);
  
  % Allocate space for the matrix form of P(t) and the gain matrix K(t).
  N = length(ts);
  Ps = zeros([ size(A), N ]);
  Ks = zeros([ size(B'), N ]);

  % Unpack the matrix form of P and compute the gain matrix at each time.
  % There might be some way of vectorizing this process, but it should be
  % faster than the (looping) ODE solver, so why bother.
  for i = 1 : N
    P = reshape(ps(i, :)', size(A));
    Ps(:,:,i) = P;
    Ks(:,:,i) = -Rinv * B' * 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('Continuous 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('Continuous 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;

      % The ODEs.
      odeOpen = @(t, y) A * y;
      odeClosed = @(t, y) odeClosedGeneral(t, y, A, B, ts, Ks);

      % Solve the ODEs.
      [ t_open, x_open ] = ode45(odeOpen, [ 0, tf ], x0);
      [ t_closed, x_closed ] = ode45(odeClosed, [ 0, tf ], x0);
      
      % 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({ [ 'Continuous time evolution of open ' ...
                'and optimal closed loop trajectories' ], ...
              [ 'from initial condition [ ' num2str(x0') ' ]^T' ] })
      grid on;
      axis equal;
    end
  
  end
  

%-----------------------------------------------------------------------
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%-----------------------------------------------------------------------
function ydot = odeLQRgeneral(t, y, A, B, Q, Rinv)
% odeLQRgeneral: implement the Riccati ODE.
%
%   ydot = odeLQRgeneral(t, y, A, B, Q, Rinv)
%
% Implements the Riccati ODE (except that we multiply by -1 to take account
% of the fact that the ODE solvers prefer to work forward in time).
%
% Input Parameters:
%   
%   t: Double.  Current time.  Ignored.
%
%   y: n^2 vector.  A vector representation of the matrix P(t).
%
%   A: n x n matrix.
%
%   B: n x m matrix.
%
%   Q: n x n matrix, Q^T = Q >= 0.
%
%   Rinv: m x m matrix, inverse of the R matrix, where R^T = R > 0.
%
% Output Parameters:
%
%   ydot: time derivative of the P matrix, in vector form.
  
  n = size(A, 1);
  P = reshape(y, n, n);
  
  % Don't forget that we're multiplying by -1 to reverse time.
  Pdot = A' * P + P * A - P * B * Rinv * B' * P + Q;
  
  ydot = Pdot(:);
  

  
%-----------------------------------------------------------------------
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%-----------------------------------------------------------------------
function ydot = odeClosedGeneral(t, y, A, B, ts, Ks)
% odeCloseGeneral: implement the optimal closed loop dynamics.
%
%   ydot = odeClosedGeneral(t, y, A, B, ts, Ks)
%
% Implements the optimal closed loop dynamics:
%
%   \dot x = A * x(t) + B * K(t) * x(t)
%
% The reason that we need to do this ODE as a separate subfunction is that
% we need to interpolate the value of K(t) based on the results from solving
% for P(t) and K(t) at timesteps chosen by the ODE integrator.
%
% Input Parameters:
%   
%   t: Double.  Current time.
%
%   y: n x 1 vector.  Current state.
%
%   A: n x n matrix.
%
%   B: n x m matrix.
%
%   ts: vector of times at which the gain matrix is known.
%
%   Ks: array storing the gain matrix at each of those times.
%
% Output Parameters:
%
%   ydot: time derivative of the current state.

  % Interpolate the current gain matrix.
  K = zeros(size(B'));
  for i = 1 : size(B', 1)
    for j = 1 : size(B', 2)
      K(i,j) = interp1(ts, squeeze(Ks(i,j,:)), t);
    end
  end
  
  % Closed loop dynamics.
  ydot = A * y + B * K * y;
  
