%% How to calculate Jacobians
%
% Suppose we are developing an extended Kalman filter (EKF) for the pendulum:
%
%
% We have the state vector \hat{\mathbf{x}}=\left(\begin{array}{c}
% \theta\\
% \dot{\theta}
% \end{array}\right) and prediction step:
% \hat{\mathbf{x}}_{k|k-1}=f\left(\hat{\mathbf{x}}_{k-1}\right)=\left(\begin{array}{c}
% \theta_{k-1}+\dot{\theta}_{k-1}\Delta t-\frac{g}{2L}\sin\left(\theta_{k-1}\right)\left(\Delta t\right)^{2}\\
% \dot{\theta}_{k-1}-\frac{g}{2L}\left(\sin\theta_{k-1}+\sin\theta_{k|k-1}\right)\Delta t
% \end{array}\right).
%
% We can use symbolic differentiation to calculate the Jacobian for this function.
%

% declare symbolic variables 
syms theta theta_dot deltaT g L m

% leapfrog integration
theta_hat = theta + theta_dot*deltaT - g/(2*L)*sin(theta)*deltaT^2;
f = [theta_hat;
     theta_dot - g/(2*L) * (sin(theta) + sin(theta_hat)) * deltaT]
% calculate the Jacobian with respect to the state vector
x = [theta; theta_dot];
F = jacobian(f, x)

%
% Hint: Once you have run the code, you can copy the Jacobian matrix as Matlab code by using the "..." button to the right of the output, and selecting "copy".
%
% Another example is the measurement function
% h\left(\hat{x}\right)=\left(\begin{array}{c}
% \frac{1}{2}m\left(L\dot{\theta}\right)^{2}\\
% mg\left(L-L\cos\theta\right)
% \end{array}\right)
% Here's how we would ask Matlab for the Jacobian of this function:

% symbolic variables were already declared above so we can just use them
h = [(1/2) * m * (L * theta_dot)^2
     m*g*(L - L*cos(theta))];

H = jacobian(h, x)

%
%% Example implementation of the EKF
%
% Specify the parameters:

g = 9.81; % m/s^2
m = 0.3; % kg 
L = 0.5; % m
theta_0 = 170 * pi / 180; % initial angle, radians
theta_dot_0 = 0; % initial angular velocity, radians/s
deltaT = 0.05; % s
t_max = 8; % s
time = 0:deltaT:t_max; % time interval for simulation

%
% Find the true solution:

[~,exact] = ode45(@(t,x)[x(2); -g/L * sin(x(1))], time, [theta_0; theta_dot_0]);
exact_theta = exact(:,1);
exact_theta_0 = exact(:,2);

figure();
ax1 = subplot(2,1,1);
plot(time, exact_theta);
ylabel('\theta (radians)');
xlabel('Time (seconds)');
title('Pendulum position');
ax2 = subplot(2,1,2);
plot(time, exact_theta_0);
ylabel('d\theta/dt (rad/s)');
xlabel('Time (seconds)');
title('Pendulum angular velocity');
linkaxes([ax1 ax2], 'x');
sgtitle('Accurate solution');

%
% Simulate noisy sensors:

% calculate the exact kinetic and potential energy
exact_energy = [(1/2)*m*(L*exact(:,2)).^2, m*g*(L-L*cos(exact(:,1)))];

% generate random numbers with given standard deviations
sigma_KE = 0.1;
sigma_PE = 0.1;
rho = 0.1;
R = [sigma_KE^2               rho*sigma_KE*sigma_PE
    rho*sigma_KE*sigma_PE     sigma_PE^2];
measurements = mvnrnd(exact_energy, R);

figure();
ax1 = subplot(2,1,1);
plot(time, exact_energy(:,1));
hold on;
plot(time, measurements(:,1), '.');
ylabel('Kinetic energy (J)');
xlabel('Time (seconds)');
title('Kinetic energy');
ax2 = subplot(2,1,2);
plot(time, exact_energy(:,2));
hold on;
plot(time, measurements(:,2), '.');

ylabel('Potential energy (J)');
xlabel('Time (seconds)');
title('Potential energy');
linkaxes([ax1 ax2], 'x');
sgtitle('Noisy sensors');

%
%
% Run the EKF

% specify the initial condition for each state variable
x = [theta_0; theta_dot_0];
EKF_output = zeros(length(time), length(x)); % preallocate memory
EKF_output(1,:) = x; % fill in the initial condition

% specify initial state covariance
P = [0    0;
     0    0];

% specify process covariance
Q = [(0.01 * deltaT)^2    0;
     0                   (0.01 * deltaT)^2];

for i = 2:length(time)
    % give friendly names to each element in the state vector
    theta = x(1);
    theta_dot = x(2);
    
    % Calculate the jacobian for the prediction step
    % This Jacobian was found using the symbolic differentiation above, and
    % reformatted into separate lines for readability.
    F = [1 - (deltaT^2*g*cos(theta))/(2*L), ...
        deltaT; ...
        
        -(deltaT*g*(cos(theta) - cos(theta + deltaT*theta_dot - (deltaT^2*g*sin(theta))/(2*L))*((deltaT^2*g*cos(theta))/(2*L) - 1)))/(2*L), ...
        1 - (deltaT^2*g*cos(theta + deltaT*theta_dot - (deltaT^2*g*sin(theta))/(2*L)))/(2*L)];

    % calculate the Jacobian for the measurement step 
    H = [0                  m*L*theta_dot
         m*g*L*sin(theta)   0];
    
    % state prediction:
    x_pred = [theta + deltaT*theta_dot - (deltaT^2*g*sin(theta))/(2*L);
              theta_dot - (deltaT*g*(sin(theta + deltaT*theta_dot - (deltaT^2*g*sin(theta))/(2*L)) + sin(theta)))/(2*L)];

    % covariance prediction:
    P = F*P*F' + Q;
  
    % measurement residual:
    % Make sure to use the predicted value of the state when calculating
    % this.
    theta_pred = x_pred(1);
    theta_dot_pred = x_pred(2);
    
    z_pred = [(1/2)*m*(L*theta_dot_pred)^2 % kinetic energy
              m*g*(L-L*cos(theta_pred))]; % potential energy
    y = measurements(i,:)' - z_pred;
    % measurement residual covariance:
    S = H*P*H' + R;
    % Kalman gain:
    K = P*H' / S;
    % updated state :
    x = x_pred + K*y;
    % updated covariance:
    P = (eye(2) - K*H) * P;
    
    % save calculated state
    EKF_output(i,:) = x;
end

% Plot the result 

figure();
ax1 = subplot(2,1,1);
plot(time, exact_theta);
hold on;
% plot(time, measurements(:,1), '.');
plot(time, EKF_output(:,1));
ylabel('\theta (radians)');
xlabel('Time (seconds)');
title('Pendulum position');
legend({'Exact', 'EKF'})
ax2 = subplot(2,1,2);
plot(time, exact_theta_0);
hold on;
% plot(time, measurements(:,2), '.');
plot(time, EKF_output(:,2));
ylabel('d\theta/dt (rad/s)');
xlabel('Time (seconds)');
title('Pendulum angular velocity');
linkaxes([ax1 ax2], 'x');
sgtitle('EKF');