function [q, p, u] = varInt(S, q0, p0, qT, pT, T, N, qig, uig, options)
%VARINT integrates a finite dimensional mechanical system or an optimal
% control problem governed by such a mechanical system.
%
%   VARINT integrates a mech. system governed by a Lagrangian and a force
%
%       L :  TQ  --->  R                  f :  TQ  --->  T*Q
%          (q,v) |-->  L(q,v)                (q,v) |-->  f(q,v)
%
%   where Q is the finite dimensional configuration space.
%
%   VARINT solves the OCP given by the objective functional
%
%       J(q(.),u(.)) := Int_0^T C(q(t),v(t),u(t)) dt
%
%   associated to the cost function
%
%       C : TQxU  --->  R
%         (q,v,u) |-->  C(q,v,u)
%
%   and subject to the mechanical system governed by the Lagrangian and
%   control force
%
%       L :  TQ  --->  R              f : RxTQxU  --->  T*Q
%          (q,v) |-->  L(q,v)           (t,q,v,u) |-->  f(t,q,v,u)
%
%   where Q is the finite dimensional configuration space.   
%
%   VARINT implements a Galerkin method: polynomial approximations for 
%   the mechanical trajectories and numerical quadrature for the integral
%   action as for the cost function. More precisely, VARINT uses Lagrange
%   polynomials and Lobatto's quadrature.
%   For more details, refer to [1] and [2].
%
%   Inputs:
%
%       S   Mechanical system parameters (see below).
%      q0   Initial point (column vector).
%      p0   Initial momentum (idem).
%      qT   Final point (idem).
%      pT   Final momentum (idem).
%       T   Final time or time length (initial time is 0).
%      N+1  Number of macronodes (q0 and qT included, N>0).
%      qig  Initial guess for the trajectory.
%      uig  Initial guess for the controls.
%      options  Option structure.
%
%   Outputs:
%
%        q  Computed trajectory.
%        p  Computed momenta for the trajectory.
%        u  Computed optimal controls.
%
%   Details:
%
%       S is a structure that must have at least the following elements:
%
%           S.DqL  A function reference to the jacobian of L w.r.t q.
%           S.DvL  A function reference to the jacobian of L w.r.t v.
%           S.f    A function reference to the external force.
%           S.cost A function reference to the cost function.
%
%       qig and uig follow the same structure than q and u (below).
%
%       options is a structure passed to FSOLVE/FMINCON that furthermore
%         has the following fields:
%
%           options.s  Order of the polynomial aproximation (0<s<6). They
%              are determined by s+1 points, i.e. s-1 micronodes and 2
%              macronodes.
%           options.data Type of the boundary conditions, posible values
%              are 'momentum' (default), 'velocity' and 'disrete'.
%                'momentum' (q0,p0,qT,pT) are indeed (q0,p0,qT,pT).
%                'velocity' (q0,p0,qT,pT) are actually (q0,v0,qT,vT).
%                'momentum' (q0,p0,qT,pT) are actually (q0,q1,qN-1,qN).
%
%       q is a (dim,s,N+1)-tensor matrix where q(:,1,i) is the i-th point
%         of the discrete trajectory and q(:,j+1,i) is the j-th midpoint
%         at the i-th step. Note that q(:,2:s,N+1) are not be considered.
%       p is a (dim,N+1)matrix where p(:,i) is the momenta associated to
%         q(:,1,i).
%       u is a (dim,s,N+1)-tensor matrix where u(:,j,i) is the control
%         acting on q(:,j,i).
%
%   Calls:
%
%      varInt(S,q0,p0,qT,pT,T,N,qig,uig,options*)
%        Standard call.
%
%      varInt(S,q0,p0,qT,pT,T,N,options*)
%        Initial guesses are assumed to be empty (qig=[], uig=[]).
%
%      varInt(S,N,qig,uig,options*)
%        The system elements (q0,p0,qT,pT,T) are included in the system
%        structure S itslef as fields (S.q0, S.p0, ...).
%
%      varInt(S,N,options*)
%        The system elements (q0,p0,qT,pT,T) are included in the system
%        structure S itslef as fields (S.q0, S.p0, ...).
%        Initial guesses are assumed to be empty (qig=[], uig=[]).
%
%      * The option structure is always optional.
%
%   Modes: depending on the system inputs within the call (explicit or
%   implicitely given in S), the function works in two different modes,
%   mechanical integration or optimal control. If doubts, the 'standard
%   call' with 'momentum' boundary constraints is recommended.
%
%      (q0,p0,qT,pT) control integration.
%      (q0,p0,[],[]) forward mechanical integration.
%      ([],[],qT,pT) backward mechanical integration.
%      (q0,[],qT,[]) mechanical integration.
%
%
%
%   Refer to [1], section 2.6.6, and to [2].
%
%   [1] J.E. Marsden, M. West: "Discrete mechanics and variational
%       integrators", Acta Numerica (2001).
%   [2] S. Ober-Bloebaum, O. Junge, J.E. Marsden: "Discrete Mechanics and
%       Optimal Control: an Analysis", ESAIM: COCV (2011)
%
%   See also FSOLVE, FMINCON, OPTIMSET.

