function f = cem_car()

% time horizon and segments
tf = 5;
S.N = 16;
S.h = tf/S.N;   %--- dt

% cost function parameters
S.Q = .0*diag([5, 5, 1, 1, 1]);
S.R = .01*diag([1, 5]);
S.Qf = diag([5, 5, 1, 1, 1]);

S.f = @car_f;   %--- Dynamics
S.L = @car_L;   %--- running cost
S.Lf = @car_Lf; %--- terminal cost
S.mu = 1;       

% control dimension
S.m=2;

% initial state
x0 = [-5; -4; 0; 1; 0];

% initial control sequence
us = zeros(2,S.N);
%us = .1*[repmat([.1;0], 1, S.N/2), repmat(-[.1;0], 1, S.N/2)];
%us = [repmat([.1;0], 1, N/2), repmat(-[.1;0], 1, N/2)]/5;

opts.v = .8;      % mixing factor
opts.iter = 1;    % internal iterations
opts.N = 50;      % # of samples
opts.sigma = 0;  % the sigma-point version
%opts.tilt = 1;
%opts.sigma = 1;

C0 = diag(repmat([1; 0.1], S.N, 1));

iters = 20;
cs=zeros(iters,1);
x=x0;
opts.C = C0;

S.x0 = x0;

xs = sys_traj(x0, us, S);

figure

subplot(1,3,1)

plot(xs(1,:), xs(2,:), '-b')
hold on

mu = reshape(us, 2*S.N, 1);

for i=1:iters  
  [us, c, mu, C] = cem(@car_cost, mu, opts, S);
  cs(i) = c;
  opts.C = C;
  us = reshape(us, 2,S.N);
  
  % update trajectory
  xs = sys_traj(x0, us, S);
  
  plot(xs(1,:), xs(2,:), '-b');

end

plot(xs(1,:), xs(2,:), '-g');

J = car_cost(us, S);
J'*J

xlabel('x')
ylabel('y')
axis equal

subplot(1,3,2)

plot(0:S.h:tf-S.h, us(1,:),0:S.h:tf-S.h, us(2,:));
xlabel('sec.')
legend('u_1','u_2')

subplot(1,3,3)
plot(cs);

function y = car_cost(us, S)

us = reshape(us, S.m, S.N);

xs = sys_traj(S.x0, us, S);

N=size(us,2);

%y = zeros(S.N*(S.n+S.m), 1);
y=[];
for k=1:N
  y = [y; S.R*us(:,k)];
end

for k=2:N
  y = [y; S.Q*xs(:,k)];
end
y = [y; S.Qf*xs(:,N+1)];



function [L, Lx, Lxx, Lu, Luu] = car_L(k, x, u, S)
% car cost (just standard quadratic cost)

if (k == S.N+1)
  L = x'*S.Qf*x/2;
  Lx = S.Qf*x;
  Lxx = S.Qf;
  Lu = [];
  Luu = [];
else
  L = S.h/2*(x'*S.Q*x + u'*S.R*u);
  Lx = S.h*S.Q*x;
  Lxx = S.h*S.Q;
  Lu = S.h*S.R*u;
  Luu = S.h*S.R;
end


function [x, A, B] = car_f(k, x, u, S)
% unicycle dynamics and jacobians (optional)

%--- x = [ X, Y, theta, v, theta_dot ]

h = S.h;
c = cos(x(3));
s = sin(x(3));
v = x(4);
w = x(5);

if nargout > 1
    A = [1 0 -h*s*v h*c 0;
         0 1 h*c*v h*s 0;
         0 0 1 0 h;
         0 0 0 1 0;
         0 0 0 0 1];

    B = [0 0;
         0 0;
         0 0;
         h 0;
         0 h];
end

x = [x(1) + h*c*v;
     x(2) + h*s*v;
     x(3) + h*w;
     v + h*u(1);
     w + h*u(2)];


function xs = sys_traj(x0, us, S)

N = size(us, 2);
xs(:,1) = x0;

for k=1:N,
  xs(:, k+1) = S.f(k, xs(:,k), us(:,k), S);
end
