%% Dynamic Model

clc
clear

% 0: cart; 1-3: joints
syms q0 q1 q2 q3 real
syms q_dot_0 q_dot_1 q_dot_2 q_dot_3 real
syms q_ddot_0 q_ddot_1 q_ddot_2 q_ddot_3  real

syms t real positive
syms f real % control input

q = [q0; q1; q2; q3]
dq = [q_dot_0; q_dot_1; q_dot_2; q_dot_3]
ddq = [q_ddot_0; q_ddot_1; q_ddot_2; q_ddot_3]

% kinematic parameters
syms l1 l2 l3 a1 a2 a3 real positive
% dynamic parameters
syms J1 J2 J3 m0 m1 m2 m3 d0 d1 d2 d3 real positive
syms g real positive

function [M,C,G,D,S] = model_full(q,dq,parameters)
q0 = q(1);
q1 = q(2);
q2 = q(3);
q3 = q(4);
q_dot_0 = dq(1);
q_dot_1 = dq(2);
q_dot_2 = dq(3);
q_dot_3 = dq(4);

m0 = parameters.m0;
m1 = parameters.m1;
m2 = parameters.m2;
m3 = parameters.m3;
d0 = parameters.d0;
d1 = parameters.d1;
d2 = parameters.d2;
d3 = parameters.d3;
J1 = parameters.J1;
J2 = parameters.J2;
J3 = parameters.J3;
l1 = parameters.l1;
l2 = parameters.l2;
l3 = parameters.l3;
a0 = parameters.a0;
a1 = parameters.a1;
a2 = parameters.a2;
a3 = parameters.a3;
g = parameters.g;

M = [
    m0+m1+m2+m3, -a1*m1*cos(q1)-l1*m2*cos(q1)-l1*m3*cos(q1), -a2*m2*cos(q2)-l2*m3*cos(q2), -a3*m3*cos(q3);
    -a1*m1*cos(q1)-l1*m2*cos(q1)-l1*m3*cos(q1), J1+a1^2*m1+l1^2*m2+l1^2*m3, a2*l1*m2*cos(q1-q2)+l1*l2*m3*cos(q1-q2), a3*l1*m3*cos(q1-q3);
    -a2*m2*cos(q2)-l2*m3*cos(q2), a2*l1*m2*cos(q1-q2)+l1*l2*m3*cos(q1-q2), J2+a2^2*m2+l2^2*m3, a3*l2*m3*cos(q2-q3);
    -a3*m3*cos(q3), a3*l1*m3*cos(q1-q3), a3*l2*m3*cos(q2-q3), J3+a3^2*m3
    ];

C = [
    0, q_dot_1*a1*m1*sin(q1)+q_dot_1*l1*m2*sin(q1)+q_dot_1*l1*m3*sin(q1), q_dot_2*a2*m2*sin(q2)+q_dot_2*l2*m3*sin(q2), q_dot_3*a3*m3*sin(q3);
    0,0,q_dot_2*a2*l1*m2*sin(q1-q2)+q_dot_2*l1*l2*m3*sin(q1-q2), q_dot_3*a3*l1*m3*sin(q1-q3);
    0,-q_dot_1*a2*l1*m2*sin(q1-q2)-q_dot_1*l1*l2*m3*sin(q1-q2), 0, q_dot_3*a3*l2*m3*sin(q2-q3);
    0,-q_dot_1*a3*l1*m3*sin(q1-q3),-q_dot_2*a3*l2*m3*sin(q2-q3),0
    ];

G = -g*[0;
    sin(q1)*(a1*m1+l1*m2+l1*m3);
    sin(q2)*(a2*m2 + l2*m3);
    sin(q3)*a3*m3
    ];

% damping terms
D = [
    d0,0,0,0;
    0,d1+d2,-d2,0;
    0,-d2,d2+d3,-d3;
    0,0,-d3,d3
    ];

S = [1;0;0;0]; % selection matrix

end

%
%% Simplifications
% The authors assume that the cart behaves like a double integrator. Any intertial coupling and effects deriving from the poles acting on the cart are ignored, and a lower-level feedback loop (e.g. speed controller) will take care of that.
% The authors use a speed control loop for the motor, but the control input for the triple pendulum on cart is cart acceleration, so they dynamically extend the output of the controller.