%   Copyleft 2011-2013 Cedric M. Campos, Dr., joint work with
%                      Oliver Junge, Prof. Dr., and
%                      Sina Ober-Bloebaum, Prof. Dr.
%                      Emmanuel Trélat, Prof. Dr.
%   $Revision: 0.2.0.00 $  $Date: 2013/02/02 16:19:11 $

% narg{,in,out}{,chk} error
switch nargin
    case {2,2+1} %varInt(S,N,options)
         N = q0;
        if nargin == 3
            options = p0;
        end
        q0 = S.q0;
        p0 = S.p0;
        qT = S.qT;
        pT = S.pT;
         T = S.T;
        qig = [];
        uig = [];
    case {4,4+1} %varInt(S,N,qig,uig,options)
         N = q0;
        qig = p0;
        uig = qT;
        if nargin == 5
            options = pT;
        end
        q0 = S.q0;
        p0 = S.p0;
        qT = S.qT;
        pT = S.pT;
         T = S.T;
    case {7,7+1} %varInt(S,q0,p0,qT,pT,T,N,options)
        if nargin == 8
            options = qig;
        end
        qig = [];
        uig = [];
    case {9,9+1} %varInt(S,q0,p0,qT,pT,T,N,qig,uig,options)
    otherwise
        error('Wrong number of arguments.');
end

if ~isfield(options,'s') || isempty(options.s)
    options.s = 5;
end
s = options.s;

if ~isfield(options,'data') || isempty(options.data)
    options.data = 'momentum';
else
    data = options.data;
    if ~(strcmp(data,'momentum') || strcmp(data,'velocity') || strcmp(data,'discrete'))
        error('options.data not properly set.');
    end
end

% ------------Initialization----------------
global lobattoPoint   % hack: overrides unnecessary computations
global lobattoWeight
global lagrangeCoeffX

             n = s+1;
         [x,w] = lobatto(n);
             x = (x+1)/2;
  lobattoPoint = x;
 lobattoWeight = w/2;
lagrangeCoeffX = zeros(n,n,2);
for i = 1:n
    for j = 1:n
        lagrangeCoeffX(i,j,false+1) = lagrangeCoeff(x, i, x(j), false);
        lagrangeCoeffX(i,j,true +1) = lagrangeCoeff(x, i, x(j), true);
    end
end

if ~(isempty(q0) || isempty(p0))
    switch options.data
        case 'velocity' % p0 = v0
            p0 = S.DvL(q0,p0,S);
        case 'discrete' % p0 = q1
            q0_0_s        = zeros(dim,s,2);
            q0_0_s(:,:,1) = q0*ones(1,s);
            q0_0_s(:,1,2) = p0;
            if ~(isempty(qig))
                q0_0_s = qig(:,:,1:2);
            end
            [~, p] = varInt_SI(S, q0, p0, h, 1, q0_0_s, options);
            p0 = p(:,1);
    end
end

if ~(isempty(qT) || isempty(pT))
    switch options.data
        case 'velocity' % pT = vT
            pT = S.DvL(qT,pT,S);
        case 'discrete' % qT = q(T-h); pT = qT
            q0_0_s        = zeros(dim,s,2);
            q0_0_s(:,:,1) = qT*ones(1,s);
            q0_0_s(:,1,2) = pT;
            if ~(isempty(qig))
                q0_0_s = qig(:,:,end-1:end);
            end
            [~, p] = varInt_SI(S, qT, pT, h, 1, q0_0_s, options);
            pT = p(:,2);
    end
