% Code to test the Unscented Kalman filter on a specific nonlinear time-series generated by the following model:
% x(k) = x(k-1) / 2 + 25 * x(k-1) / ( 1 + x2(k-1)) + 8 * cos (1.2 * k) + v(k-1)
% y(k) = (x(k)^2)/20 + n(k)
%
% Author: Saurabh Prasad
% Last Modified: October 31, 2005.

clear all;
N = 100;
L = 1; % Dimensionality of the signal
num_sigma = 2*L+1; %Number of Sigma Points
lambda = 3 - L;

%Defining the process and measurement noise Std. Deviations
sigma_n = 1;
sigma_v = sqrt(10);

% Defining the weights:
W_m = zeros(1,(2*L+1));
W_c = zeros(1,(2*L+1));
W_m(1)=lambda/(L+lambda);
W_c(1)=(lambda/(L+lambda));
for j = 2 : 2*L+1
    W_m(j)=1/(2*(L+lambda));
    W_c(j)=1/(2*(L+lambda));
end;

% Store the 'clean' states generated by the process model 
% x(k) = x(k-1) / 2 + 25 * x(k-1) / ( 1 + x2(k-1)) + 8 * cos (1.2 * k) + v(k-1)
state(1) = 0;
for data_index = 2 : N
    state(data_index) = (state(data_index-1)/2) + 25*state(data_index-1)/(1+(state(data_index-1))^2)+8*cos(1.2*(data_index)) + normrnd(0,sigma_v);
end;

error_avg = 0;
for loop = 1 : 100

    % Generate the noisy observations from the noisy observations
    y_corrupted = ((state(:)).^2)/20 + normrnd(0,sigma_n,N,1);

    % Augment with process and measurement noise
    X_a = [y_corrupted(1) 0 0]';

    P_0 = var(state);

    P_a = [P_0, 0, 0; 0, sigma_v, 0; 0, 0, sigma_n];
    mu_a = [0, 0, 0]';


    y_filtered(1) = y_corrupted(1); 
    state_estimate(1) = state(1);

    for data_index = 2 : length(y_corrupted)    
        % Generate the sigma points about the current estimate of the state-vector 
        [sigma] = sigmapt(X_a,mu_a',abs(P_a),L);   
        for j = 1 : num_sigma
            sigma_i = sigma(j,:);
            sigma_state = sigma_i(1);  % in a multi-dimensional case, replace (1) with length of state vector.
            state_sigma(j,:) = sigma_state + normrnd(0,sigma_v);
            clear sigma_state;
        end;

        % Compute the function Y = g(X) on the sigma points:
        % Let the sample nonlinearity be y = (x(k)^2)/20 + n(k);
        for i = 1 : L
            for j = 1 : 2*L+1
                Y(j,i) = ((sigma(j,i))^2)/20 + normrnd(0,sigma_v); % Y holds the transformed sigma_points
            end;
        end;

        [mu_x, sigma_x] = sigmawt_process(state_sigma,L);
        x_predicted = mu_x; ;
        P_predicted = sigma_x; ;

        [mu_y, sigma_y] = sigmawt(sigma,L);

        y_predicted = mu_y;

        % The following updates the innovations covariance matrix
        P_yy = zeros(L,L);
        for i = 1 : 2*L+1
            P_yy = P_yy + W_c(i)*((Y(i,:)-y_predicted)')*((Y(i,:)-y_predicted));
        end;

        % The following updates the Pxy
        P_xy = zeros(L,L);
        for i = 1 : 2*L+1
            P_xy = P_xy + W_c(i)*((state_sigma(i,:)-x_predicted)')*((Y(i,:)-y_predicted));
        end;

        % Computing the Kalman gain
        K = P_xy*inv(P_yy);

        % Estimating the state given the current innovations process
        x_estimate = x_predicted + K*(y_corrupted(data_index)-y_predicted);
        state_estimate(data_index) = x_estimate;

        P_estimate = P_predicted - K*P_yy*K';
        y_filtered(data_index) = ((x_estimate)^2)/20;

        % Updating the Augmented Covariance Marix
        P_a = [P_estimate, 0, 0; 0, sigma_v, 0; 0, 0, sigma_n]; 

    end;

    error = 0;
    for i = 1 : 100
        error(i) = ((state(i) - state_estimate(i)))^2;
    end;
    error = error./max((state(:) - state_estimate(:)).^2);
    error_avg = error+error_avg;

end;

error_avg = error_avg./100;

subplot(2,1,1); 
plot(state); hold on; plot(state_estimate);
ylabel('Amplitude'); title('Unscented Kalman Filtering');
legend('Clean series' , 'Filtered Estimate');

figure; plot(error_avg); axis([0 100 0 1]);
xlabel('Time index'); ylabel('Normalized MSE'); title('MSE');
