function wavelet_example
% Compute approximation using first-order interpolet's
%
% C.N. Jones, EPFL, 2013
% Revised by M. Herceg, ETH Zurich 2014
%
% requires MPT3, YALMIP

close all

%% printing parameters
label_font_size = 14;
tick_font_size = 10;
line_width = 0.8;
axeswidth=0.2;
% create new folder for storing figures
mkdir('figures');


figure_name1 = ['.',filesep,'figures',filesep,'wavelet_partition'];
figure_name2 = ['.',filesep,'figures',filesep,'wavelet_approxJ'];
figure_name3 = ['.',filesep,'figures',filesep,'wavelet_approxU'];
figure_name4 = ['.',filesep,'figures',filesep,'wavelet_errU'];
figure_name5 = ['.',filesep,'figures',filesep,'wavelet_errJ'];
figure_name6 = ['.',filesep,'figures',filesep,'wavelet_secondOrderInterpolant'];


%% core code
fprintf('---> Setting up problem\n');

% load data
dat = defineEx;
% Compute vertices of the feasible set
var = [dat.variables.x(:,2:end)]; var = [var(:);dat.variables.u'];
o = Opt(dat.constraints,dat.objective,dat.variables.x(:,1),var);
oi = o.eliminateEquations;
P = Polyhedron([-oi.pB oi.A], oi.b);
% feasible set
p = P.projection(1:size(oi.pB,2));

% extract the parametrized cost function 
dat.opt  = optimizer(dat.constraints, dat.objective, [], dat.variables.x(:,1), dat.objective);
% add cost function
p.addFunction(Function(@(x)optFunc(x, dat)),'cost');
% add penalty
epsilon = 2;
Q = dat.ctrl.model.x.penalty.H;
p.addFunction(Function(@(x)optFunc(x, dat)+epsilon*x(:)'*Q*x(:)),'cost_gam');

% Compute maximum invariant set (terminal set)
termSet = dat.ctrl.model.LQRSet;
dat.termSet = termSet;

% Define the soft-constrained form
sDat = defineEx(true);
var = [sDat.variables.x(:,2:end)]; var = [var(:);sDat.variables.u'];
dat.softOpt = optimizer(sDat.constraints, sDat.objective, [], sDat.variables.x(:,1), sDat.objective);
dat.optU = optimizer(sDat.constraints, sDat.objective, [], sDat.variables.x(:,1), sDat.variables.u(:,1));


% The function that we're trying to approximate
dat.theFunc = @(x) optFunc(x, dat);

dat.Q = Q;
dat.epsilon = epsilon;
dat.lam = sdpvar(4,1);

% Define hierarchical grid
B = p.outerApprox;
grd.lb = B.Internal.lb; grd.ub = B.Internal.ub;

% Compute approximation
clf;
p.plot('alpha',0.1);
hold on
fprintf('Approximating...')
tic
[box,err] = waveletApprox(grd, dat);
fprintf('done.\n')
close all
toc

% get the result
W = flattenWavelet(box);
W = PolyUnion(W);

% Compute the optimal inputs
tic
for i = 1:W.Num
  if toc > 3
    fprintf('%i of %i\n', i, W.Num);
    tic
  end
  W.Set(i).Data.U = dat.optU{W.Set(i).Data.V'}';
end
t=[W.Set.H]; % Make sure we have H-reps

p.addFunction(@(x) evalPWInterpolet(x, W, 'J'), 'approxJ');
p.addFunction(@(x) abs(evalPWInterpolet(x, W, 'J') - dat.opt{x}), 'errJ');
p.addFunction(@(x) evalPWInterpolet(x, W, 'U'), 'approxU');
p.addFunction(@(x) evalPWInterpolet(x, W, 'U') - dat.ctrl.evaluate(x), 'errU');

% plot the partition
disp('hit any key to plot the partition');
pause
fig1 = figure;
hold on
grid on
for i = 1:W.Num
  if any(isnan(W.Set(i).Data.J))
    col = 0.9*[1 1 1];
  else
    col = 0.7*[1 1 1];
  end
  W.Set(i).plot('color',col,'edgecolor',0.5*[1 1 1],'linewidth',line_width);
end
v = p.V;
i = orderForPlot([v zeros(size(v,1),1)],[0 0 1]);
plot(v(i,1),v(i,2),'k--','linewidth',line_width);
axis([-5 5 -5 5]);

set(gca,'LineWidth',axeswidth)
set(gca,'FontSize', tick_font_size);

hx1 = xlabel('$x_1$');
set(hx1, 'FontSize', label_font_size);
hy1 = ylabel('$x_2$');
set(hy1, 'FontSize', label_font_size);


% plot the approximated value function
disp('hit any key to plot the value function');
pause
fig2 = figure;
grid on
p.fplot('approxJ','grid',20,'linewidth',line_width,'color',[0.9 0.9 0.9],'showgrid',true);

set(gca,'LineWidth',axeswidth)
set(gca,'FontSize', tick_font_size);
set(gca,'XTick',[-5 0 5]);
set(gca,'YTick',[-5 0 5]);
set(gca,'ZTick',0:10:30);
axis([-6 6 -6 6 0 30]);

hx2 = xlabel('$x_1$');
set(hx2, 'FontSize', label_font_size);
hy2 = ylabel('$x_2$');
set(hy2, 'FontSize', label_font_size);
hz2 = zlabel('$\tilde{J}(x)$');
set(hz2, 'FontSize', label_font_size);
pos = get(gca,'Position');
set(gca,'Position',[0.17 0.17 pos(3:4)]);
posz2 = get(hz2,'Position');
set(hz2,'Position',[posz2(1)-1 posz2(2:3)]);




% plot the approximated control law
disp('hit any key to plot the approximated control law');
pause
fig3 = figure;
grid on
p.fplot('approxU','grid',30,'linewidth',line_width,'color',[0.9 0.9 0.9],'showgrid',true);
view(-121,20)

set(gca,'LineWidth',axeswidth)
set(gca,'FontSize', tick_font_size);
set(gca,'XTick',[-5 0 5]);
set(gca,'YTick',[-5 0 5]);
set(gca,'ZTick',-1:0.5:1);
axis([-6 6 -6 6 -1 1]);

hx3 = xlabel('$x_1$');
set(hx3, 'FontSize', label_font_size);
hy3 = ylabel('$x_2$');
set(hy3, 'FontSize', label_font_size);
hz3 = zlabel('$\tilde{u}(x)$');
set(hz3, 'FontSize', label_font_size);
pos = get(gca,'Position');
set(gca,'Position',[0.17 0.17 pos(3:4)]);
posz3 = get(hz3,'Position');
set(hz3,'Position',[posz3(1)+5 posz3(2:3)]);



% Plot input error
disp('hit any key to plot the input error');
pause
fig4 = figure;
grid on
p.fplot('errU','grid',30,'linewidth',line_width,'color',[0.9 0.9 0.9],'showgrid',true);

view(-121,20)
set(gca,'LineWidth',axeswidth)
set(gca,'FontSize', tick_font_size);
set(gca,'XTick',[-5 0 5]);
set(gca,'YTick',[-5 0 5]);
set(gca,'ZTick',-1:0.5:1);
axis([-6 6 -6 6 -1 1]);

hx4 = xlabel('$x_1$');
set(hx4, 'FontSize', label_font_size);
hy4 = ylabel('$x_2$');
set(hy4, 'FontSize', label_font_size);
hz4 = zlabel('$\tilde{u}(x)-u^{*}(x)$');
set(hz4, 'FontSize', label_font_size);
pos = get(gca,'Position');
set(gca,'Position',[0.17 0.17 pos(3:4)]);
posz4 = get(hz4,'Position');
set(hz4,'Position',[posz4(1)+5 posz4(2:3)]);



% Plot cost error
disp('hit any key to plot the cost error');
pause
fig5 = figure;
grid on
p.fplot('errJ','grid',30,'linewidth',line_width,'color',[0.9 0.9 0.9],'showgrid',true);

set(gca,'LineWidth',axeswidth)
set(gca,'FontSize', tick_font_size);
pos = get(gca,'Position');
set(gca,'Position',[0.17 0.17 pos(3:4)]);
set(gca,'XTick',[-5 0 5]);
set(gca,'YTick',[-5 0 5]);
set(gca,'ZTick',0:4);
axis([-6 6 -6 6 0 4]);


hx5 = xlabel('$x_1$');
set(hx5, 'FontSize', label_font_size);
hy5 = ylabel('$x_2$');
set(hy5, 'FontSize', label_font_size);
hz5 = zlabel('$\tilde{J}(x)-J^{*}(x)$');
set(hz5, 'FontSize', label_font_size);
posz5 = get(hz5,'Position');
set(hz5,'Position',[posz5(1)-1 posz5(2:3)]);



% Plot basis function
disp('hit any key to plot the basis function');
pause
fig6 = figure;
grid on
q = Polyhedron('ub',[1;1],'lb',-[1;1]);
q.addFunction(@(x) prod(1-abs(x)),'secondOrderInterp');
q.fplot('secondOrderInterp','grid',25,'linewidth',line_width,'color',[0.9 0.9 0.9],'showgrid',true);

set(gca,'LineWidth',axeswidth)
set(gca,'FontSize', tick_font_size);
set(gca,'XTick',[-1 0 1]);
set(gca,'YTick',[-1 0 1]);
set(gca,'ZTick',0:0.2:1);
axis([-1.2 1.2 -1.2 1.2 0 1]);


hx6 = xlabel('$x_1$');
set(hx6, 'FontSize', label_font_size);
hy6 = ylabel('$x_2$');
set(hy6, 'FontSize', label_font_size);
hz6 = zlabel('$\phi_{0,1}(x)$');
set(hz6, 'FontSize', label_font_size);
pos = get(gca,'Position');
set(gca,'Position',[0.17 0.17 pos(3:4)]);
posz6 = get(hz6,'Position');
set(hz6,'Position',[posz6(1)-0.2 posz6(2:3)]);



%% save book figures
disp('Hit any key to save the figures.');
pause
% save
saveas(fig1, [figure_name1,'_matlab'], 'fig');

% print 
laprint(fig1, [figure_name1,'_tex'],'scalefonts','off');

% save
saveas(fig2, [figure_name2,'_matlab'], 'fig');

% print 
laprint(fig2, [figure_name2,'_tex'],'scalefonts','off','figcopy','off');

% save
saveas(fig3, [figure_name3,'_matlab'], 'fig');

% print 
laprint(fig3, [figure_name3,'_tex'],'scalefonts','off','figcopy','off');

% save
saveas(fig4, [figure_name4,'_matlab'], 'fig');

% print 
laprint(fig4, [figure_name4,'_tex'],'scalefonts','off','figcopy','off');

% save
saveas(fig5, [figure_name5,'_matlab'], 'fig');

% print 
laprint(fig5, [figure_name5,'_tex'],'scalefonts','off','figcopy','off');

% save
saveas(fig6, [figure_name6,'_matlab'], 'fig');

% print 
laprint(fig6, [figure_name6,'_tex'],'scalefonts','off','figcopy','off');

end

function dat = defineEx(soft)
%
% Define optimizer for double integrator
%
% C.N. Jones, EPFL, 2013
% Revised by M. Herceg, ETH Zurich 2014
%
% requires MPT3, YALMIP


if nargin < 1
  soft = false;
end

% Define an LTI dynamics x^+ = A x + B u
% (note that we can ommit the output equation)
A = [1 1; 0 1];
B = [1; 0.5];
lti = LTISystem('A', A, 'B', B);

% Add constraints on predicted states
lti.x.min = [-5; -5];
lti.x.max = [5; 5];

if soft
  lti.x.with('softMax');
  lti.x.softMax.maximalViolation = 100*[1;1];
  lti.x.softMax.penalty = QuadFunction(eye(2));
  lti.x.with('softMin');
  lti.x.softMin.maximalViolation = 100*[1;1];
  lti.x.softMin.penalty = QuadFunction(eye(2));
end

% Add constraints on predicted control inputs
lti.u.min = -1;
lti.u.max = 1;

% Use quadratic state penalty with identity weighting matrix
lti.x.penalty = QuadFunction(0.1*eye(2));
lti.u.penalty = QuadFunction(1);


% Define an MPC controller using "lti" as the prediction model
ctrl = MPCController(lti, 10);

% export to YALMIP 
dat = ctrl.toYALMIP;
dat.ctrl = ctrl;
dat.lti  = lti;

end

function [box,err] = waveletApprox(box, dat)
%
% Compute error within the box, and split if there's a problem
%
% C.N. Jones, EPFL, 2013
% Revised by M. Herceg, ETH Zurich 2014
%
% requires MPT3, YALMIP


% Convert to vertices
V(:,1) = [0 0 1 0;1 0 0 0;1 0 0 0;0 0 1 0]*[box.ub;box.lb]*0.9999;
V(:,2) = [0 1 0 0;0 1 0 0;0 0 0 1;0 0 0 1]*[box.ub;box.lb]*0.9999;
box.V   = V;

hold on
h=Polyhedron(V).plot('color',[0.9 0.9 0.9]);
set(h,'FaceColor','none');
drawnow

% Compute the values at the vertices if we don't know them
J = nan*ones(4,1); box.J = J;
for i = 1:length(box.J)
  if isnan(box.J(i))
    J(i,1) = dat.theFunc(V(i,:)');
  end
end
box.J   = J;

% Check if box is inside the terminal set
if all(dat.termSet.contains(V'))
%   fprintf('Terminating on terminal set containment\n');
  set(h,'facecolor',[0.7 0.7 0.7]);
  box.err = 0;
  err = 0;
  return
end

% Compute the maximum error in this region
if all(isnan(J)) % Whole box is outside
  set(h,'facecolor',[0.8 0.8 0.8]);
  box.err = 0;
  err = 0;
  return
end

if prod(box.ub-box.lb) < 1e-2
%   fprintf('Terminating on minimum box size\n');
  set(h,'facecolor',[0.9 0.9 0.9]);
  box.err = 0;
  err = 0;
  return
end

if any(isnan(J)) % Some feasible, some not
  % Take the soft-constrained value for the infeasible points
  for i = 1:length(box.J)
    if isnan(box.J(i))
      J(i,1) = dat.softOpt{V(i,:)'};
    end
  end
  box.J   = J;  
end


if any(isnan(J))
  box.err = -inf;
else
  % Compute error
  x = dat.variables.x(:,1);
  obj = dat.objective + dat.epsilon*x'*dat.Q*x - J'*dat.lam;
  con = dat.constraints + set(dat.lam >= 0) + set(sum(dat.lam) == 1);
  con = con + set(x == V'*dat.lam);
  
  dd = solvesdp(con, obj, sdpsettings('verbose',0));
  if dd.problem
    warning('Could not solve error computation problem')
    set(h,'facecolor',[0.8 0.8 0.8]);
    box.err = 0;
    err = 0;
    return
  end
  
  box.err = double(obj);
end

if box.err >= 0
  err = box.err;
  set(h,'facecolor',[0.7 0.7 0.7]);
  return
else
  % Split box and recursively compute approximations
  mn = mean([box.lb box.ub],2);
  b(1).lb = box.lb;
  b(1).ub = mn;
  
  b(2).lb = mn;
  b(2).ub = box.ub;
  
  b(3).lb = [mn(1);box.lb(2)];
  b(3).ub = [box.ub(1);mn(2)];
  
  b(4).lb = [box.lb(1);mn(2)];
  b(4).ub = [mn(1);box.ub(2)];
  
  
  for i = 1:4
    [box.children{i},childErr] = waveletApprox(b(i), dat);
    err = -inf;
    if childErr > -inf
      err = max([childErr, err]);
    end
  end
end

end

function J = optFunc(x, dat)
%
% Compute the function that we're trying to evaluate
%
% J(x) = Jstar(x) - x'*P*x
% 
% C.N. Jones, EPFL, 2013
% Revised by M. Herceg, ETH Zurich 2014
%
% requires MPT3, YALMIP

[J,isProblem] = dat.opt{x(:)};
% if isProblem, error('Could not solve optimization'); end
% if isProblem, fprintf('Could not solve optimization\n'); end
if iscell(J(1))
  J = J{1};
end

end

function P = flattenWavelet(box)
%
% Compute a PolyUnion from the wavelet tree
%
% C.N. Jones, EPFL, 2013
%
% requires MPT3, YALMIP


exp = box;
leaves = [];
k = 1;
while ~isempty(exp)
%   fprintf('Length of exp = %i\n', length(exp));
  e = exp(1); exp(1) = [];
  
  for i = 1:length(e.children)
    if isfield(e.children{i}, 'children')
      exp = [exp [e.children{i}]];
    else
      leaves = [leaves;e.children{i}];
      P(k) = Polyhedron(e.children{i}.V);
      P(k).Data = e.children{i};
      P(k).addFunction(@(x) evalInterpolet(x, P(k).Data), 'cost');
      k = k + 1;
    end
  end  
end

% P = PolyUnion(P);

end

function J = evalPWInterpolet(x, W, func)
% evaluate interpolet over W regions

x = x(:);
[isin,inwhich,closest]=W.Set.isInside(x);
if ~isin
    inwhich = closest;
    %error('Point is not covered by the set W')
end

i=inwhich(1);
J = evalInterpolet(x,W.Set(i).Data,func);


end

function J = evalInterpolet(x, dat, func)

% V(:,1) = [0 0 1 0;1 0 0 0;1 0 0 0;0 0 1 0]*[dat.ub;dat.lb];
% V(:,2) = [0 1 0 0;0 1 0 0;0 0 0 1;0 0 0 1]*[dat.ub;dat.lb];

x = x(:);

% Compute percentage from lower left
per = (x - dat.lb) ./ (dat.ub - dat.lb);

% Percentage distance from each corner
% top left, top right, bottom right, bottom left
per = [...
  per(1)   1-per(2);...
  1-per(1) 1-per(2);...
  1-per(1) per(2);...
  per(1)   per(2)];
per = 1-per;
per = prod(per,2);

if func == 'U'
  J = per'*dat.U;
else
  J = per'*dat.J;
end

end