% this control law makes the cart dynamics a double integrator, with input
% v = ddq0
syms v real
syms phi_1 phi_2 phi_3 real
syms phi_dot_1 phi_dot_2 phi_dot_3 real

phi = [phi_1; phi_2; phi_3]
dphi = [phi_dot_1; phi_dot_2; phi_dot_3]

function [M,C,G,D,bs]=model_reduced(q,dq,parameters)
    m0 = parameters.m0;
    m1 = parameters.m1;
    m2 = parameters.m2;
    m3 = parameters.m3;
    a1 = parameters.a1;
    a2 = parameters.a2;
    a3 = parameters.a3;
    l1 = parameters.l1;
    l2 = parameters.l2;

    q1 = q(1);
    q2 = q(2);
    q3 = q(3);

    % expand q,dq to make sense for the full model
    q_ = [0; q];
    dq_ = [0; dq];

    [M,C,G,D,~]=model_full(q_,dq_,parameters);

    %f = D(1,:)*dq + G(1,:) + c(1) + M(1,:)*ddq + M(1,1)*v;

    M = M(2:4, 2:4);
    C = C(2:4, 2:4);
    G = G(2:4);
    D = D(2:4, 2:4);

    bs = [
        a1*m1*cos(q1) + l1*m2*cos(q1) + l1*m3*cos(q1);
        a2*m2*cos(q2) + l2*m3*cos(q2);
        a3*m3*cos(q3)
        ];
end

%
%% Falling pendulum simulation

% parameters from author's indentification process
parameters_num.l1 = 0.319; parameters_num.l2 = 0.419; parameters_num.l3 = 0.480; % link length [m]
parameters_num.a1 = 0.221; parameters_num.a2 = 0.225; parameters_num.a3 = 0.185; % distnace to center of mass [m]
parameters_num.m1 = 0.653; parameters_num.m2 = 0.852; parameters_num.m3 = 0.615; % mass [kg]
parameters_num.J1 = 0.009; parameters_num.J2 = 0.024; parameters_num.J3 = 0.021; % moment of inertia [Nms^2]
parameters_num.d1 = 0.016; parameters_num.d2 = 0.005; parameters_num.d3 = 0.001; % viscous friction coefficient [Nm (rad/s)^-1]
parameters_num.m0=0;parameters_num.l0=0;parameters_num.J0=0;parameters_num.a0=0;parameters_num.d0=0;
parameters_num.g = 9.81; % gravitational constant [m^3 (kg s^2)^-1]

% only the pendulum
function dy=sim_reduced(t,y,parameters,ffw)
    p = y(1:3);
    dp = y(4:end);
    
    [M,C,G,D,bs]=model_reduced(p,dp,parameters);

    ddp = M\(-C*dp -G -D*dp + bs*ffw);
    dy = [dp; ddp];
end

%


dt = 1e-3
tend = 8
tspan = 0:dt:tend
y0 = [0;0;0;0;0;0.4];
options = odeset('RelTol', 1e-6, 'AbsTol', 1e-6);
[t, y] = ode45(@(t,y) sim_reduced(t,y,parameters_num,0), tspan, y0, options); % I think this fails because the pendulum falling is too chaotic

figure(1)
clf;
hold on
plot(t, y(:, 1))
plot(t, y(:, 2))
plot(t, y(:, 3))
grid minor;

%
%% Feedforward

