clear all
rad = pi/180;
a = 6371.2;

% CCDB values of Euler angles
q_A = [-0.013237940900440 0.618884832920591 0.022645867500753 0.785043619926120]; % Swarm Alpha
q_B = [0.011655028800383 0.619619452020370 -0.014490071300476 0.784682058425797]; % Swarm Bravo
q_C = [-0.001679525100088 0.621376719032451 0.003690655100193 0.783501455840917]; % Swarm Charlie

% read 4-9 May 2021 Swarm Alpha
year = 2021; month = 5;
t = []; r = []; theta = []; phi = []; B_NEC = []; F = []; B_VFM = []; q_NEC_CRF = [];
Step = 30; % downsample data to 30 sec
for day = 4:9
    d_B = dir(['SW_OPER_MAGA_LR_1B_' num2str(year, '%04d') num2str(month, '%02d') num2str(day, '%02d') '*.cdf']);
    filename_in = [d_B(1).name];
    data = cdfread(filename_in, 'CombineRecords', 1, 'variable', {'Timestamp', 'Latitude', 'Longitude', ...
        'Radius', 'B_NEC', 'F', 'B_VFM', 'q_NEC_CRF'});
    t = [t; todatenum(data{1}(1:Step:end)) - 730486];
    r = [r; double(data{4}(1:Step:end))/1e3]; % radius in km
    theta = [theta; 90-double(data{2}(1:Step:end))];
    phi = [phi; double(data{3}(1:Step:end))];
    B_NEC = [B_NEC; double(data{5}(1:Step:end,:))];
    F = [F; double(data{6}(1:Step:end,:))];
    B_VFM = [B_VFM; double(data{7}(1:Step:end,:))];
    q_NEC_CRF = [q_NEC_CRF; double(data{8}(1:Step:end,:))];
end

B_r = -B_NEC(:,3); B_theta = -B_NEC(:,1); B_phi = +B_NEC(:,2);

% CHAOS model
load ('\nio\m\mmod\CHAOS-7.8.mat')
gh_CHAOS = fnval(pp, md2000(2021, 5, 1)); % CHAOS Gauss coefficients for 1 May 2021.



%% simple field modeling: estimate a static (snapshot) magnetic field from 6 days of Swarm Alpha data in B_NEC
%%% first example: estimate static internal field model entirely from B_r (all latitudes)

% only B_r at polar latitudes, otherwise all three vector components
theta_d = gg2gm(theta, phi);
idx = abs(90-theta_d) < 90;

N = 20; % maximum spherical harmonic degree
A_r = design_SHA(r/a, theta*rad, phi*rad, N, 'int');

m_i = A_r\B_r;

semilogy([1:N], powerspec(m_i), [1:N], powerspec(gh_CHAOS), [1:N], powerspec(m_i - gh_CHAOS))
xlabel 'degree \itn'
ylabel 'R_n [nT^2]'
grid on
legend('estimated model', 'CHAOS-7.8', 'spectrum of difference')

%% second example: estimate static internal field model from B_r (all latitudes) and (B_theta, B_phi) at non-polar latitudes

N_coeff_MF = N*(N+2); % number of internal field Gauss coefficients
N_coeff_ext = 3; % number of external field Gauss coefficients
N_coeff = N_coeff_MF + N_coeff_ext; % number of SHA coefficients

[A_r, A_theta, A_phi] = design_SHA(r/a, theta*rad, phi*rad, N, 'int');
[A_r_ext, A_theta_ext, A_phi_ext] = design_SHA(r/a, theta*rad, phi*rad, 1, 'ext');
A_r = [A_r A_r_ext];
A_theta = [A_theta A_theta_ext];
A_phi = [A_phi A_phi_ext];

G = [A_r; A_theta(idx,:); A_phi(idx,:)];
d = [B_r; B_theta(idx); B_phi(idx)];
m_i = G\d;

semilogy([1:N], powerspec(m_i(1:N_coeff_MF)), [1:N], powerspec(gh_CHAOS), [1:N], powerspec(m_i(1:N_coeff_MF) - gh_CHAOS))
xlabel 'degree \itn'
ylabel 'R_n [nT^2]'
grid on
legend('estimated model', 'CHAOS-7.8', 'spectrum of difference')

%% third example: as 2nd example, but IRLS

w_vec = [sind(theta); sind(theta); sind(theta)]/sigma_VFM; % data weightsto simulate equal area distribution
N_data = size(w_vec, 1);

