%Script for calculating 3D positions from measured azimuth and elevation
%angles (AoA method) using U-blox XPLR-AOA-2 kit

%Using Newton-Raphson method

%Author: Petr Šimek
%This script is part of a diploma thesis "3D lokalizace BLE modulů úhloměrnou
%metodou využívající Kálmánův filtr"

%Inputs:
%true_pos = True position of tag: [x,y,z] (input directly in script)
%Anchs = Anchor positions and rotations in following format:
    %[x1,y1,z1,rot_horizontal1,rot_vertical1] rotation input in degrees
    %[x2,y2,z2,rot_horizontal2,rot_vertical2]
    %...
    %input as txt file
%z_matrix = Measured azimuth and elevation values in following format: (in degrees)
    %[az1,az2,az3,az4,el1,el2,el3,el4]
    %each row represents one measurement from all anchors
    %index represents anchor number
    %if value is not measured, replace it with NaN
    %input as txt file

%Outputs
%Graph with calculated positions
%x_all matrix with all calculated positions
%TOTAL_RMSE = calculated RMSE value
%TOTAL_RMSExy = calculated RMSE value in xy plain
%TOTAL_RMSEz = calculated RMSE value in z axis (height)
%TOTAL_smodch = Total STD value


clc; clear all; close all; format compact;

% Initial parameters
console_output = false; % choose to write output coordinates into console
progress_bar = true; % choose to put progress bar into console
progress_step = 100; % choose how often to display progress

% Measurement dimensions
state_dim = 3;    % [x, y, z] coordinates

% Load all anchor data
Anchs = load("Anchors_laborka.txt");

%Input true position of tag
true_pos = [0.6,3,1.66];


% Anchors
anchor_positions = Anchs(:, 1:3);
num_anchors = size(anchor_positions, 1); % Number of anchors

% Plot anchor positions
figure(1)
p1 = stem3(anchor_positions(:, 1), anchor_positions(:, 2), anchor_positions(:, 3), ...
    "Color", "#D95319", LineWidth=1.6, DisplayName="Anchor positions");
hold on;

% Initial guess of position
x = [1.5; 2.5; 1.5]; % approximate [x, y, z] guessed
% Plot initial guess
p2 = plot3(x(1), x(2), x(3), 'O', 'MarkerSize', 8, Color="#77AC30", ...
    MarkerFaceColor="#77AC30", DisplayName="Initial guess");

% Load measurements
z_matrix = load("E.txt"); % replace with chosen file of mesured values
z_matrix(:,1:4) = -z_matrix(:,1:4); %aadjust orientation of azimuth angles

% Adjust measurements with anchor rotation
rot = cat(2, Anchs(:, 4)', Anchs(:, 5)');
for i = 1:size(z_matrix, 1)
    z_matrix(i, :) = z_matrix(i, :) + rot;
end

z_matrix = deg2rad(z_matrix); % convert from degrees to radians

% Normalize z_matrix angles to [-pi, pi] not used here
% [m, n] = size(z_matrix);
% for m_i = 1:m
%     for n_i = 1:n
%         z_matrix(m_i, n_i) = mod(z_matrix(m_i, n_i) + pi, 2 * pi) - pi;
%     end
% end

% Newton-Raphson loop over all measurements
tol = 1e-4; % convergence tolerance
max_iter = 80; % maximum iterations for Newton-Raphson

for i = 1:size(z_matrix, 1)
    % Get i-th measurement
    z_i = z_matrix(i, :)';
    
    % Filter out NaNs
    valid_idx = ~isnan(z_i); % logical array of valid (non-NaN) indices
    z_i = z_i(valid_idx); % remove NaN values from z_i
    
    % Iterative Newton-Raphson process
    iter = 0;
    while iter < max_iter
        % Compute residual and Jacobian
        z_pred = h(x, anchor_positions, valid_idx);
        H = compute_jacobian(x, anchor_positions, valid_idx);
        
        % Compute update step
        residual = z_i - z_pred;
        
        % Normalization procedure
        [m, n] = size(residual);
        for m_i = 1:m
            for n_i = 1:n
                residual(m_i, n_i) = mod(residual(m_i, n_i) + pi, 2 * pi) - pi;
            end
        end
        %dx = (H' * H) \ (H' * residual);

        epsilon = 1e-2;  % Small regularization value
        dx = (H' * H + epsilon * eye(size(H,2))) \ (H' * residual);
       
        
        % Update estimate
        x = x + dx;
        
        % Check for convergence
        if norm(dx) < tol
            break;
        end
        iter = iter + 1;

    end
    
    % Display the updated position estimate
    if console_output
        fprintf('Estimated position %d: [%.2f, %.2f, %.2f]\n', i, x(1), x(2), x(3));
    end
    
    % Progress into console
    if progress_bar && (mod(i, progress_step) == 0)
        fprintf('Iteration %d of %d\n', i, size(z_matrix, 1));
    end
    
    % Completion of cycle into console
    if i == size(z_matrix, 1)
        fprintf('Iterations Complete\n');
    end
    
    % Save updated values of x
    x_all(i, :) = x;
