% This script runs simulations of the original and updated Na/K ATPase
% models, and compares them to the equations in the LRd00 model (Figure 4).

clear;
clc;
close all;

%% Directory names
current_dir = cd;
Idx_backslash = find(current_dir == filesep);
main_dir = current_dir(1:Idx_backslash(end));
data_dir = [main_dir 'data' filesep];
output_dir = [main_dir 'output' filesep];
storage_dir = [main_dir 'storage' filesep];

%% Define constants and options
R = 8.314;
T = 310;
F = 96485;

W_i = 38;
W_e = 5.182;
W_total = W_i + W_e;

print_figures = true;

%% Load stoichiometric information and constants
load([storage_dir 'kinetic_fitting_results.mat']);

%% Run simulation of Na/K pump (using 1 fmol of Na/K pump)
NaK_density = params_vec(15); % unit um^-2
NaK_density_original = 1913.29; % unit um^-2

% Load AP waveform from Faber and Rudy (2000)
CellML_data = csvread([data_dir 'faber_rudy_2000_data.csv'],1,0);

t_LRd = CellML_data(:,1); % unit ms
V_LRd = CellML_data(:,5); % unit mV

t_start = -100;

% Restrict to an action potential
t_LRd = t_LRd(t_LRd <= 400) + t_start;
V_LRd = V_LRd(t_LRd <= 400);

struct_input.V = -0.08;
struct_input.MgATP = 6.95;
struct_input.Pi_total = 0.8;
struct_input.Nai = 10;
struct_input.Nae = 140;
struct_input.Ke = 5.4;
struct_input.Ki = 145;
struct_input.H_conc = 10^(-4.095);
struct_input.MgADP = 0.035;
struct_input.T = 310;

v_NaK = zeros(length(t_LRd),1);
v_NaK_original = zeros(length(t_LRd),1);
for i_time = 1:length(t_LRd)
    V_m = V_LRd(i_time)/1000;
    struct_input.V = V_m;
    v_NaK(i_time) = NaK_fitting_vcyc(params_vec,struct_input);
    v_NaK_original(i_time) = Terkildsen_NaK_original(struct_input);
end

I_NaK = NaK_density*v_NaK*1.6e-5; % unit uA/cm^2
I_NaK_original = NaK_density_original*v_NaK_original*1.6e-5; % unit uA/cm^2

%% Run a simulation of the LRd Na/K pump under the same conditions
I_NaK_max = 2.25; % unit uA/uF
K_mNai = 10; % unit mM
K_mKo = 1.5; % unit mM

V_m = V_LRd/1000;
sigma = 1/7*(exp(struct_input.Nae/67.3)-1);
f_NaK = 1./(1+0.1245*exp(-0.1*V_m*F/(R*T))+0.0365*sigma*exp(-V_m*F/(R*T)));
I_NaK_LRd = I_NaK_max*f_NaK*1/(1+(K_mNai/struct_input.Nai)^2)*struct_input.Ke/(struct_input.Ke+K_mKo);

%% Compare currents of original and updated models
h1 = figure;
plot(t_LRd,I_NaK,'r','LineWidth',3);
hold on;
plot(t_LRd,I_NaK_original,'b','LineWidth',3);
ylim([0.1 0.25])
xlabel('Time (ms)');
ylabel('Current density ({\mu}A/cm^2)');
legend('Updated','Original');
set(gca,'FontSize',22);

%% Compare currents of original and updated models when scaled to match LRd currents
NaK_density_scaled = 3.4*NaK_density;
NaK_density_original_scaled = 3.9*NaK_density_original;

I_NaK_scaled = NaK_density_scaled*v_NaK*1.6e-5; % unit uA/cm^2
I_NaK_original_scaled = NaK_density_original_scaled*v_NaK_original*1.6e-5; % unit uA/cm^2

h2 = figure;
plot(t_LRd,I_NaK_scaled,'r','LineWidth',3);
hold on;
plot(t_LRd,I_NaK_original_scaled,'b','LineWidth',3);
plot(t_LRd,I_NaK_LRd,'k','LineWidth',3);
ylim([0.4 0.85])
xlabel('Time (ms)');
ylabel('Current density ({\mu}A/cm^2)');
legend('Updated','Original','LRd00');
set(gca,'FontSize',22);

%% Plot the membrane voltage waveform
h3 = figure;
plot(t_LRd,V_LRd,'k','LineWidth',3);
xlabel('Time (ms)');
ylabel('Voltage (mV)');
set(gca,'FontSize',22);

%% Print figures
if print_figures
    print_figure(h1,output_dir,'NaK_current');
    print_figure(h2,output_dir,'NaK_current_scaled');
    print_figure(h3,output_dir,'LRd_AP_waveform');
end