%% Set up basic symbols.
% Use sprintf() and char(sym()) to automate.

syms x dx ddx a da dda
    
%% Physical parameters of the segway - masses, inertias, lengths, gravity
syms Mw Iw Rw Hr Ir Mr g

%% Wheel kinetic and potential energy
Tw = (1/2) * Mw * dx^2 + (1/2)*Iw*(dx/Rw)^2
Vw = sym('0')

%% Rider kinetic and potential energy
Tr = (1/2)*Ir*(da^2) + (1/2)*Mr*((dx+Hr*da*cos(a))^2+(Hr*da*sin(a))^2)
Vr = Mr * g * Hr * cos(a)


%% The Lagrangian of the segway
T = Tw + Tr
V = Vw + Vr
L = T - V

%% Differentiate the second term of the lagrangian.
dLdx = diff(L, x)
dLda = diff(L, a)

%% Inner differentiation of the first term of the lagrangian.
dLdxdot = diff(L, dx)
dLdadot = diff(L, da)

%% Differentiate ugly first terms with respect to time.
% Matlab doesn't know that we would like e.g. dx = diff(x,t)
% so we can't do the following
%dtdLdxdot = diff(dLdxdot, t)
%dtdLdadot = diff(dLdadot, t)
%
% Instead we use the chain rule:
% df/dt = df/dx * dx/dt
% In Matlab, that's
% diff(d,x) * dx
dtdLdxdot = diff(dLdxdot, x) * dx + diff(dLdxdot, dx) * ddx ...
    + diff(dLdxdot, a) * da + diff(dLdxdot, da) * dda;
% Or just the use the jacobian to do it more compactly.
dtdLdadot = jacobian( dLdadot, [ x; a; dx; da ] ) * [ dx; da; ddx; dda ];

%% Form expressions for generalized forces

qx = dtdLdxdot - dLdx
qa = dtdLdadot - dLda

%% Gather into matrix form
% Forces are linear in ddx, dda
qx = expand(simplify(qx))
qa = expand(simplify(qa))
MM = sym(zeros(2));

%% Demonstrate use coeffs to get mass matrix
qxtmp = qx;
[Cc,Tc] = coeffs(qxtmp, ddx);
if length(Cc) == 2 && Tc(2) == ddx
  MM(1,1) = Cc(2);
  qxtmp = expand(qxtmp - Cc(2) * ddx);
end

%% Get mass matrix and rest in a loop.
qs = { qx, qa };
ddvars = { ddx, dda };
for i = 1:2
    for j = 1:2
        [Cc,Tc] = coeffs(qs{i}, ddvars{j});
        if length(Cc) == 2 && Tc(2) == ddvars{j}
          MM(i,j) = simplify(Cc(2));
          qs{i} = expand(qs{i} - Cc(2) * ddvars{j});
        end
    end
end

rest = [ simplify(qs{1}); simplify(qs{2}) ];