end

% Plot calculated positions
p3 = plot3(x_all(:, 1), x_all(:, 2), x_all(:, 3), 'O', 'MarkerSize', 1.2, ...
    Color="#0072BD", MarkerFaceColor="#0072BD", DisplayName="Calculated positions");
axis equal
legend([p1, p2, p3], 'Location', 'best');
xlabel("x [m]");
ylabel("y [m]");
zlabel("z [m]");
title("Calculated positions (Newton-Raphson)");

% Sums for RMSE calculation
sum = 0;
sumxy = 0;
sumz = 0;

% RMSE calculation
for i = 1:(size(x_all,1))
    sum = sum + (x_all(i,1)-true_pos(1,1))^2 + (x_all(i,2)-true_pos(1,2))^2 + (x_all(i,3)-true_pos(1,3))^2;
    sumxy = sumxy + (x_all(i,1)-true_pos(1,1))^2 + (x_all(i,2)-true_pos(1,2))^2;
    sumz = sumz + (x_all(i,3)-true_pos(1,3))^2;
end

%RMSE results
TOTAL_RMSE = sqrt(1/(size(x_all,1))*sum);
TOTAL_RMSExy = sqrt(1/(size(x_all,1))*sumxy);
TOTAL_RMSEz = sqrt(1/(size(x_all,1))*sumz);


%STD calculation
size_result = size(x_all,1);
sumxo = 0;
sumyo = 0;
sumzo = 0;
for i = 1:size_result
    sumxo = sumxo + x_all(i,1);
    sumyo = sumyo + x_all(i,2);
    sumzo = sumzo + x_all(i,3);
end

avgxo = sumxo / size_result;
avgyo = sumyo / size_result;
avgzo = sumzo / size_result;

sumxokv = 0;
sumyokv = 0;
sumzokv = 0;

for i = 1:size_result
    sumxokv = sumxokv + (x_all(i,1)-avgxo)^2;
    sumyokv = sumyokv + (x_all(i,2)-avgyo)^2;
    sumzokv = sumzokv + (x_all(i,3)-avgzo)^2;
end

%Partial STDs
smodchx = sqrt((1/size_result)*sumxokv);
smodchy = sqrt((1/size_result)*sumyokv);
smodchz = sqrt((1/size_result)*sumzokv);

%Resulting STD
TOTAL_smodch = sqrt(smodchx^2+smodchy^2+smodchz^2);

%% Compute azimuth and elevation predictions (only for valid measurements)
function z_pred = h(x, anchor_positions, valid_idx)
    elevations = compute_elevation(x, anchor_positions);
    azimuths = compute_azimuth(x, anchor_positions);
    
    % Combine azimuth and elevation
    z_pred_full = [azimuths; elevations];
    z_pred = z_pred_full(valid_idx);
end

%% Jacobian matrix for valid measurements
function H = compute_jacobian(x, anchor_positions, valid_idx)
    num_anchors = size(anchor_positions, 1);
    H_full = zeros(2 * num_anchors, 3); % initialize full Jacobian
    
    % Compute partial derivatives for azimuth and elevation
    for i = 1:num_anchors
        dx = x(1) - anchor_positions(i, 1);
        dy = x(2) - anchor_positions(i, 2);
        dz = x(3) - anchor_positions(i, 3);
        
        r_xy = sqrt(dx^2 + dy^2);
        r = sqrt(dx^2 + dy^2 + dz^2);
        
        % Azimuth partials
        H_full(i, 1) = -dy / (dx^2 + dy^2);
        H_full(i, 2) = dx / (dx^2 + dy^2);
        H_full(i, 3) = 0;
        
        % Elevation partials
        H_full(i + num_anchors, 1) = -(dx * dz) / (r_xy * r^2);
        H_full(i + num_anchors, 2) = -(dy * dz) / (r_xy * r^2);
        H_full(i + num_anchors, 3) = r_xy / (r^2);
    end
    
    % Select valid rows of H
    H = H_full(valid_idx, :);
end

%% Predicted elevation angle from estimated position
function elevations = compute_elevation(x, anchor_positions)
    num_anchors = size(anchor_positions, 1);
    elevations = zeros(num_anchors, 1);
    for i = 1:num_anchors
        dx = x(1) - anchor_positions(i, 1);
        dy = x(2) - anchor_positions(i, 2);
        dz = x(3) - anchor_positions(i, 3);
        r_xy = sqrt(dx^2 + dy^2);
        elevations(i) = atan2(dz, r_xy);
    end
end

%% Predicted azimuth angle from estimated position
function azimuths = compute_azimuth(x, anchor_positions)
    num_anchors = size(anchor_positions, 1);
    azimuths = zeros(num_anchors, 1);
    for i = 1:num_anchors
        dx = x(1) - anchor_positions(i, 1);
        dy = x(2) - anchor_positions(i, 2);
        azimuths(i) = atan2(dy, dx);
    end
end