% trajectory that makes the third link rotate
function [bvp_sol, cart_fun] =BVP(x0, xT, T_real, pguess, parameters_num, N)

    t = sym("t", [1, 1], ["real", "positive"]);
    T = sym("T", [1, 1], ["real", "positive"]);
    p = sym("p", [6, 1], 'real');
    %syms p0 p1 p2 p3 p4 p5 real
    %p = [p0;p1;p2;p3;p4;p5];
    
    % TODO: use this. bvp fails with Error: Difference order N must be a positive integer scalar
    %arg = pi*mod(t,T)/T;
    
    % use sine series as initial guess for bvp5c
    xc_t = 0;
    for i=1:6
        xc_t = xc_t + p(i)*sin((2+i)*pi*t / T);
    end
    % authors add these two terms to the initial guess
    xc_t = xc_t + (-3*p(1)-5*p(3)-7*p(5))*sin(pi*t/T) + (-2*p(2)-3*p(4)-4*p(6))*sin(2*pi*t/T);
    
    dxc_t = diff(xc_t, t);
    ddxc_t = diff(dxc_t, t);
    
    % boundary conditions. ya and yb should be initial and final values, resp.
    % Must be in g(y(a), y(b))=0 form
    function res=bcfcn(ya,yb,x0,xT)
        res = [
            ya-x0;
            yb-xT
            ];
    end
    
    tmesh = linspace(0,T_real,N);
    % 3. Construct initial solution as linear interpolation of the BCs on the
    % grid
    xguess = @(t) x0 + (t/T_real)*(xT-x0);    
    
    xc_handler = matlabFunction(xc_t, 'Vars', {t, T, p});
    dxc_handler = matlabFunction(dxc_t, 'Vars', {t, T, p});
    ddxc_handler = matlabFunction(ddxc_t, 'Vars', {t, T, p});
    
    ffw_handler = ddxc_handler;
    ode_handler = @(t,y,p) sim_reduced(t, y, parameters_num, ffw_handler(t,T_real,p));
    bcfcn_handler = @(ya, yb, p) bcfcn(ya, yb, x0, xT);
    
    solinit = bvpinit(tmesh,xguess,pguess);
    options = bvpset('NMax', 20000, 'RelTol', 1e-6, 'AbsTol', 1e-6, 'Stats', 'on');
    bvp_sol = bvp5c(ode_handler, bcfcn_handler, solinit, options);

    cart_fun.xc_handler = xc_handler;
    cart_fun.dxc_handler = dxc_handler;
    cart_fun.ddxc_handler = ddxc_handler;
    cart_fun.ffw_handler = ffw_handler;
    
    p_sol = bvp_sol.parameters;
    cart_fun.xc = @(t) cart_fun.xc_handler(mod(t, T_real), T_real, p_sol);
    cart_fun.dxc = @(t) cart_fun.dxc_handler(mod(t, T_real), T_real, p_sol);
    cart_fun.ddxc = @(t) cart_fun.ddxc_handler(mod(t, T_real), T_real, p_sol);
    cart_fun.ffw = @(t) cart_fun.ffw_handler(mod(t, T_real), T_real, p_sol);
end

function plot_bvp(bvp_sol, cart_fun)
    t = bvp_sol.x;
    phi = bvp_sol.y(1:3, :);
    x = cart_fun.xc(t);
    
    figure()
    subplot(221)
    title("BVP solution")
    plot(t, x)
    subplot(222)
    title("BVP solution")
    plot(t, cart_fun.dxc(t))
    subplot(223)
    title("BVP solution")
    plot(t, cart_fun.ddxc(t))
    subplot(224)
    hold on
    for i=1:3
        plot(t, phi(i, :))
    end
end

%
%% Feedforward simulation

% full simulations

% simulation of full model, including cart as double integrator (with new
% input v)
function dy=sim_full_simplified(t,y,parameters,v)
    x = y(1);
    dx = y(2);
    p = y(3:5);
    dp = y(6:8);

    [M,C,G,D,bs]=model_reduced(p,dp,parameters);

    ddp = M\(-C*dp -G -D*dp + bs*v);
    dy = [dx; v; dp; ddp];
end

%
%% Feedback LQR

% compute linear approximation around final point of the trajectory
x = [q0; q_dot_0; phi; dphi];
f_full = sim_full_simplified(t, x, parameters_num, v);
% Linearized dynamics
A = jacobian(f_full, x);
A_fun = matlabFunction(A, 'vars', {q0, q_dot_0, phi, dphi, v})
b = jacobian(f_full, v);
b_fun = matlabFunction(b, 'vars', {q0, q_dot_0, phi, dphi, v})

%