end

options.data = 'momentum';

if ~(isempty(q0) || isempty(p0) || isempty(qT) || isempty(pT))
    disp('optimal control problem')
    [q, p, u, ~, ~, ~, ~] = varInt_OC(S, q0, p0, qT, pT, T, N, qig, uig, options);
elseif ~(isempty(q0) || isempty(p0))
    disp('forward integration')
    [q, p] = varInt_FW(S, q0, p0, T, N, options);
elseif ~(isempty(qT) || isempty(pT))
    disp('backward integration')
    [q, p] = varInt_BW(S, qT, pT, T, N, options);
elseif ~(isempty(q0) || isempty(qT))
    disp('single integration')
    [q, p] = varInt_SI(S, q0, qT, T, N, qig, options);
else
    disp('free')
    disp('not yet implemented')
end

end

%%%%%
% The following set of functions, varInt_{FW,BW,SI,OC} are the main
% algorithms of integration. Note that varInt_OC uses the same functions as
% varInt_{FW,BW,SI} do, but with slight variations (control inclusion).

function [q, p] = varInt_FW(S, q0, p0, T, N, options)
dim = S.dimQ;
  s = options.s;
  h = T/double(N);

q = zeros(dim, s, N+1);
p = zeros(dim,    N+1);

ti     = 0;
qi_0   = q0;
qi_1_s = q0*ones(1,s);
pi     = p0;
qi_1_s = fsolve(@(qi_1_s)DEL(ti,[qi_0,qi_1_s],pi,h,S), qi_1_s, options);

q(:,:,1)   = [qi_0,qi_1_s(:,1:s-1)];
q(:,1,2)   = qi_1_s(:,s);
p(:,  1)   = pi;
p(:,  2)   = DELi(s+1, ti, [qi_0,qi_1_s], h, S);

for i = 1+1:N
        ti = (i-1)*h;
      qi_0 = q(:,0+1,i);
        pi = p(:,    i);
        vi = (q(:,1,i) - q(:,1,i-1))/h;
    qi_1_s = initialGuess(qi_1_s,vi,h);
    qi_1_s = fsolve(@(qi_1_s)DEL(ti,[qi_0,qi_1_s],pi,h,S), qi_1_s, options);
    q(:,2:s,i)   = qi_1_s(:,1:s-1);
    q(:,0+1,i+1) = qi_1_s(:,s);
    p(:,    i+1) = DELi(s+1, ti, [qi_0,qi_1_s], h, S);
end
end

function [q, p] = varInt_BW(S, qT, pT, T, N, options)
% self-adjoint method
[q, p] = varInt_FW(S, qT, -pT, T, N, options);

q =  q(:,end:-1:1,end:-1:1);
p = -p(:,         end:-1:1);
end

function [q, p] = varInt_SI(S, q0, qT, T, N, qig, options)
dim = S.dimQ;
  s = options.s;
  h = T/double(N);

q = reshape(qig(:,:,1:N),[dim,s*N]);
q = q(:,2:end);
qN_0_s = [qT, zeros(dim,s-1)];
q = fsolve(@(q)DELFull(reshape([q0,q,qN_0_s],[dim,s,N+1]),h,S),q,options);
q = reshape([q0,q,qN_0_s],[dim,s,N+1]);

p = zeros(dim,N+1);
for i = 0+1:N
    ti = (i-1)*h;
    p(:,i) = -DELi(1,ti,[q(:,:,i),q(:,0+1,i+1)],h,S); %pre-momenta
end
p(:,N+1) = DELi(s+1,ti,[q(:,:,N),q(:,0+1,N+1)],h,S); %post-momentum
end

function [q,p,u,fval,exitflag,output,lambda] = varInt_OC(S,q0,p0,qT,pT,...
                                                       T,N,qig,uig,options)
dQ = S.dimQ;
dU = S.dimU;
 s = options.s;
 h = T/double(N);