Max_iter = 10;
step = 1;
m_i = zeros(size(G, 2), 1);
m_i(1) = -30000; % starting model
for iter = 1:Max_iter
    B_mod_r = A_r*m_i;
    B_mod_theta = A_theta*m_i;
    B_mod_phi = A_phi*m_i;
    delta_d_vec = [B_r(:) - B_mod_r(:); ...
        B_theta(idx) - B_mod_theta(idx); ...
        B_phi(idx) - B_mod_phi(idx)];
    if iter > 1
        rms = sqrt(((w'.*delta_d_vec') * (w.*delta_d_vec)) / sum(w));
        w = min(1.5./abs(delta_d_vec./rms),1);  % robust weights
    else
        w = ones(size(delta_d_vec));
    end
    iC_y = spdiags(w.*w_vec, 0, N_data, N_data);
    wG = iC_y*G;
    delta_m = (G'*wG)\(wG'*delta_d_vec);
    m_i = m_i + step*delta_m;
    fprintf(2,'Iteration: %3i, rms (nT):%10.4f\n', iter, rms_std(delta_d_vec, w))
end

semilogy([1:N], powerspec(m_i(1:N_coeff_MF)), [1:N], powerspec(gh_CHAOS), [1:N], powerspec(m_i(1:N_coeff_MF) - gh_CHAOS))
xlabel 'degree \itn'
ylabel 'R_n [nT^2]'
grid on
legend('estimated model', 'CHAOS-7.8', 'spectrum of difference')


%% 4th example: coestimation of field model and Euler angles using IRLS
% using Swarm Alpha magnetic vector data B_VFM

% R_rtp is matrix which transforms from CRF (star imager acommon reference frame) to (r, theta, phi) frame
R_rtp = squeeze(qu2arot(qmult(q_NEC_CRF, qnorm(dcm2q([0 0 -1; -1 0 0; 0 1 0])))));

Max_iter = 20;
step = 1;

sigma_VFM = 2; % error of vector data

N = 20; % maximum spherical harmonic degree
N_para = N_coeff + 3; % total number of model parameters (including Euler angles)

m_i = zeros(N_coeff, 1); m_i(1) = -30000; % starting model: axial dipole
alpha = 11.8; beta = 283.8074; gamma = 347.4419; % starting values
% alpha = 1v0; beta = 270; gamma = 350; % convergence even with these starting values
m_i = [m_i; alpha*rad; beta*rad; gamma*rad]; % add Euler angles

% a-priori model
m_prior = m_i;
m_weights = zeros(size(m_i));
iC_prior = spdiags(m_weights, 0, N_para, N_para);

w_vec = [sind(theta); sind(theta); sind(theta)]/sigma_VFM; % data weights
N_data = size(w_vec, 1);
for iter = 1:Max_iter
    B_mod_r             = A_r     * m_i(1:N_coeff);
    B_mod_theta         = A_theta * m_i(1:N_coeff);
    B_mod_phi           = A_phi   * m_i(1:N_coeff);

    [R_VFM, R_a, R_b, R_c] = euler_123_SCARF(m_i(N_coeff+1), m_i(N_coeff+2), m_i(N_coeff+3)); % rotation matrix VFM -> CRF

    B_rtp = mat_mul_mat(R_rtp, mat_mul_mat_3(R_VFM, B_VFM), -1, -1);
    dB_dalpha = mat_mul_mat(R_rtp, mat_mul_mat_3(-R_a, B_VFM), -1, -1);
    dB_dbeta  = mat_mul_mat(R_rtp, mat_mul_mat_3(-R_b, B_VFM), -1, -1);
    dB_dgamma = mat_mul_mat(R_rtp, mat_mul_mat_3(-R_c, B_VFM), -1, -1);

    dB_r = B_rtp(:,1) - B_mod_r;
    dB_theta = B_rtp(:,2) - B_mod_theta;
    dB_phi = B_rtp(:,3) - B_mod_phi;

    delta_d_vec = [dB_r; dB_theta; dB_phi];

    if iter > 1
        rms = sqrt(((w'.*delta_d_vec') * (w.*delta_d_vec)) / sum(w));
        w = min(1.5./abs(delta_d_vec./rms),1);  % robust weights
    else
        w = ones(size(delta_d_vec));
    end

    G = [[A_r; A_theta; A_phi] dB_dalpha(:) dB_dbeta(:) dB_dgamma(:)];
    iC_y = spdiags(w.*w_vec, 0, N_data, N_data);
    Gty = (iC_y*G)'*[dB_r; dB_theta; dB_phi];
    GtG = G'*iC_y*G;

    delta_m = (GtG + iC_prior)\(Gty + iC_prior*(m_prior-m_i));

    m_i = m_i + step*delta_m;

    fprintf(2,'Iteration: %3i, rms (nT):%10.4f\n', iter, rms_std(delta_d_vec, w))
end

semilogy([1:N], powerspec(m_i(1:N_coeff_MF)), [1:N], powerspec(gh_CHAOS), [1:N], powerspec(m_i(1:N_coeff_MF) - gh_CHAOS))
xlabel 'degree \itn'
ylabel 'R_n [nT^2]'
grid on
legend('estimated model', 'CHAOS-7.8', 'spectrum of difference')

fprintf(2,'%6.2f %6.2f %6.2f %8.4f %8.4f %8.4f %6.2f\n', alpha, beta, gamma, mod(m_i(N_coeff+1:end)/rad, 360), rms_std(delta_d_vec, w));
