% Fit bond graph parameters to model

% Use Pan functions to calculate alpha beta fits

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

% units: mV

% try all z = 1 so fit kf and kr only. Use detailed balance.

tic;
clear;

storage_dir = 'storage/';

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

pso = true; % do lsqcurvefit if no pso


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 = 250; % swarm size
nvars = 2;

% initialise guess
x0v = [10497.18514	191.7007552   % [=] 1/s
9297.831061	255.8590544
8475.49086	360.0495824
3.79E-07	8.347303336
9.18E+03	0.00981794
];

% z_alpha = [1 1 1 1 1];
% z_beta = [-1 -1 -1 -1 5];

z_alpha = [1 0 1 0 1];
z_beta = [-1 0 -1 -1 4];

% x0v = zeros(size(x0v)); 

x0v_og = x0v;

% loop of PSO iteratively, replacing the initial guess for alpha and beta
% with soln from previous iteration.
for  j = 1:4 % how many iterations until convergence?
    for i = 1:size(alpha,2)
        if pso
            error_func_alpha = @(params) square_error(alpha(:,i) - calc_alpha_z1(params,V/1000,z_alpha(i))); % fit to alpha - k.exp(zFV/RT) : do lsqminerror.
            error_func_beta = @(params) square_error(beta(:,i) - calc_beta_z1(params,V/1000,z_beta(i)));

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

            A = [];
            b = [];
            Aeq = [];
            beq = [];
            lb = [0e6; 0e6;];
            ub = [1e5;1e5];

            xIC = x0v(i,:).*ones(nsize, nvars);
            options_unc = optimoptions('fminunc','MaxFunEvals',2000);
            options_ps = optimoptions('particleswarm','UseParallel',false,'HybridFcn',@fmincon, ...
                'SwarmSize', nsize, 'InitialSwarmMatrix',xIC);

            [params_vec(i,:,j),fval,exitflag,output] = particleswarm(error_func,nvars,lb,ub,options_ps);
        
        else 
            xdata = V/1000;
            lsqoptions = optimoptions('lsqcurvefit','FunctionTolerance',1e-6);
            fun = @(kfr, xdata)kfr*exp(1*F*xdata/(R*T));
            xIC = x0v(i,:);

            % fit to existing kf kr
            alpha_kz = calc_alpha_z1(x0v,V/1000,z_alpha(i));
            beta_kz = calc_beta_z1(x0v,V/1000,z_beta(i));
            ydata_alpha = alpha_kz(:,i);
            ydata_beta = beta_kz(:,i);
            
            params_vec(i,1,j) = lsqcurvefit(fun,xIC(1),xdata,ydata_alpha,eps,inf,lsqoptions);
            if i == 5
%                 ydata_beta = alpha_kz(:,3).*alpha_kz(:,4).*alpha_kz(:,5)./(beta_kz(:,4).*beta_kz(:,3));
                params_vec(i,2,j) = params_vec(5,1,j)*params_vec(4,1,j)*params_vec(3,1,j)/...
                (params_vec(4,2,j)*params_vec(3,2,j));
            else
                params_vec(i,2,j) = lsqcurvefit(fun,xIC(2),xdata,ydata_beta,eps,inf,lsqoptions);
            end

        end
    end
   
    % detailed balance for remaining beta in closed loop
    if true
        if false
            kr_mu = params_vec(5,1,j)*params_vec(4,1,j)*params_vec(3,1,j)/...
                (params_vec(4,2,j)*params_vec(3,2,j));
        else
            
            af_c1o = calc_alpha_z1(params_vec(3,:,j),V/1000,z_alpha(3));
            af_oi = calc_alpha_z1(params_vec(4,:,j),V/1000,z_alpha(4));
            af_ic1 = calc_alpha_z1(params_vec(5,:,j),V/1000,z_alpha(5));
            br_io = calc_beta_z1(params_vec(4,:,j),V/1000,z_beta(4));
            br_oc1 = calc_beta_z1(params_vec(3,:,j),V/1000,z_beta(3));
            
            br_mu = af_c1o.*af_oi.*af_ic1./(br_io.*br_oc1);
            % kr_mu should equal kr_b2 = params_vec(5,1,j)*params_vec(4,1,j)*params_vec(3,1,j)/(params_vec(4,2,j)*params_vec(3,2,j)); 
            kr_mu = mean(br_mu ./ exp(z_beta(5)*F*V*1e-3/(R*T))); 
            
        end

        params_vec(5,2,j) = [kr_mu];
    end
    
    x0v = params_vec(:,:,j);
end

% params_vec_flat: a list of all alpha and beta found using PSO iteratively
params_vec_flat = x0v_og;
if false
    save([storage_dir 'Kr_all_parameters.mat'],'params_vec');
else
    for jj = 1:j
        params_vec_flat = [params_vec_flat; params_vec(:,:,jj)];
    end
    alpha = params_vec(:,1,end);
    beta = params_vec(:,2,end);
    save([storage_dir 'Kr_alphabeta.mat'],'alpha','beta');
    
    dlmwrite('output/pso_out_v3_detailedbalance_zconserved.txt', params_vec_flat, 'delimiter', '\t','precision',8);
end

toc;