% initial guess
if isempty(qig)
    Suncntr = S;
    Suncntr.f = Suncntr.funcontrolled;
    [q, ~] = varInt(Suncntr, q0, p0, T, N, s, options);
    q = [reshape(q(:,:,1:N),dQ,s*N), q(:,1,end)];
else
    q = [reshape(qig(:,:,1:N),dQ,s*N), qig(:,1,end)];
end
if isempty(uig)
    u = zeros(dU, s*N+1);
else
    u = [reshape(uig(:,:,1:N),dQ,s*N), uig(:,1,end)];
end

[x,fval,exitflag,output,lambda] = ...
    fmincon(@(x)discrObjFun(reshape([x(1:dQ,:),zeros(dQ,s-1)], ...
                                    [dQ, s, N+1]), ...
                            reshape([x(dQ+1:dQ+dU,:),zeros(dU,s-1)], ...
                                    [dU , s, N+1]), s,N,T,S), ...
            [q;u], [], [], [], [], [], [], ...
            @(x)constrFun(reshape([x(1:dQ,:),zeros(dQ,s-1)], ...
                                  [dQ, s, N+1]), ...
                          reshape([x(dQ+1:dQ+dU,:),zeros(dU,s-1)], ...
                                  [dU , s, N+1]), ...
                          q0,p0,qT,pT,h,s,N,S), options);

q = reshape([x(1:dQ,:),zeros(dQ,s-1)], [dQ, s, N+1]);
u = reshape([x(dQ+1:dQ+dU,:),zeros(dU,s-1)], [dU , s, N+1]);
p = zeros(dQ, N+1);
p(:,0+1) = p0;
for i=1:N;
    ti = (i-1)*h;
    qi_0_s = [q(:,:,i), q(:,0+1,i+1)];
    ui_0_s = [u(:,:,i), u(:,0+1,i+1)];
    p(:,i+1) = DELi_OC(s+1, ti, qi_0_s, ui_0_s, h, S);
end
end

%%%%%
% Block of functions used by varInt_{FW,BW,SI}.
% - DEL consists of the general set of equations from node to node.
% - DELFull consists of DEL along the whole trajectory.
% - DELi is the equation at a single micronode, wich is specific to the
%   chosen quadrature rule and polynomial class.

function q1_1_s = initialGuess(q0_1_s, v0, h)
% basically: q1_1_s = q0_1_s + h*v0
% note: q0_0 is not explicitly used here (but in v0)

q1_1_s = zeros(size(q0_1_s));
for i = 1:size(q0_1_s,2)
    q1_1_s(:,i) = q0_1_s(:,i) + h*v0;
end
end

function f = DELFull(q,h,S)
[dim, s, N] = size(q);
          N = N-1;

f = zeros(dim,s,N);
t = 0:h:N*h;

      qk_0_s = [q(:,:,0+1),q(:,0+1,1+1)];
%f(:,0+1,0+1) = qk_0_s(:,0+1,0+1)-S.q0;
for i = 1+1:s
    f(:,i,1) =  DELi(i,t(0+1),qk_0_s,h,S);
end
for k = 1+1:N
      qk_1_0_s = qk_0_s;
      qk_0_s   = [q(:,:,k),q(:,0+1,k+1)];
    f(:,0+1,k) = DELi(s+1,t(k-1),qk_1_0_s,h,S) + DELi(1,t(k),qk_0_s,h,S);
    for i = 1+1:s
        f(:,i,k) = DELi(i,t(k),qk_0_s,h,S);
    end
end
f = reshape(f,[dim,s*N]);
f = f(:,2:end);
end

function f = DEL(t0,q0_0_s,p0,h,S)
%        f == 0 are the (dim,s)-DEL equations
%   f(:,1) == 0 are the dim-equations from the discrete mechanics
% f(:,2:s) == 0 are the (dim,s-1) equations from the quadrature scheme

[dim, s] = size(q0_0_s);
       s = s-1;

f = zeros(dim,s);
f(:,0+1) =  DELi(1,t0,q0_0_s,h,S) + p0;
for i = 1+1:s
    f(:,i) =  DELi(i,t0,q0_0_s,h,S);
end
end

function f = DELi(i,t,q,h,S)
% i-th part of the DEL equations (with 0+1<=i<=s+1)

