%% Estimation of EEJ model parameters
clear all
rad = pi/180;

load EEJ_A_extended_1.mat
tmp = load('EEJ_B_extended_1');
t_all = [t; tmp.t]; phi_all = [phi; tmp.phi]; J_all = [J; tmp.J];

idx = t > md2000(2023, 2, 11);
t_all(idx) = []; phi_all(idx) = []; J_all(idx,:) = [];


nu_all = moon_phase(t_all);

% read daily values of F10.7
filename = 'SW_OPER_AUX_F10_2__20060101T000000_20231015T000000_0001.DBL';
[t_F107_all, F107_all]  = textread(filename, '%f %f', 'commentstyle', 'shell');

idx = ~isnan(F107_all);
F107 = interp1(t_F107_all(idx), zp_filter(F107_all(idx), 'boxcar', 81), t_all, 'linear'); % interpolate F107 to time instants of "synthetic" data
Fst_all = (F107 + interp1(t_F107_all(idx), F107_all(idx), t_all, 'linear'))/2;

n_t = 4; n_s = 4; n_phi = 4; n_l = 2;

M = 0; % number of coefficients a,b
p_all = []; m_all = []; k_all = []; l_all = []; ab_all = [];
for p = 0:n_t
    for m = -n_phi:n_phi
        for k = -n_s:n_s
            if n_l == 1
                l_sel = [0];
            elseif n_l == 2
                l_sel = [0 -2];
            end
            for l = l_sel
                if ~( (p == 0 && l == 0 && m < 0) || (p == 0 && l == 0 && m == 0 && k < 0) )
                    p_all = [p_all; p];
                    m_all = [m_all; m];
                    k_all = [k_all; k];
                    l_all = [l_all; l];
                    ab_all = [ab_all; 'a'];
                    M = M + 1;
                end
                if ~( (m < 0 && p == 0 && l == 0) || (p == 0 && m == 0 && l == 0 && k < 0) || (p == 0 && m == 0 && l == 0 && k == 0) )
                    p_all = [p_all; p];
                    m_all = [m_all; m];
                    k_all = [k_all; k];
                    l_all = [l_all; l];
                    ab_all = [ab_all; 'b'];
                    M = M + 1;
                end
            end % l
        end % k
    end % m
end % p

m_i_all = zeros(M+1, 81);
rms_all = zeros(81,1);

