%% Unscented Kalman filter (UKF)
%
% We are developing an UKF 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).
%
% The measurement model is:
% 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)
%
%
% 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.01; % 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 UKF:

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

% Specify initial state covariance.
% This cannot be 0 for the UKF.
P = [0.1^2   0;
    0        0.2^2];

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

% Unscented transform parameters
alpha = 0.001;
beta = 2;
n = length(x);
lambda = alpha^2 * n - n;

% Unscented transform weights
w_m = [lambda/(n+lambda) ones(1,2*n)*(1/(2*(n+lambda)))];
w_c = w_m;
w_c(1) = lambda/(n+lambda)+1-alpha^2+beta;

for k = 2:length(time)
    % Prediction step
    
    % Calculate sigma points
    A = sqrt(n+lambda)*chol(P);
    sigma_points = [x (x+A') (x-A')];
    % each column is one of the sigma points
    
    % propagate sigma points through the prediction function
    F = zeros(2, 2*n+1);
    for i = 1:(2*n+1)
        % process each column by pulling out the state variables
        theta = sigma_points(1,i);
        theta_dot = sigma_points(2,i);

        F(1,i) = theta + theta_dot*deltaT - g/(2*L)*sin(theta)*deltaT^2; % prediction for theta
        F(2,i) = theta_dot - g/(2*L) * (sin(theta) + sin(F(1,i))) * deltaT; % prediction for theta_dot
        % ^^ this could in fact be any arbitrary computer code
    end
    
    % compute output statistics
    x_pred = F*w_m';
    P_pred = w_c .* (F - x_pred) * (F - x_pred)' + Q; % add process covariance
    
    % Propagate the predction through the measurement function
    
    % Calculate new sigma points 
    A = sqrt(n+lambda)*chol(P_pred);
    sigma_points = [x_pred (x_pred+A') (x_pred-A')];
    % Pass each column through the measurement function
    H = zeros(2, 2*n+1);
    for i = 1:(2*n+1)
        % process each column by pulling out the state variables
        theta = sigma_points(1,i);
        theta_dot = sigma_points(2,i);
        
        H(1,i) = (1/2)*m*(L*theta_dot).^2; % kinetic energy
        H(2,i) = m*g*(L-L*cos(theta)); % potential energy
    end
    % Predicted measurements
    z_pred = sum(w_m .* H, 2);
    % Predicted measurement covariance:
    S = w_c .* (H - z_pred) * (H - z_pred)' + R; % add measurement covariance
    % Measurement residual:
    y = measurements(k,:)' - z_pred;
    
    % Kalman gain:
    T = w_c .* (F - x_pred) * (H - z_pred)';
    K = T / S;
    % updated state :
    x = x_pred + K*y;
    % updated covariance:
    P = P_pred - K*S*K';
   
    % save calculated state
    UKF_output(k,:) = x;
end

% Plot the result

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