clc;clear all;close all;

%========== Model output ===========%
filename="/data3/yiwen/cesm2/JGR_ocean/Model_output/G_JRA55_test1.pop.h.0031-04.nc"
LAT=ncread(filename,'TLAT');lat=LAT(1,:);     
LON=ncread(filename,'TLONG');lon=LON(:,1);    
depth=ncread(filename,'z_t');                 %unit=cm 
temp=ncread(filename,'TEMP');mask=temp;
mask(~isnan(mask))=1; %1=ocean & nan=land


pathi='....~/scratch/cesm2/archive/G_JRA55_test1/ocn/hist/' % pop2 model archive path

file_template = 'G_JRA55_test1.pop.h.%04d-%02d.nc';
nyra_start = 184;nyra_end = 366;

nx = 320;ny = 384;
n_months = (nyra_end - nyra_start + 1) * 12;
SHF = zeros(nx, ny, n_months);
SFWF = zeros(nx, ny, n_months);
SALT = zeros(nx, ny, n_months);
TEMP = zeros(nx, ny, n_months);
for iyr = nyra_start:nyra_end
    for im = 1:12
    count = (iyr - nyra_start) * 12 + im;
    current_file = sprintf(file_template, iyr, im);
    full_path = [pathi  current_file];
    shf = ncread(full_path, 'SHF');    %negtive(-) == upward
    evap = ncread(full_path, 'EVAP_F');%negtive(-)
    prec = ncread(full_path, 'PREC_F');%postive(+)
    salt = ncread(full_path, 'SALT');
    temp = ncread(full_path, 'TEMP');

    SHF(:, :, count) = shf.*mask(:,:,1);
    SFWF(:, :, count) = (evap+prec).*mask(:,:,1);
    SALT(:, :, count) = salt(:,:,1).*mask(:,:,1);        
    TEMP(:, :, count) = temp(:,:,1).*mask(:,:,1);
    end
end
%
%============= VOLUME AVG ===============%
area=ncread(filename,'TAREA'); %choose 1st-layer(cm^2)
loc1=find(LON<=360-37.5 & LON>=360-55 & LAT<=57.5 & LAT>=47.5); %study region_N
%------------------------------------------------------%
for k=1:60
areamask{k}=area.*mask(:,:,k);    %2D
areamask_IN{k}=areamask{k}(loc1); %1D
sumarea(k)=squeeze(nansum(areamask_IN{k}));
end

for count=1:size(SHF,3)
shf_k{count} =SHF(:,:,count);
sfwf_k{count} =SFWF(:,:,count);
temp_k{count} =TEMP(:,:,count);
salt_k{count} =SALT(:,:,count);

shf_IN{count}=shf_k{count}(loc1);
sfwf_IN{count}=sfwf_k{count}(loc1);
temp_IN{count}=temp_k{count}(loc1);
salt_IN{count}=salt_k{count}(loc1);

SurfQDOT(count)=nansum(nansum(shf_IN{count}.*areamask_IN{1}))/sumarea(1);
SurfNWater(count)=nansum(nansum(sfwf_IN{count}.*areamask_IN{1}))/sumarea(1);
SurfSUMS(count)=nansum(nansum(salt_IN{count}.*areamask_IN{1}))/sumarea(1);
SurfSUMT(count)=nansum(nansum(temp_IN{count}.*areamask_IN{1}))/sumarea(1);
end

%-------------------%
rho0=1025;Cp=4000;
%cal T&H
for count=1:size(SHF,3)
%(1)
alpha(count)=-1/rho0.*(0.049762-0.0030063*SurfSUMS(count)+...
                (2*SurfSUMT(count)*(-0.0075911+0.000037*SurfSUMS(count)))...
                +3*SurfSUMT(count)^2*0.000035).*2*SurfSUMT(count); %chain rule

beta(count)=1/rho0.*(0.8056-0.0030063*SurfSUMT(count)+0.000037*SurfSUMT(count)^2);
%(2) || %E-P(kg*m-2*s-1)
thermal(count)= -(SurfQDOT(count)/Cp).*alpha(count);
haline(count)=((SurfSUMS(count)/(1-SurfSUMS(count)))*SurfNWater(count).*(rho0)).*beta(count);

%(3) sum :: D=thermal+haline;
D(count)=thermal(count)+haline(count);
end

%(1)----- annual mean ------%
for iyr = 1:length(D)/12
  DD(iyr)= nanmean(D((iyr-1)*12+1:(iyr-1)*12+12));
  DT(iyr)= nanmean(thermal((iyr-1)*12+1:(iyr-1)*12+12));
  DH(iyr)= nanmean(haline((iyr-1)*12+1:(iyr-1)*12+12));
end

% minus mean state %
DD_M=DD-mean(DD);
DT_M=DT-mean(DT);
DH_M=DH-mean(DH);

%(2)----- Winter time[December–March] ------%
D_shift=circshift(D,-1);
T_shift=circshift(thermal,-1);
H_shift=circshift(haline,-1);

ST=1;END=2196;
% month12=ST+11:12:END;
month01=ST:12:END;
month02=ST+1:12:END;
month03=ST+2:12:END;
month04=ST+3:12:END;
 
DD_winter=D_shift(month01);
DT_winter=T_shift(month01);
DH_winter=H_shift(month01);


%-- 5_yr moving mean --%
%DD5=movmean(DD_M,5,"omitnan");
%DT5=movmean(DT_M,5,"omitnan");
%DH5=movmean(DH_M,5,"omitnan");

%======================================================
% save to mat
%======================================================
patho='/data3/yiwen/cesm2/JGR_ocean/Model_output/CTRL/'
fn1 = [patho 'Flux_4_6th_1D.mat'];
save(fn1,'DD_M','DT_M','DH_M','DD','DT','DH',...
         'DD_winter','DT_winter','DH_winter','-v7.3')