global lobattoPoint
global lobattoWeight
global lagrangeCoeffX

n = size(q,2);
x = lobattoPoint;
w = lobattoWeight;

f=0;
for j=1:n
    tj = t + x(j);
     P = q(:,j); %= lagrangePolyX(q,j,false);
    DP = lagrangePolyX(q,j,true)/h;
     C = lagrangeCoeffX(i,j,false+1);
    DC = lagrangeCoeffX(i,j,true +1);
    f = f + w(j)*(   h*S.DqL(P,DP,S)*C  ...
                   +   S.DvL(P,DP,S)*DC ...
                   + h*S.f(P,DP,S)*C  );
end
end


%%%%%
% Block of functions used by varInt_{FW,BW,SI}.
% - discrObjFun is the discrete objective function.
% - constrFun is the collection of dynamical and boundary constraints.
% - {DELFull,DEL,DELi}_OC are the specific versions for the control case of
%   the previous functions DELFull,DEL,DELi.

% function f = initialGuessCtrlTrjct()
% something intelligent could be implemented
% end

function f = discrObjFun(q,u,s,N,T,S)

global lobattoPoint
global lobattoWeight

x = lobattoPoint;
w = lobattoWeight;

h = T/double(N);
t = 0:h:T;

f=0;
for i=1:N
    qi_0_s = [q(:,:,i), q(:,0+1,i+1)];
    ui_0_s = [u(:,:,i), u(:,0+1,i+1)];
    for j=1:s+1
         tij = t(i) + x(j);
         qij = qi_0_s(:,j); %= lagrangePolyX(qi_0_s,j,false);
        Dqij = lagrangePolyX(qi_0_s,j,true)/h;
         uij = ui_0_s(:,j); %= lagrangePolyX(ui_0_s,j,false);
           f = f + w(j)*S.cost(tij,qij,Dqij,uij,S);
    end
end
f = h*f;
end

function [empty, f] = constrFun(q,u,q0,p0,qT,pT,h,s,N,S)

    constrDEL_aux = DELFull_OC(q,u,p0,h,s,N,S);
    [dim, s, N] = size(constrDEL_aux);
    constrDEL = reshape(constrDEL_aux,dim*s*N,1);
    q0_0_s = [q(:,:,0+1), q(:,0+1,1+1)];
%   u0_0_s = [u(:,:,0+1), u(:,0+1,1+1)];
    constrInit_q0 = q0_0_s(:,0+1) - q0;
%   constrInit_p0 = DELi_OC(0+1,0,q0_0_s,u0_0_s,h,S) + p0;
    qN1_0_s = [q(:,:,N),   q(:,0+1,N+1)];
    uN1_0_s = [u(:,:,N),   u(:,0+1,N+1)];
    constrFinal_qT = qN1_0_s(:,s+1) - qT;
    constrFinal_pT = DELi_OC(s+1,(N-1)*h,qN1_0_s,uN1_0_s,h,S) - pT;

    f = [constrInit_q0; constrDEL; constrFinal_qT; constrFinal_pT];
    % NOTE: p0 is "already" enforced by DEL
    empty =[];
end

function f = DELFull_OC(q,u,p0,h,s,N,S)
% f == 0 is the full set of DEL equations along the optimal trajectory

dim = size(p0,1);
f = zeros(dim,s,N);
pi = p0;
for i = 1:N
    ti = (i-1)*h;
    qi_0_s = [q(:,:,i), q(:,0+1,i+1)];
    ui_0_s = [u(:,:,i), u(:,0+1,i+1)];
    f(:,:,i) = DEL_OC(ti, qi_0_s, ui_0_s, pi, h, S);
    pi = DELi_OC(s+1, ti, qi_0_s, ui_0_s, h, S);
end
end

function f = DEL_OC(t0,q0_0_s,u0_0_s,p0,h,S)
%        f == 0 are the (dim,s)-DEL equations
%   f(:,1) == 0 are the dim-equations from the discrete mechanics
% f(:,2:s) == 0 are the (dim,s-1) equations from the quadrature scheme

[dim, s] = size(q0_0_s);
       s = s-1;

