%% Set up basic symbols.
% Use sprintf() and char(sym()) to automate.
t = sym('t');

x = sym('x(t)');
dx=diff(x,t);
ddx=diff(dx,t);

a = sym('a(t)');
da = diff(a,t);
dda = diff(da,t);

prettyx = sym('x');
prettydx = sym('dx');
prettyddx = sym('ddx');
prettya = sym('a');
prettyda = sym('da');
prettydda = sym('dda');


%% Set up display symbols
% Use sprintf() and char(sym()) to automate.
ddprettyvars = { prettyddx, prettydda };
dprettyvars = { prettydx, prettyda };
prettyvars = { prettyx, prettya };
%allprettyvars = { prettydda, prettyddx, prettyda, prettydx, prettya, prettyx };

dduglyvars = { ddx, dda };
duglyvars = { dx, da };
uglyvars = { x, a };
%alluglyvars = { dda, ddx, da, dx, a, x };


%% 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

% Matlab can't differentiate with respect to a variable in function(ugly)
% form, e.g. sym('x(t)')
% CAUTION: Must translate from ugly accelerations first to get
% substitutions right.
prettyL = L;
prettyL = subs(prettyL, dduglyvars, ddprettyvars);
prettyL = subs(prettyL, duglyvars, dprettyvars);
prettyL = subs(prettyL, uglyvars, prettyvars);

%% Differentiate the second term of the lagrangian.
dLdx = diff(prettyL, prettyx)
dLda = diff(prettyL, prettya)

%% Inner differentiation of the first term of the lagrangian.
dLdxdot = diff(prettyL, prettydx)
dLdadot = diff(prettyL, prettyda)

%% Uglify first term so Matlab knows x, a, etc. are functions of time.
% (not constants)
% CAUTION: must translate this time from pretty positions first to get
% substitutions right
uglydLdxdot = dLdxdot;
uglydLdxdot = subs(uglydLdxdot, prettyvars, uglyvars)
uglydLdxdot = subs(uglydLdxdot, dprettyvars, duglyvars);
uglydLdxdot = subs(uglydLdxdot, ddprettyvars, dduglyvars);

uglydLdadot = dLdadot;
uglydLdadot = subs(uglydLdadot, prettyvars, uglyvars)
uglydLdadot = subs(uglydLdadot, dprettyvars, duglyvars);
uglydLdadot = subs(uglydLdadot, ddprettyvars, dduglyvars);


%% Differentiate ugly first terms with respect to time.
uglydtdLdxdot = diff(uglydLdxdot, t)
uglydtdLdadot = diff(uglydLdadot, t)

%% Prettify first terms
dtdLdxdot = uglydtdLdxdot;
dtdLdxdot = subs(dtdLdxdot, dduglyvars, ddprettyvars);
dtdLdxdot = subs(dtdLdxdot, duglyvars, dprettyvars);
dtdLdxdot = subs(dtdLdxdot, uglyvars, prettyvars);

dtdLdadot = uglydtdLdadot;
dtdLdadot = subs(dtdLdadot, dduglyvars, ddprettyvars);
dtdLdadot = subs(dtdLdadot, duglyvars, dprettyvars);
dtdLdadot = subs(dtdLdadot, uglyvars, prettyvars);

%% Form pretty 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, prettyddx);
if length(Cc) == 2 && Tc(2) == prettyddx
  MM(1,1) = Cc(2);
  qxtmp = expand(qxtmp - Cc(2) * prettyddx);
end

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

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