%%%%%%%%% AUXILIARY FUNCTIONS %%%%%%%%%%%%%%%%%

function [xdot,u] = rhs(x,u,params)
%RHS Dynamics of an inverted nonlinear pendulum on a cart.

% Extract states
x_cart = x(1);  % cart position
x_dot_cart = x(2); % cart velocity
theta = x(3);   % pendulum angle
theta_dot = x(4); % pendulum angular velocity

%extract parameters
m_p = params.m_p;
l = params.l;
m_c = params.m_c;
g = params.g;
b_c = params.b_c;
I = params.I;

% x_ddot_cart = (m_c+m_p)\( u -b_c*x_dot_cart - m_p*l*theta_ddot*cos(theta) + m_p*l*theta_dot^2*sin(theta) );
% theta_ddot = (I + m_p*l^2)\( -m_p*g*l*sin(theta) - m_p*l*x_ddot_cart*cos(theta) );
% these equations were solved for (x_ddot_cart,theta_ddot) below

theta_ddot = -l * m_p * (cos(theta) * sin(theta) * l * m_p * theta_dot ^ 2 - cos(theta) * b_c * x_dot_cart + sin(theta) * g * m_c + sin(theta) * g * m_p + cos(theta) * u) / (-m_p ^ 2 * l ^ 2 * cos(theta) ^ 2 + l ^ 2 * m_c * m_p + l ^ 2 * m_p ^ 2 + I * m_c + I * m_p);
x_ddot_cart = 0.1e1 / (-m_p ^ 2 * l ^ 2 * cos(theta) ^ 2 + l ^ 2 * m_c * m_p + l ^ 2 * m_p ^ 2 + I * m_c + I * m_p) * (sin(theta) * l ^ 3 * m_p ^ 2 * theta_dot ^ 2 + m_p ^ 2 * l ^ 2 * g * sin(theta) * cos(theta) + I * sin(theta) * l * m_p * theta_dot ^ 2 - b_c * l ^ 2 * m_p * x_dot_cart + l ^ 2 * m_p * u - I * b_c * x_dot_cart + I * u);

%write to output
xdot = [x_dot_cart; x_ddot_cart;theta_dot; theta_ddot];

end

function [A,b,C,d] = linearized_sys_matrices_continuous(params)
% system matrices for the inverted pendulum linearized around theta = pi

%extract parameters
m_p = params.m_p;
l = params.l;
m_c = params.m_c;
g = params.g;
b_c = params.b_c;
I = params.I;

%system matries
A = [0,1,0,0;
    0,-(I+m_p*l^2)*b_c/(I*(m_c+m_p)+m_c*m_p*l^2),m_p^2*g*l^2/(I*(m_c+m_p)+m_c*m_p*l^2),0;
    0,0,0,1;
    0,-m_p*l*b_c/(I*(m_c+m_p)+m_c*m_p*l^2),m_p*g*l*(m_c+m_p)/(I*(m_c+m_p)+m_c*m_p*l^2),0];
b = [0;(I+m_p*l^2)/(I*(m_c+m_p)+m_c*m_p*l^2);0;m_p*l/(I*(m_c+m_p)+m_c*m_p*l^2)];
C = [0,0,1,0];
d = 0;

end

function [x_traj,u_traj,t_physical,x_physical] = simulate_system(u_FF_traj,x_0,t_array,params,k_T_FB,x_R)
% Simulate the inverted pendulum using given time array and input traj.
% The outer loop is running on the sample time (of our control hardware), internally, the model is
% simulated accurately.

n_steps = numel(t_array);
n_states = numel(x_0);
x_traj = zeros(n_states,n_steps);
x_traj(:,1) = x_0(:);
u_traj = zeros(1,n_steps);
t_physical = [];
x_physical = [];

for k = 1:n_steps-1
    u_traj(k) = u_FF_traj(k) + k_T_FB* (x_traj(:,k)-x_R);
    [t_interval,x_interval] = ode45(@(t,x) rhs(x,u_traj(k),params), [t_array(k),t_array(k+1)], x_traj(:,k));
    t_physical = [t_physical,t_interval'];
    x_physical = [x_physical, x_interval']; % Store the trajectory for analysis
    % Store the final state in the trajectory
    x_traj(:,k+1) = x_interval(end,:);
end

end


%%%%%%%%%%% MAIN SCRIPT %%%%%%%%%%%%%%%%%%%%%%%%%

% Parameters struct (modify as needed)
params.m_p = 0.2;      % pendulum mass
params.l = 0.3;        % pendulum length (distance to center of mass)
params.m_c = 0.5;      % cart mass
params.g = 9.81;       % gravity
params.b_c = 0.4;      % cart viscous damping
params.I = params.m_p * params.l^2;  % moment of inertia of the (mathematical) pendulum

% Initial conditions for the simulation
x_0 = [0; -1; pi*0.99; 0]; % [x_cart, x_dot_cart, theta, theta_dot]
% Steady state position (for FB controller)
x_R = [0; 0; pi; 0];

% sampling time of digital control loop
T_a = 0.3;
% Time array for the control loop's simulation (physical model is simulated more finely)
t_array = (0:T_a:10);

% Define control input (for example, a step input)
u_FF_traj = 0*sin(t_array); % constant force applied to the cart

% design Ackermann controller in continuous time for desired poles
[A,b,C,d] = linearized_sys_matrices_continuous(params);
sys_continuous = ss(A,b,C,d);
A_eigenvalues = eig(A);
poles_des_continuous = [A_eigenvalues(A_eigenvalues<0);-10;-20];
k_T_continuous = -place(A,b,poles_des_continuous);

% discretize system and design Ackermann controller in discrete time
sys_discrete = c2d(sys_continuous,T_a);
Phi = sys_discrete.A;
Phi_eigenvalues = eig(Phi);
Gamma = sys_discrete.B;
poles_des_discrete = [Phi_eigenvalues(abs(Phi_eigenvalues)<1);0.5;0.3];
k_T_discrete = -place(Phi,Gamma,poles_des_discrete);

% solve system forward
[x,u,t_physical,x_physical] = simulate_system(u_FF_traj,x_0,t_array,params,k_T_discrete,x_R);

% plot the simulation results
figure()
plot(t_physical,x_physical(1,:))
title("x postion of the cart")
xlabel('t')
ylabel('x_{cart}')
figure()
plot(t_physical,x_physical(3,:))
title("Angle of the pendulum")
xlabel('t')
ylabel('\theta')
figure()
stairs(t_array,u, 'LineWidth', 1.2)
xlabel('t')
ylabel('u')
title('Input force on cart (piecewise-constant)')