f = zeros(dim,s);
f(:,0+1) =  DELi_OC(1,t0,q0_0_s,u0_0_s,h,S) + p0;
for i = 1+1:s
    f(:,i) =  DELi_OC(i,t0,q0_0_s,u0_0_s,h,S);
end
end

function f = DELi_OC(i,t,q,u,h,S)
% i-th part of the DEL equations (with 1<=i<=s)

global lobattoPoint
global lobattoWeight
global lagrangeCoeffX

n = size(q,2);
x = lobattoPoint;
w = lobattoWeight;

f=0;
for j=1:n
    tj = t + x(j);
     P = q(:,j); %= lagrangePolyX(q,j,false);
    DP = lagrangePolyX(q,j,true)/h;
     C = lagrangeCoeffX(i,j,false+1);
    DC = lagrangeCoeffX(i,j,true +1);
     U = u(:,j); %= lagrangePolyX(u,j,false);
     f = f + w(j)*(   h*S.DqL(P,DP,S)*C  ...
                    +   S.DvL(P,DP,S)*DC ...
                    + h*S.f(tj,P,DP,U,S)*C  );
end
end

%%%%%
% Block of "utilitary" functions.

function [x,w] = lobatto(n)

w = zeros(1,n);
x = zeros(1,n);

switch n
    case 2
        x(1) = -1;
        x(2) =  1;

        w(1) =  1;
        w(2) =  1;

    case 3
        x(1) = -1;
        x(2) =  0;
        x(3) =  1;

        w(1) = 1/3;
        w(2) = 4/3;
        w(3) = 1/3;

    case 4
        x(1) = -1;
        x(2) = -sqrt(5)/5;
        x(3) =  sqrt(5)/5;
        x(4) =  1;

        w(1) = 1/6;
        w(2) = 5/6;
        w(3) = 5/6;
        w(4) = 1/6;

    case 5
        x(1) = -1;
        x(2) = -sqrt(21)/7;
        x(3) =  0;
        x(4) =  sqrt(21)/7;
        x(5) =  1;

        w(1) = 1/10;
        w(2) = 49/90;
        w(3) = 32/45;
        w(4) = 49/90;
        w(5) = 1/10;

    case 6
        x(1) = -1;
        x(2) = -sqrt((7+2*sqrt(7))/21);
        x(3) = -sqrt((7-2*sqrt(7))/21);
        x(4) =  sqrt((7-2*sqrt(7))/21);
        x(5) =  sqrt((7+2*sqrt(7))/21);
        x(6) =  1;

        w(1) = 1/15;
        w(2) = (14-sqrt(7))/30;
        w(3) = (14+sqrt(7))/30;
        w(4) = (14+sqrt(7))/30;
        w(5) = (14-sqrt(7))/30;
        w(6) = 1/15;
end
end

function f = lagrangePolyX(y, j, flag)
% flag == true : derivative w.r.t. time

global lagrangeCoeffX

[dim,n] = size(y);

f = zeros(dim,1);
for i=1:n
    f = f + y(:,i)*lagrangeCoeffX(i,j,flag+1);
end
end

function f = lagrangePoly(x, y, t, flag)
% flag == true : derivative w.r.t. time

[~, n] = size(x);

f = 0;
for i=1:n
    f = f + y(:,i)*lagrangeCoeff(x,i,t,flag);
end
end

function f = lagrangeCoeff(x, i, t, flag)
% flag == true : derivative w.r.t. time

[~, n] = size(x);

if flag==false
    f = 1;
    for j = 1:n
        if j~=i
            f = f * (t-x(j)) / (x(i)-x(j));
        end
    end
elseif flag==true
%     f = 0;
%     for k = 1:n
%         if k~=i
%             g = 1/(x(i)-x(k));
%             for j = 1:n
%                 if (j~=i) && (j~=k)
%                     g = g * (t-x(j)) / (x(i)-x(j));
%                 end
%             end
%             f = f + g;
%         end
%     end
    numer = 0;
    denom = 1;
    for k = 1:n
        if k~=i
            term = 1;
            for j = 1:n
                if (j~=i) && (j~=k)
                    term = term * (t-x(j));
                end
            end
            numer = numer + term;
            denom = denom * (x(i)-x(k));
        end
    end
    f = numer/denom;
end
end