%Eddy energy budget and eddy eneregy of different component.
%h5disp('scalers/scalers_s1.h5')
ky_read = h5read('scalers/scalers_s1.h5','/scales/ky');
k_total = [ky_read(1:256);2*ky_read(256)-ky_read(255);ky_read(257:end)];
kx = k_total';
ky = k_total;
psi1_time = h5read('scalers/scalers_s1.h5','/tasks/psi1');
psi2_time = h5read('scalers/scalers_s1.h5','/tasks/psi2');
time_series = h5read('scalers/scalers_s1.h5','/scales/sim_time');
%quadratic_drag_coefficient = 0.03;
u1_square_time = zeros(1000,1);
v1_square_time = zeros(1000,1);
u2_square_time = zeros(1000,1);
v2_square_time = zeros(1000,1);
u_barotropic_square_time = zeros(1000,1);
u_baroclinic_square_time = zeros(1000,1);
v_barotropic_square_time = zeros(1000,1);
v_baroclinic_square_time = zeros(1000,1);
tau_baroclinic_square_time = zeros(1000,1);
psi_barotropic_square_time = zeros(1000,1);
for i = 1:1000
    psi1 = psi1_time(:,:,i);
    psi2 = psi2_time(:,:,i);
    psi1_hat = fft2(psi1);
    psi2_hat = fft2(psi2);
    u1_hat = -1i*ky.*psi1_hat;
    v1_hat = 1i*kx.*psi1_hat;
    u2_hat = -1i*ky.*psi2_hat;
    v2_hat = 1i*kx.*psi2_hat;
    u1 = real(ifft2(u1_hat));
    v1 = real(ifft2(v1_hat));
    u2 = real(ifft2(u2_hat));
    v2 = real(ifft2(v2_hat));
    u_barotropic = 0.5*(u1 + u2);
    u_baroclinic = 0.5*(u1 - u2);
    v_barotropic = 0.5*(v1 + v2);
    v_baroclinic = 0.5*(v1 - v2);
    tau = 0.5*(psi1 - psi2);
    psi = 0.5*(psi1 + psi2);
    u1_square_time(i) = mean(u1.^2,'all');
    v1_square_time(i) = mean(v1.^2,'all');
    u2_square_time(i) = mean(u2.^2,'all');
    v2_square_time(i) = mean(v2.^2,'all');
    u_barotropic_square_time(i) = mean(u_barotropic.^2,'all');
    u_baroclinic_square_time(i) = mean(u_baroclinic.^2,'all');
    v_barotropic_square_time(i) = mean(v_barotropic.^2,'all');
    v_baroclinic_square_time(i) = mean(v_baroclinic.^2,'all');
    tau_baroclinic_square_time(i) = mean(tau.^2,'all');
    psi_barotropic_square_time(i) = mean(psi.^2,'all');
end


u1_square_time_total = [u1_square_time];
v1_square_time_total = [v1_square_time];
u2_square_time_total = [u2_square_time];
v2_square_time_total = [v2_square_time];
u_barotropic_square_time_total = [u_barotropic_square_time];
v_barotropic_square_time_total = [v_barotropic_square_time];
u_baroclinic_square_time_total = [u_baroclinic_square_time];
v_baroclinic_square_time_total = [v_baroclinic_square_time];
tau_baroclinic_square_time_total = [tau_baroclinic_square_time];
psi_barotropic_square_time_total = [psi_barotropic_square_time];
time_series_total = [time_series];


for j = 2:50
    disp(j)
    filename = sprintf("scalers/scalers_s%d.h5",j);
    psi1_time = h5read(filename,'/tasks/psi1');
    psi2_time = h5read(filename,'/tasks/psi2');
    time_series = h5read(filename,'/scales/sim_time');
    u1_square_time = zeros(1000,1);
    v1_square_time = zeros(1000,1);
    u2_square_time = zeros(1000,1);
    v2_square_time = zeros(1000,1);
    u_barotropic_square_time = zeros(1000,1);
    u_baroclinic_square_time = zeros(1000,1);
    v_barotropic_square_time = zeros(1000,1);
    v_baroclinic_square_time = zeros(1000,1);
    tau_baroclinic_square_time = zeros(1000,1);
    psi_barotropic_square_time = zeros(1000,1);
    for i = 1:1000
        psi1 = psi1_time(:,:,i);
        psi2 = psi2_time(:,:,i);
        psi1_hat = fft2(psi1);
        psi2_hat = fft2(psi2);
        u1_hat = -1i*ky.*psi1_hat;
        v1_hat = 1i*kx.*psi1_hat;
        u2_hat = -1i*ky.*psi2_hat;
        v2_hat = 1i*kx.*psi2_hat;
        u1 = real(ifft2(u1_hat));
        v1 = real(ifft2(v1_hat));
        u2 = real(ifft2(u2_hat));
        v2 = real(ifft2(v2_hat));
        u_barotropic = 0.5*(u1 + u2);
        u_baroclinic = 0.5*(u1 - u2);
        v_barotropic = 0.5*(v1 + v2);
        v_baroclinic = 0.5*(v1 - v2);
        tau = 0.5*(psi1 - psi2);
        psi = 0.5*(psi1 + psi2);
        u1_square_time(i) = mean(u1.^2,'all');
        v1_square_time(i) = mean(v1.^2,'all');
        u2_square_time(i) = mean(u2.^2,'all');
        v2_square_time(i) = mean(v2.^2,'all');
        u_barotropic_square_time(i) = mean(u_barotropic.^2,'all');
        u_baroclinic_square_time(i) = mean(u_baroclinic.^2,'all');
        v_barotropic_square_time(i) = mean(v_barotropic.^2,'all');
        v_baroclinic_square_time(i) = mean(v_baroclinic.^2,'all');
        tau_baroclinic_square_time(i) = mean(tau.^2,'all');
        psi_barotropic_square_time(i) = mean(psi.^2,'all');
    end
    u1_square_time_total = cat(1,u1_square_time_total,u1_square_time);
    v1_square_time_total = cat(1,v1_square_time_total,v1_square_time);
    u2_square_time_total = cat(1,u2_square_time_total,u2_square_time);
    v2_square_time_total = cat(1,v2_square_time_total,v2_square_time);
    u_barotropic_square_time_total = cat(1,u_barotropic_square_time_total,u_barotropic_square_time);
    u_baroclinic_square_time_total = cat(1,u_baroclinic_square_time_total,u_baroclinic_square_time);
    v_barotropic_square_time_total = cat(1,v_barotropic_square_time_total,v_barotropic_square_time);
    v_baroclinic_square_time_total = cat(1,v_baroclinic_square_time_total,v_baroclinic_square_time);
    tau_baroclinic_square_time_total = cat(1,tau_baroclinic_square_time_total,tau_baroclinic_square_time);
    psi_barotropic_square_time_total = cat(1,psi_barotropic_square_time_total,psi_barotropic_square_time);
    time_series_total = cat(1,time_series_total,time_series);
    
end
save('u1_square.mat','u1_square_time_total');
save('v1_square.mat','v1_square_time_total');
save('u2_square.mat','u2_square_time_total');
save('v2_square.mat','v2_square_time_total');
save('u_barotropic_square.mat','u_barotropic_square_time_total');
save('u_baroclinic_square.mat','u_baroclinic_square_time_total');
save('v_barotropic_square.mat','v_barotropic_square_time_total');
save('v_baroclinic_square.mat','v_baroclinic_square_time_total');
save('tau_baroclinic_square.mat','tau_baroclinic_square_time_total');
save('psi_barotropic_square.mat','psi_barotropic_square_time_total');