% returns an handler to compute LQR gains at any point in time
function K_t=LQR(A_fun, b_fun, q_fun, ffw, Q, r, T_real, periods)    
    % Differential Riccati Equation
    % use ode45 to solve DRE, so define the function K_dot
    function dP=dP(t,P,Q,r,A_fun, b_fun, x, v)
        % compute linearization around current point in trajectory
        q0_ = x(1);
        dq0_ = x(2);
        phi_ = x(3:5);
        dphi_ = x(6:8);
        A = A_fun(q0_, dq0_, phi_, dphi_, v);
        b = b_fun(q0_, dq0_, phi_, dphi_, v);
    
        P = reshape(P, [8,8]); %matlab reshapes the matrix as a vector
        dP = -A'*P - P*A + P*(b*b')*P/r - Q;
        dP  = reshape(dP, [], 1); % make derivative back into a vector
    end
        
    % Solve DRE backwards in time
    dt = 1e-3;
    tend = T_real*periods;
    tspan = tend:-dt:0;
    
    options = odeset('RelTol', 1e-6, 'AbsTol', 1e-16);
    [t, P] = ode45(@(t,P) dP(t,P,Q,r,A_fun, b_fun, q_fun(t), ffw(t)), tspan, zeros(8));%, options);    
    
    %if periods > 1
        % Solve DRE backwards in time again, now using the knowledge of how many
        % cycles it needs to converge
        tspan=T_real*2:-dt:0;
        [t, P] = ode45(@(t,P) dP(t,P,Q,r,A_fun, b_fun, q_fun(t), ffw(t)), tspan, reshape(P(end, :),8,8), options);
    %end

    K = zeros(size(P, 1), 8);
    K_spline = cell(8, 1);
    for i=1:size(t,1)
        p = reshape(P(i, :), [8,8]);
        t_ = t(i);
        x_ = q_fun(t_);
        q0_ = x_(1);
        dq0_ = x_(2);
        phi_ = x_(3:5);
        dphi_ = x_(6:8);
        b_ = b_fun(q0_, dq0_, phi_, dphi_, ffw(t_));
        k = -r\(b_'*p);
        K(i, :) = k;
    end
    for i=1:8
        K_spline{i} = spline(t, K(:, i));
    end
    K_t = @(t) [ppval(K_spline{1}, t) ppval(K_spline{2}, t) ppval(K_spline{3}, t) ppval(K_spline{4}, t) ppval(K_spline{5}, t) ppval(K_spline{6}, t) ppval(K_spline{7}, t) ppval(K_spline{8}, t)];
end

%
%% Feedforward + feedback simulation functions

% Feedback + Feedforward simulation

% LQR feedback is local to deviation from reference trajectory
function diff=circular_error(x,y)
    diff = mod(x-y, 2*pi);
    if diff > pi
        diff = diff - 2*pi;
    end
end

function n=process_noise(sigma)
    n = [0; randn(1)*sigma; zeros(3,1); randn(3,1)*sigma];
end

function u=feedback(t, y, K, qref)
    x = y(1);
    dx = y(2);
    p = y(3:5);
    dp = y(6:8);
    xc = qref(1);
    dxc = qref(2);
    phi_ref = qref(3:end);

    err_x = x - xc;
    err_dx = dx - dxc;
    err_phi1 = circular_error(p(1), phi_ref(1));
    err_phi2 = circular_error(p(2), phi_ref(2));
    err_phi3 = circular_error(p(3), phi_ref(3));
    err_phi = [err_phi1; err_phi2; err_phi3];
    err_dphi = dp - phi_ref(4:6);

    err = [err_x; err_dx; err_phi; err_dphi];
    u = K*err;
end

function dy=sim_full_simplified_noisy(t,y,parameters,v)
    dy = sim_full_simplified(t,y,parameters,v) + process_noise(0.08);
end

%
%% Periodic trajectory simulation

% Setup
traj_periodic.pguess = ones(6,1)*0.01;
traj_periodic.T_real = 1.096;
traj_periodic.deta3 = 1.89;
traj_periodic.x0 = [0;0;0;0;0;traj_periodic.deta3];
traj_periodic.xT = [0;0;2*pi;0;0;traj_periodic.deta3];
traj_periodic.N = 500;

[bvp_sol, cart_fun] = BVP(traj_periodic.x0, traj_periodic.xT, traj_periodic.T_real, traj_periodic.pguess, parameters_num, traj_periodic.N);
traj_periodic.bvp_sol = bvp_sol;
traj_periodic.cart_fun = cart_fun;
plot_bvp(bvp_sol, cart_fun)

phi_interp = spline(traj_periodic.bvp_sol.x, traj_periodic.bvp_sol.y);
traj_periodic.phi_fun = @(t) ppval(phi_interp, mod(t, traj_periodic.T_real));
traj_periodic.q_fun = @(t) [traj_periodic.cart_fun.xc(t); traj_periodic.cart_fun.dxc(t); traj_periodic.phi_fun(t)];

%


traj_periodic.periods = 10;
traj_periodic.Q = diag([10, 1, 1, 1, 10, 0, 0, 0])
traj_periodic.r = 1;
traj_periodic.K_t = LQR(A_fun, b_fun, traj_periodic.q_fun, traj_periodic.cart_fun.ffw, traj_periodic.Q, traj_periodic.r, traj_periodic.T_real, traj_periodic.periods);

% Simulation
dt = 1e-3
tend = traj_periodic.T_real*2
tspan = 0:dt:tend
traj_periodic.y0 = [0;0;0;0;0;0;0;traj_periodic.deta3];
options = odeset('RelTol', 1e-12, 'AbsTol', 1e-16);
ode_handler = @(t,y,p) sim_full_simplified_noisy(t, y, parameters_num, traj_periodic.cart_fun.ffw(t)+feedback(t, y, traj_periodic.K_t(t), traj_periodic.q_fun(t)));
[t, y] = ode45(ode_handler, tspan, traj_periodic.y0);
y_period = y

figure()
clf;
subplot(211)
plot(t, y(:, 1))
subplot(212)
hold on
plot(t, y(:, 3))
plot(t, y(:, 4))
plot(t, y(:, 5))
grid minor;

%
%% Entry maneuver simulation

% Simulate the entry maneuver first, then simulate periodic trajectory
% using as initial condition the final state of the entry maneuver

% Setup
traj_entry.pguess = zeros(6,1);
traj_entry.T_real = 2.5025;%2.65;
traj_entry.x0 = [0;0;0;0;0;0];
traj_entry.xT = [0;0;2*pi;0;0;traj_periodic.deta3];
traj_entry.N = 500;

[bvp_sol, cart_fun] = BVP(traj_entry.x0, traj_entry.xT, traj_entry.T_real, traj_entry.pguess, parameters_num, traj_entry.N);
traj_entry.bvp_sol = bvp_sol;
traj_entry.cart_fun = cart_fun;

phi_interp = spline(traj_entry.bvp_sol.x, traj_entry.bvp_sol.y);
traj_entry.phi_fun = @(t) ppval(phi_interp, t);
traj_entry.q_fun = @(t) [traj_entry.cart_fun.xc(t); traj_entry.cart_fun.dxc(t); traj_entry.phi_fun(t)];

plot_bvp(traj_entry.bvp_sol, traj_entry.cart_fun)

%

traj_entry.Q = diag([1, 0.01, 0.01, 0.1, 1, 0, 0, 0])*0.1;
traj_entry.r = 1;
traj_entry.K_t = LQR(A_fun, b_fun, traj_entry.q_fun, traj_entry.cart_fun.ffw, traj_entry.Q, traj_entry.r, traj_entry.T_real, 1);

figure()
hold on
KK = traj_entry.K_t(bvp_sol.x')
for i = 1:8
    plot(bvp_sol.x, KK(:, i), 'DisplayName', ['k_', num2str(i)])
end
legend('Interpreter', 'latex')

%

% Simulation of entry maneuver
dt = 1e-3
tend = traj_entry.T_real
tspan = 0:dt:tend
traj_entry.y0 = [0;0;0;0;0;0;0;0];
ode_handler = @(t,y,p) sim_full_simplified_noisy(t, y, parameters_num, cart_fun.ffw(t)+feedback(t, y, traj_entry.K_t(t), traj_entry.q_fun(t)));
[t1, y1] = ode45(ode_handler, tspan, traj_entry.y0);

figure()
clf;
subplot(211)
plot(t, y(:, 1))
xlabel("t")
ylabel("x(t)")
grid minor;
subplot(212)
hold on
plot(t1, y1(:, 3),"DisplayName","\phi_1")
plot(t1, y1(:, 4),"DisplayName","\phi_2")
plot(t1, y1(:, 5),"DisplayName","\phi_3")
xlabel("t")
legend("Interpreter","latex","Location","northwest")
grid minor;

%disp('Final configuration', y(end, :))

%
%% Entry maneuver + periodic trajectory simulation

% Simulation of periodic trajectory starting from end of entry maneuver
dt = 1e-3
tend = traj_periodic.T_real*2
tspan = 0:dt:tend
ode_handler = @(t,y,p) sim_full_simplified_noisy(t, y, parameters_num, traj_periodic.cart_fun.ffw(t)+feedback(t, y, traj_periodic.K_t(t), traj_periodic.q_fun(t)));
[t2, y2] = ode45(ode_handler, tspan, y1(end, :));

t2 = t2+t1(end)

T = [t1;t2];
Y = [y1;y2];

figure()
clf;
subplot(211)
plot(T, Y(:, 1))
subplot(212)
hold on
plot(T, Y(:, 3))
plot(T, Y(:, 4))
plot(T, Y(:, 5))
grid minor;

%
%% Swing-up trajectory

% Simulate the swingup manevur first, then entry maneuver, then simulate periodic trajectory
% using as initial condition the final state of the entry maneuver

% Setup
traj_swingup.pguess = zeros(6,1);
traj_swingup.T_real = 2.4;
traj_swingup.x0 = [-pi;-pi;-pi;0;0;0];
traj_swingup.xT = zeros(6,1);
traj_swingup.N = 500;

[bvp_sol, cart_fun] = BVP(traj_swingup.x0, traj_swingup.xT, traj_swingup.T_real, traj_swingup.pguess, parameters_num, traj_swingup.N);
traj_swingup.bvp_sol = bvp_sol;
traj_swingup.cart_fun = cart_fun;

phi_interp = spline(traj_swingup.bvp_sol.x, traj_swingup.bvp_sol.y);
traj_swingup.phi_fun = @(t) ppval(phi_interp, t);
traj_swingup.q_fun = @(t) [traj_swingup.cart_fun.xc(t); traj_swingup.cart_fun.dxc(t); traj_swingup.phi_fun(t)];

plot_bvp(traj_swingup.bvp_sol, traj_swingup.cart_fun)

%

traj_swingup.Q = diag([10, 1, 0.1, 0.1, 1, 0, 0, 0])*0.1; % Check these because they don't really work
traj_swingup.r = 1;
traj_swingup.K_t = LQR(A_fun, b_fun, traj_swingup.q_fun, traj_swingup.cart_fun.ffw, traj_swingup.Q, traj_swingup.r, traj_swingup.T_real, 1);

figure()
hold on
KK = traj_swingup.K_t(bvp_sol.x')
for i = 1:8
    plot(bvp_sol.x, KK(:, i), 'DisplayName', ['k_', num2str(i)])
end
legend('Interpreter', 'latex', 'Location', 'southwest')

%

% Simulation of swingup
dt = 1e-3
tend = traj_swingup.T_real
tspan = 0:dt:tend
options = odeset('RelTol', 1e-6, 'AbsTol', 1e-6);
% TODO: check LQR gains because noisy simulation diverges
ode_handler = @(t,y,p) sim_full_simplified_noisy(t, y, parameters_num, traj_swingup.cart_fun.ffw(t)+feedback(t, y, traj_swingup.K_t(t), traj_swingup.q_fun(t)));
[t1, y1] = ode45(ode_handler, tspan, [0;0;traj_swingup.x0], options);

T = t1;

Y = y1;
% Simulation of entry maneuver after swingup
dt = 1e-3
tend = traj_entry.T_real
tspan = 0:dt:tend
ode_handler = @(t,y,p) sim_full_simplified_noisy(t, y, parameters_num, traj_entry.cart_fun.ffw(t)+feedback(t, y, traj_entry.K_t(t), traj_entry.q_fun(t)));
[t2, y2] = ode45(ode_handler, tspan, y1(end,:));

T = [T; t2+T(end)]
Y = [Y; y2];

% Simulation of periodic trajectory after entry maneuver after swingup
dt = 1e-3
tend = traj_periodic.T_real*2
tspan = 0:dt:tend
ode_handler = @(t,y,p) sim_full_simplified_noisy(t, y, parameters_num, traj_periodic.cart_fun.ffw(t)+feedback(t, y, traj_periodic.K_t(t), traj_periodic.q_fun(t)));
[t3, y3] = ode45(ode_handler, tspan, y2(end,:));

T = [T; t3+T(end)]
Y = [Y;y3];

figure()
clf;
subplot(211)
plot(T, Y(:, 1))
subplot(212)
hold on
plot(T, Y(:, 3))
plot(T, Y(:, 4))
plot(T, Y(:, 5))
grid minor;

%
%% Defuse maneuver

% Simulate the defuse manevur first, then entry maneuver, then simulate periodic trajectory
% using as initial condition the final state of the entry maneuver

% Setup
%traj_defuse.pguess = zeros(6,1);
%traj_defuse.T_real = 2.526;%2.519;
%traj_defuse.x0 = [0;0;0;0;0;traj_periodic.deta3];
%traj_defuse.xT = zeros(6,1);
%traj_defuse.xT = [0;0;2*pi;0;0;0];
%traj_defuse.N = 500;

% traj_defuse.pguess = zeros(6,1);
% traj_defuse.T_real = 2.5025;
% traj_defuse.x0 = [0;0;0;0;0;traj_periodic.deta3];
% traj_defuse.xT = [0;0;2*pi;0;0;0];
% traj_defuse.N = 500;

traj_defuse.pguess = zeros(6,1)*0.01;
traj_defuse.T_real = 2.958; %2.95722 %2.967495 %1.12*2.65;
%traj_defuse.T_real = 1.25; 1.87513; %1.180211 %1.2118; %2.954;
traj_defuse.x0 = [0;0;0;0;0;traj_periodic.deta3];
traj_defuse.xT = [0;0;2*pi;0;0;0];
traj_defuse.N = 500;

[bvp_sol, cart_fun] = BVP(traj_defuse.x0, traj_defuse.xT, traj_defuse.T_real, traj_defuse.pguess, parameters_num, traj_defuse.N);
traj_defuse.bvp_sol = bvp_sol;
traj_defuse.cart_fun = cart_fun;

phi_interp = spline(traj_defuse.bvp_sol.x, traj_defuse.bvp_sol.y);
traj_defuse.phi_fun = @(t) ppval(phi_interp, t);
traj_defuse.q_fun = @(t) [traj_defuse.cart_fun.xc(t); traj_defuse.cart_fun.dxc(t); traj_defuse.phi_fun(t)];

plot_bvp(traj_defuse.bvp_sol, traj_defuse.cart_fun)

%

traj_defuse.Q = diag([50, 1, 5, 5, 5, 0, 0, 0]);
traj_defuse.r = 1;
traj_defuse.K_t = LQR(A_fun, b_fun, traj_defuse.q_fun, traj_defuse.cart_fun.ffw, traj_defuse.Q, traj_defuse.r, traj_defuse.T_real, 1);

figure()
hold on
KK = traj_defuse.K_t(bvp_sol.x')
for i = 1:8

    plot(bvp_sol.x, KK(:, i), 'DisplayName', ['k_', num2str(i)])
end
legend('Interpreter', 'latex', 'Location', 'southwest')

dt = 1e-3;
tend = traj_defuse.T_real;
tspan = 0:dt:tend;

% TODO: make this noisy
ode_handler = @(t,y,p) sim_full_simplified(t, y, parameters_num, traj_defuse.cart_fun.ffw(t)+feedback(t, y, traj_defuse.K_t(t), traj_defuse.q_fun(t)));
options = odeset('RelTol', 1e-10, 'AbsTol', 1e-10);
[t4, y4] = ode45(ode_handler, tspan, y3(end, :), options);


T = [T; t4+T(end)];
Y = [Y;y4];

figure()
clf;
subplot(311)
plot(t4, y4(:, 1))
subplot(312)
hold on
plot(t4, y4(:, 3))
plot(t4, y4(:, 4))
plot(t4, y4(:, 5))
grid minor;

subplot(313)
hold on
plot(t4, y4(:, 6))
plot(t4, y4(:, 7))
plot(t4, y4(:, 8))
grid minor;

%

% keep upright
Q = diag([ 12, 0.1, 10, 10, 10, 1, 1, 1])
r = 1;
% solve ARE
A_upright = A_fun(0,0,zeros(3,1),zeros(3,1), 0)
b_upright = b_fun(0,0,zeros(3,1),zeros(3,1), 0)
K_are = lqr(A_upright, b_upright, Q, traj_periodic.r)
ode_handler = @(t,y,p) sim_full_simplified_noisy(t, y, parameters_num, feedback(t, y, -K_are, zeros(8,1)));
[t5, y5] = ode45(ode_handler, tspan, rand(8,1)*0.01);

T = [T; t5+T(end)]
Y = [Y;y5];


figure()
clf;
subplot(411)
plot(t5, y5(:, 1))
subplot(412)
plot(t5, y5(:, 2))
subplot(413)
hold on
plot(t5, y5(:, 3))
plot(t5, y5(:, 4))
plot(t5, y5(:, 5))
grid minor;
subplot(414)
hold on
plot(t5, y5(:, 6))
plot(t5, y5(:, 7))
plot(t5, y5(:, 8))
grid minor;

%

% Swing down

traj_swingdown.pguess = -ones(6,1)*0.01;
traj_swingdown.T_real = 2.802;
traj_swingdown.x0 = [0;0;0;0;0;0];
traj_swingdown.xT = [pi;pi;pi;0;0;0];
traj_swingdown.N = 500;


[bvp_sol, cart_fun] = BVP(traj_swingdown.x0, traj_swingdown.xT, traj_swingdown.T_real, traj_swingdown.pguess, parameters_num, traj_swingdown.N);
traj_swingdown.bvp_sol = bvp_sol;
traj_swingdown.cart_fun = cart_fun;

phi_interp = spline(traj_swingdown.bvp_sol.x, traj_swingdown.bvp_sol.y);
traj_swingdown.phi_fun = @(t) ppval(phi_interp, t);
traj_swingdown.q_fun = @(t) [traj_swingdown.cart_fun.xc(t); traj_swingdown.cart_fun.dxc(t); traj_swingdown.phi_fun(t)];

plot_bvp(traj_swingdown.bvp_sol, traj_swingdown.cart_fun)

%

traj_swingdown.Q = diag([1, 1, 10, 1, 1, 1, 5, 1])*1; % weighting of states errors
traj_swingdown.r = 1; % weighting of control output
traj_swingdown.K_t = LQR(A_fun, b_fun, traj_swingdown.q_fun, traj_swingdown.cart_fun.ffw, traj_swingdown.Q, traj_swingdown.r, traj_swingdown.T_real, 1);

%

%TODO: make this noisy
 
ode_handler = @(t,y,p) sim_full_simplified(t, y, parameters_num, traj_swingdown.cart_fun.ffw(t)+feedback(t, y, traj_swingdown.K_t(t), traj_swingdown.q_fun(t)));
options = odeset('RelTol', 1e-10, 'AbsTol', 1e-10);
[t6, y6] = ode45(ode_handler, tspan, y5(end, :), options);


T = [T; t6+T(end)]
Y = [Y;y6];

figure()
clf;
subplot(411)
plot(t6, y6(:, 1))
subplot(412)
plot(t6, y6(:, 2))
subplot(413)
hold on
plot(t6, y6(:, 3))
plot(t6, y6(:, 4))
plot(t6, y6(:, 5))
grid minor;
subplot(414)
hold on
plot(t6, y6(:, 6))
plot(t6, y6(:, 7))
plot(t6, y6(:, 8))
grid minor;

%