% Fit bond graph parameters to model

% Use Pan functions to calculate alpha beta fits

% params: [kf, zf, kr, zr];

% units: mV


tic;
clear;

storage_dir = 'storage/';

zero_IC = true;

V = transpose(-120:1:60); % unit mV
R = 8.314;
T = 310;
F = 96485;
Ko = 4.5; 

% alpha and beta are swapped around in _i, _mu. 
alpha = 55.5*exp(0.05547153*(V-12));
beta = 2.357*exp(-0.036588*V);
alpha_in = 2.172*ones(size(V));
beta_in = 1.077*ones(size(V));
alpha_alpha = 65.5*exp(0.05547153*(V-36));
beta_beta = 2.9357*exp(-0.02158*V);
beta_i = 439*exp(-0.02352*(V+25))*4.5/Ko;
alpha_i = 656*exp(0.000942*V).*((4.5/Ko).^ 0.3);
mu = alpha_i.*beta_beta.*alpha_alpha./(alpha_alpha.*beta_i);        

alpha = [alpha alpha_in alpha_alpha alpha_i alpha_alpha]; 
beta = [beta beta_in beta_beta beta_i mu];

nsize = 500; % swarm size
nvars = 4;

% initialise guess
x0v = [10497.186	0.50706377	191.70055	-1.3158781
9297.8311	0.50917998	255.85905	-1.3158778
8.48E+03	0.52553345	360.04979	-1.3158778
4.39E-07	-5.1036976	7.6259274	0.11640022
9.18E+03	0.90001141	0.036530664	-0.80680085
];

if zero_IC
    x0v = zeros(size(x0v)); 
end

x0v_og = x0v;

for  j = 1:4 % how many iterations until convergence?
    for i = 1:size(alpha,2)

        error_func_alpha = @(params) square_error(alpha(:,i) - calc_alpha(params,V/1000)); % fit to alpha - k.exp(zFV/RT) : do lsqminerror.
        error_func_beta = @(params) square_error(beta(:,i) - calc_beta(params,V/1000));

        error_func = @(params) error_func_alpha(params) + error_func_beta(params);

        A = [];
        b = [];
        Aeq = [];
        beq = [];
        lb = [eps; -10; eps; -10];
        ub = [Inf; 10; Inf; 10];
        
        x0 = x0v(i,:).*ones(nsize, nvars);
        options_unc = optimoptions('fminunc','MaxFunEvals',5000);
        options_ps = optimoptions('particleswarm','UseParallel',false,'HybridFcn',@fmincon, ...
            'SwarmSize', nsize, 'InitialSwarmMatrix',x0);
        
        [params_vec(i,:,j),fval,exitflag,output] = particleswarm(error_func,nvars,lb,ub,options_ps);
        
    end
   
    zr_mu = params_vec(5,2,j) + params_vec(4,2,j) + params_vec(3,2,j) - ...
        (params_vec(4,4,j) + params_vec(3,4,j));
    
    % DETAILED BALANCE
    if false
        kr_b2 = params_vec(5,1,j)*params_vec(4,1,j)*params_vec(3,1,j)/...
            (params_vec(4,3,j)*params_vec(3,3,j));
    else
        % maybe need to incorporate e(z.F.V/RT) into kr_b2
        af_c1o = calc_alpha(params_vec(3,:,j),V/1000);
        af_oi = calc_alpha(params_vec(4,:,j),V/1000);
        af_ic1 = calc_alpha(params_vec(5,:,j),V/1000);
        br_io = calc_beta(params_vec(4,:,j),V/1000);
        br_oc1 = calc_beta(params_vec(3,:,j),V/1000);
        
        br_mu = af_c1o.*af_oi.*af_ic1./(br_io.*br_oc1);
        kr_mu = mean(br_mu ./ exp(zr_mu*F*V*1e-3/(R*T)));
    end
    
    params_vec(5,[3 4],j) = [kr_mu zr_mu];
    
    x0v = params_vec(:,:,j);
end


params_vec_flat = x0v_og;
for jj = 1:j
    params_vec_flat = [params_vec_flat; params_vec(:,:,jj)];
end
alpha = params_vec(:,1,end);
zf = params_vec(:,2,end);
beta = params_vec(:,3,end);
zr = params_vec(:,4,end);
if zero_IC
    save([storage_dir 'Kr_alphabetaz_zeroIC.mat'],'alpha','zf','beta','zr');
else
    save([storage_dir 'Kr_alphabetaz_someIC.mat'],'alpha','zf','beta','zr');
end

dlmwrite('output/pso_out_v6_exp_k.txt', params_vec_flat, 'delimiter', '\t','precision',8);

toc;