for k_QD = 1:81
    disp(['k_QD:' num2str(k_QD)]);
    J = J_all(:,k_QD);
    t = t_all;
    phi = phi_all;
    nu = nu_all;
    Fst = Fst_all;

    % remove outliers
    q_1 = prctile(J,1);
    q_99 = prctile(J,99);
    dq = q_99 - q_1;
    idx = q_1-1.5*dq < J & J < q_99 + 1.5*dq & ~isnan(J);
    t = t(idx); phi = phi(idx); J = J(idx); nu = nu(idx); Fst = Fst(idx);

    T = mod(t+phi/360,1)*24;

    N = length(t);
    G = zeros(N, M);
    n = 1;
    s_0 = md2000(2000, 3, 22);
    for p = 0:n_t
        for m = -n_phi:n_phi
            for k = -n_s:n_s
                if n_l == 1
                    l_sel = [0];
                elseif n_l == 2
                    l_sel = [0 -2];
                end
                for l = l_sel
                    % coefficients a
                    if ~( (p == 0 && l == 0 && m < 0) || (p == 0 && l == 0 && m == 0 && k < 0) )
                        G(:,n) = cosd(p*t*360 + m*phi + k*(t-s_0)*360/365.25 + l*nu*15);
                        n = n + 1;
                    end
                    % coefficients b
                    if ~( (m < 0 && p == 0 && l == 0) || (p == 0 && m == 0 && l == 0 && k < 0) || (p == 0 && m == 0 && l == 0 && k == 0) )
                        G(:,n) = sind(p*t*360 + m*phi + k*(t-s_0)*360/365.25 + l*nu*15);
                        n = n + 1;
                    end
                end % l
            end % k
        end % m
    end % p

    % fixed value of R
    R = 0.0213;
    G = (1+R*Fst).*G;

    % IRLS
    m_i = zeros(M, 1);
    Max_iter = 10;
    for iter = 1:Max_iter
        d_mod = G*m_i;
        delta_d = [J - d_mod];

        if iter > 1
            rms = sqrt((delta_d' * (w.*delta_d)) / sum(w));
            w = min(1.5./abs(delta_d./rms),1);  % robust weights
        else
            rms = sqrt((delta_d' * (delta_d)) / length(t));
            w = ones(size(delta_d));
        end
        wG = w.*G;
        delta_m = (G'*wG)\(wG'*delta_d);
        m_i = m_i + delta_m;
        %     fprintf(2,'Iteration: %3i, rms:%10.4f\n', iter, rms)
    end

    m_i = [m_i; R];

    % % non-linear inversion: co-estimation of R
    % m_IRLS = m_i;
    %
    % F107 = Fst;
    %
    % m_i = [m_IRLS; 15e-3];
    % Max_iter = 10;
    % for iter = 1:Max_iter
    %     d_mod = G*m_i(1:end-1);
    %     G_F = [G.*(1+m_i(end)*F107) F107.*d_mod];
    %     d_mod = (1+m_i(end)*F107).*G*m_i(1:end-1);
    %
    %     delta_d = [J - d_mod];
    %
    %     if iter > 1
    %         rms = sqrt((delta_d' * (w.*delta_d)) / sum(w));
    %         w = min(1.5./abs(delta_d./rms),1);  % robust weights
    %     else
    %         rms = sqrt((delta_d' * (delta_d)) / length(t));
    %         w = ones(size(delta_d));
    %     end
    %     wG = w.*G_F;
    %     delta_m = (G_F'*wG)\(wG'*delta_d);
    %     m_i = m_i + delta_m;
    %     fprintf(2,'Iteration: %3i, rms:%10.4f\n', iter, rms)
    % end

    m_i_all(:,k_QD) = m_i;
    rms_all(k_QD) = rms;
end % k_QD

% save the model
QD_lat = [-20:.5:20]';
N_coeff = length(p_all);
n_QD = length(QD_lat);

header = char('# EEJ model based on Swarm A and B EEF current profile data', ...
    ['# from 25 Nov 2013 to 22 Feb 2023'], ...
    ['# Solar activity dependence described by EUVAC, fixed value R = 0.0213'], ...
    ['# A_1_extended, B_1_extended'], ...
    ['# deGeeter et al., (2023)'], ...
    ['# Model file written on ', datestr(now)]);

filename_out = 'EEJ_SW-AB_06.txt';
fid = fopen(filename_out, 'w');

% write header comment lines
for k = 1:size(header, 1)
    fprintf(fid, '%s\n', header(k,:));
end
% write the two header lines
fprintf(fid, '# n_t n_phi   n_s   n_l  n_QD\n');
fprintf(fid, '%5d %5d %5d %5d %5d\n', n_t, n_phi, n_s, n_l, n_QD);
% write QD_latitudes
fprintf(fid, '# QD-latitudes\n');
fprintf(fid, '                   ');
for k = 1:n_QD
    fprintf(fid, ' %10.3f', QD_lat(k));
end
fprintf(fid, '\n');
% write rms
fprintf(fid, '# rms data misfit  ');
for k = 1:n_QD
    fprintf(fid, ' %10.3e', rms_all(k));
end
fprintf(fid, '\n');
% write R
fprintf(fid, '# Solar flux dependencies R\n');
fprintf(fid, '                   ');
for k = 1:n_QD
    fprintf(fid, ' %10.3e', m_i_all(end,k));
end
fprintf(fid, '\n');

fprintf(fid, '# p   m   k   l a/b     coefficients\n');
% write  coefficients
for n = 1:N_coeff
    fprintf(fid, '%3d %3d %3d %3d %3s', p_all(n), m_all(n), k_all(n), l_all(n), ab_all(n));
    for k = 1:n_QD
        fprintf(fid, ' %10.3e', m_i_all(n,k));
    end % k
    fprintf(fid, '\n');
end % n
fclose(fid);


%% rms in dependence of LT
rms_LT_F = zeros(24,1);
T = mod(t+phi/360, 1)*24;
for LT = 1:24
    idx = abs(LT - T - .5) < .5;
    rms_LT_F(LT) = sqrt(delta_d(idx)'*(w(idx).*delta_d(idx))/sum(w(idx)));
end

plot([.5:24], rms_LT, [.5:24], rms_LT_F, 'LineWidth', 2)
xlabel 'LT [hours]'
ylabel 'rms [A/km]'
grid on
set(gca, 'XTick', [0:3:24])
legend('no dependence on F', 'F dependence', 'Location','south')
axis([0 24 0 Inf])

