% ------------------------------------------------------------------------- %
%  Read MITgcm NetCDF Output Fields 
% Written by: Taylor Shropshire
% Date: 01/06/21
% Description: Reads in each individual day across all tiles 
% then organizes tiles into fields.
% ------------------------------------------------------------------------- %
clear;clc;close all

% Paths 
path1='./fields/';

% Get all tile directory names
dir_names=dir([path1,'/*tile*']); 

% Model Start Time 
msd=datenum(2005,01,01);

% File Read Flags 
ptr_flag=0;		% read 30-day 3D state variable files
diag3D_flag=1;		% read 30-day 3D diagnostics
diag2D_flag=0;		% read daily 2D diagnostics 

% ============================================================================ %
% ======== Get grid coordinates and spatial dimension of output files ======== % 
% ============================================================================ %
display('Reading Model Grid');

% Get depth layers 
grid_file=dir([path1,dir_names(1).name,'/*grid*']);
nc_file=[path1,dir_names(1).name,'/',grid_file(1).name];
depth_uv=ncread(nc_file,'RC',1,Inf)*-1;
depth_w=ncread(nc_file,'RF',1,Inf)*-1;
depth_dz=diff(depth_w); 

% ------- Tile Loop ------------- % 
grd_field=[];
for dd=1:length(dir_names)
grid_file=dir([path1,dir_names(dd).name,'/*grid*']);
nc_file=[path1,dir_names(dd).name,'/',grid_file(1).name];

% Get variable names
temp=ncinfo(nc_file); var_names_all={temp.Variables.Name};
var_names={'XC','YC','Depth','rA'};     % manually

% Extract variables
grd=[];
icount_start=[1,1];    icount=[Inf,Inf];
for nn=1:length(var_names);
grd(:,:,nn)=ncread(nc_file,cell2mat(var_names(nn)),icount_start,icount);
end

% Rotate variables
grd2=[];
for pp=1:size(grd,3)
grd2(:,:,pp)=grd(:,:,pp)';
end
grd=grd2; clear grd2;

% Save Tile 
grd_field(:,:,:,dd)=grd;
end     % Directory Loop 
% -------------------------------- % 

% Define domain lat lon limits
xc_temp=squeeze(grd_field(:,:,1,:));            yc_temp=squeeze(grd_field(:,:,2,:));
temp1=reshape(xc_temp,[1,numel(xc_temp)]);      temp2=reshape(yc_temp,[1,numel(yc_temp)]);
dom_x1=min(temp1); dom_x2=max(temp1);           dom_y1=min(temp2); dom_y2=max(temp2);

% Define all tile lat lon limits
xlim_all=[]; ylim_all=[];
xlim_all(:,1)=squeeze(xc_temp(1,1,:));   % bottom left x coordinate for all tiles 
xlim_all(:,2)=squeeze(xc_temp(1,end,:)); % bottom right x corrdinate for all tiles
ylim_all(:,1)=squeeze(yc_temp(1,1,:));   % bottom left y coordinate for all tiles
ylim_all(:,2)=squeeze(yc_temp(end,1,:)); % top left y coorodinate for all tiles

% Define unique tile lat lon limits 
xlim=[]; ylim=[];
xlim(:,1)=sort(unique(xlim_all(:,1)));          xlim(:,2)=sort(unique(xlim_all(:,2)));
ylim(:,1)=sort(unique(ylim_all(:,1)));          ylim(:,2)=sort(unique(ylim_all(:,2)));

% Dimensions
txd=size(grd_field,2); tyd=size(grd_field,1); tzd=length(depth_uv);           % define tile dimension
num_tx=size(xlim,1); num_ty=size(ylim,1);                       % number of tiles in x and y

% ------- Organize into single array -------- % 
grd_field2=[];
y1=1;           y2=tyd;
for yy=1:num_ty
x1=1;           x2=txd;
for xx=1:num_tx
ind=find(xlim_all(:,1)==xlim(xx,1) & xlim_all(:,2)==xlim(xx,2)...
& ylim_all(:,1)==ylim(yy,1) & ylim_all(:,2)==ylim(yy,2));

grd_field2(y1:y2,x1:x2,:)=grd_field(:,:,:,ind);

x1=x2+1;        x2=x2+txd;
end
y1=y2+1;        y2=y2+tyd;
end
grd_field=grd_field2; clear grd_field2;
% ------- Organize into single array -------- % 

% ======================================================= % 
% ======================================================= % 
% ======================================================= % 

% Define variables 
ind=find(ismember(var_names, 'XC'));		xc=grd_field(:,:,ind);
ind=find(ismember(var_names, 'YC'));    	yc=grd_field(:,:,ind);
ind=find(ismember(var_names, 'Depth'));		bathy=grd_field(:,:,ind);
ind=find(ismember(var_names, 'rA'));    	rA=grd_field(:,:,ind);

% Set land to NaN
bathy(find(bathy==0))=NaN;

% ------ Get mask ------ % 
for dd=1:length(dir_names)
grid_file=dir([path1,dir_names(dd).name,'/*grid*']);
nc_file=[path1,dir_names(dd).name,'/',grid_file(1).name];
cfrac=ncread(nc_file,'HFacC',[1,1,1],[Inf,Inf,Inf])*-1;

cfrac2=[];
for zz=1:size(cfrac,3)
cfrac2(:,:,zz)=cfrac(:,:,zz)';
end 
cfrac=cfrac2; clear cfrac2;

cfrac_all(:,:,:,dd)=cfrac;
end 

cfrac_all2=[];
y1=1;           y2=tyd;
for yy=1:num_ty
x1=1;           x2=txd;
for xx=1:num_tx
ind=find(xlim_all(:,1)==xlim(xx,1) & xlim_all(:,2)==xlim(xx,2)...
& ylim_all(:,1)==ylim(yy,1) & ylim_all(:,2)==ylim(yy,2));

cfrac_all2(y1:y2,x1:x2,:)=cfrac_all(:,:,:,ind);

x1=x2+1;        x2=x2+txd;
end
y1=y2+1;        y2=y2+tyd;
end
cfrac=cfrac_all2; clear cfrac_all2;

mask3d=cfrac; mask3d(find(cfrac==0))=NaN; mask3d(find(~isnan(mask3d)))=1;
% ------------------ % 

% Remove boundary region
ind=find(yc<=21 & xc>=-89);
for zz=1:size(mask3d,3)
temp=mask3d(:,:,zz); temp(ind)=NaN;
mask3d(:,:,zz)=temp;
end
mask2d=mask3d(:,:,1);

% Define Model Field Dimensions
xdim=size(bathy,2); ydim=size(bathy,1); zdim=length(depth_uv);

% Calculate Cell Volume 
vol=[];
for zz=1:size(cfrac,3)
vol(:,:,zz)=rA*depth_dz(zz).*cfrac(:,:,zz);
end

display('Done Reading Model Grid'); display(''); %pause 
% ============================================================================ %
% ============================================================================ %
% ============================================================================ %

if ptr_flag==1
% ###################################################################### %
% ######################## READ 3D Monthly Ouput ####################### %
% ###################################################################### %
NO_all=[]; NH_all=[]; SI_all=[]; DON_all=[]; PON_all=[]; sp_all=[]; lp_all=[]; 
sz_all=[]; lz_all=[]; pz_all=[]; ptracers_all=[];

% Time Dimension
file_head='/*ptracers*';
%file_head='/*ptr_tave*';

files_PT=dir([path1,dir_names(dd).name,file_head]);
nc_file=[path1,dir_names(dd).name,'/',files_PT(1).name];
rec_PT=ncread(nc_file,'T',1,Inf);
display(['Reading "',file_head,'" NetCDF Files']);
display(['Total Records = ', num2str(length(rec_PT))]);

% ************** Time Loop ************** %  
count=0;
for tt=1:length(rec_PT)
display(['Record = ',num2str(tt), '/',num2str(length(rec_PT))]);tic

% ------- Tile Loop ------------- % 
ptracers_field=[];
for dd=1:length(dir_names)
files_PT=dir([path1,dir_names(dd).name,'/*ptracers*']);
nc_file=[path1,dir_names(dd).name,'/',files_PT(1).name];

% Get variable names
temp=ncinfo(nc_file); var_names={temp.Variables.Name};
var_names(1:5)=[];      % dimensions

% Extract varibles
ptracers=[];
ncid=netcdf.open(nc_file,'NOWRITE');
for nn=1:length(var_names)
varid=netcdf.inqVarID(ncid,cell2mat(var_names(nn)));
temp=netcdf.getVar(ncid,varid,[0,0,0,tt-1],[txd,tyd,tzd,1]);
ptracers(:,:,:,nn)=double(temp);
end
netcdf.close(ncid);

% Extract varibles (alternative code)
%ptracers=[];
%icount_start=[1,1,1,tt];    icount=[txd,tyd,tzd,1];
%for nn=1:length(var_names);
%ptracers(:,:,:,nn)=ncread(nc_file,cell2mat(var_names(nn)),icount_start,icount);
%end
% note: slower than extracting code above

% Rotate variables
ptracers2=[];
for pp=1:size(ptracers,4)
for zz=1:size(ptracers,3)
ptracers2(:,:,zz,pp)=ptracers(:,:,zz,pp)';
end; end
ptracers=ptracers2; clear ptracers2;

% Save Tile 
ptracers_field(:,:,:,:,dd)=ptracers;
end     % Directory Loop 
% -------------------------------------- %

% ------- Organize into single array -------- % 
ptracers_field2=[];
y1=1;           y2=tyd;
for yy=1:num_ty
x1=1;           x2=txd;
for xx=1:num_tx
ind=find(xlim_all(:,1)==xlim(xx,1) & xlim_all(:,2)==xlim(xx,2)...
& ylim_all(:,1)==ylim(yy,1) & ylim_all(:,2)==ylim(yy,2));

ptracers_field2(y1:y2,x1:x2,:,:)=ptracers_field(:,:,:,:,ind);

x1=x2+1;        x2=x2+txd;
end
y1=y2+1;        y2=y2+tyd;
end
ptracers_field=ptracers_field2; clear ptracers_field2;
% ------------------------------------- % 

% --- Check Conservation and Apply Mask--- %
for pp=1:size(ptracers_field,4)
ptracers_field(:,:,:,pp)=ptracers_field(:,:,:,pp).*mask3d;
end 
temp_ptracers=ptracers_field;

ind=[1,2,4,5,7,8,9,10,11];
temp=temp_ptracers(:,:,:,ind);
tot_n(1,tt)=sum(reshape(temp,[1,numel(temp)])); 

r_SI_N=1.0;	% SI/N model parameter

ind=[3,6,8];
temp=temp_ptracers(:,:,ind)*r_SI_N;
tot_s(1,tt)=sum(reshape(temp,[1,numel(temp)]));
% --------------------------- % 

% Average Fields 
if count==0
ptracers_avg=ptracers_field;
count=count+1;
else
ptracers_avg=ptracers_avg+ptracers_field;
count=count+1;
end

% Save specific fields 
ind=find(ismember(var_names, 'NO'));	NO_all(:,:,:,tt)=ptracers_field(:,:,:,ind);
%ptracers_all(:,:,:,:,tt)=ptracers_field;
% note: can run into system memory issues with long simulations

toc
end % Time Loop
% ********************************************** % 

% Get Average 
ptracers_avg=ptracers_avg/count;
NO_map=ptracers_avg(:,:,:,1);		NH_map=ptracers_avg(:,:,:,2);
SI_map=ptracers_avg(:,:,:,3);		DON_map=ptracers_avg(:,:,:,4);
PON_map=ptracers_avg(:,:,:,5);		OP_map=ptracers_avg(:,:,:,6);
sp_map=ptracers_avg(:,:,:,7);		lp_map=ptracers_avg(:,:,:,8);
sz_map=ptracers_avg(:,:,:,9);		lz_map=ptracers_avg(:,:,:,10);
pz_map=ptracers_avg(:,:,:,11);

display('Done Reading "ptracers_fields"'); display('');% pause
% ####################################################################### % 
% ####################################################################### % 
% #######################################################################
end % if ptr_flag==1

if diag3D_flag==1
% ###################################################################### %
% ######################## READ 3D Diagnostic  ####################### %
% ###################################################################### %
chl_all=[]; tmp_all=[]; par_all=[];
diag3D_all=[]; 

% Time Dimension
files_bio3D=dir([path1,dir_names(dd).name,'/*3D_bio*']);
nc_file=[path1,dir_names(dd).name,'/',files_bio3D(1).name];
rec_bio3D=ncread(nc_file,'T',1,Inf);
display(['Reading "3D_bio_diagnostics" NetCDF Files']);
display(['Total Records = ', num2str(length(rec_bio3D))]);

% ********** Time Record Loop *********** % 
count=0;
for tt=1:length(rec_bio3D)
display(['Record = ',num2str(tt), '/',num2str(length(rec_bio3D))]);tic

% ------- Tile Loop ------------- % 
diag3D_field=[];
for dd=1:length(dir_names)
files_bio3D=dir([path1,dir_names(dd).name,'/*3D_bio*']);
nc_file=[path1,dir_names(dd).name,'/',files_bio3D(1).name];

% Get variable names
temp=ncinfo(nc_file); var_names={temp.Variables.Name};
var_names(1:5)=[];      % dimensions

% Extract varibles
diag3D=[];
ncid=netcdf.open(nc_file,'NOWRITE');
for nn=1:length(var_names)
varid=netcdf.inqVarID(ncid,cell2mat(var_names(nn)));
temp=netcdf.getVar(ncid,varid,[0,0,0,tt-1],[txd,tyd,tzd,1]);
diag3D(:,:,:,nn)=double(temp);
end
netcdf.close(ncid);

% Rotate variables
diag3D_2=[];
for pp=1:size(diag3D,4)
for zz=1:size(diag3D,3)
diag3D_2(:,:,zz,pp)=diag3D(:,:,zz,pp)';
end; end
diag3D=diag3D_2; clear diag3D_2;

% Save Tile 
diag3D_field(:,:,:,:,dd)=diag3D;
end     % Directory Loop 
% ---------------------------- %

% ---------- Organize into single array ---------- % 
diag3D_field2=[];
y1=1; y2=tyd;
for yy=1:num_ty
x1=1; x2=txd;
for xx=1:num_tx
ind=find(xlim_all(:,1)==xlim(xx,1) & xlim_all(:,2)==xlim(xx,2)...
& ylim_all(:,1)==ylim(yy,1) & ylim_all(:,2)==ylim(yy,2));

diag3D_field2(y1:y2,x1:x2,:,:)=diag3D_field(:,:,:,:,ind);

x1=x2+1; x2=x2+txd;
end
y1=y2+1; y2=y2+tyd;
end
diag3D_field=diag3D_field2; clear diag3D_field2;
% ------------------------------------------------- % 

% Apply Mask 
for nn=1:size(diag3D_field,4);
diag3D_field(:,:,:,nn)=diag3D_field(:,:,:,nn).*mask3d;
end 

% Average Fields 
if count==0
diag3D_avg=diag3D_field;
count=count+1;
else
diag3D_avg=diag3D_avg+diag3D_field;
count=count+1;
end

% Save fields 
ind=find(ismember(var_names, 'chl_3Dfi'));    chl_all(:,:,:,tt)=diag3D_field(:,:,:,ind);
%diag3D_all(:,:,:,:,tt)=diag3D_field;

toc
end % Time Loop

% Get Average 
diag3D_avg=diag3D_avg/count;
	
display('Done Reading "diag3D_fields"'); display(''); %pause
% ####################################################################### % 
% ####################################################################### % 
% #######################################################################
end % if diag3D_flag==1

if diag2D_flag==1
% ###################################################################### %
% ######################## READ 2D Diagnostic  ####################### %
% ###################################################################### %
avgChl_all=[]; surfChl_all=[]; mld_all=[]; tmp_all=[]; 

% Time Dimension
files_bio2D=dir([path1,dir_names(dd).name,'/*2D_bio*']);
nc_file=[path1,dir_names(dd).name,'/',files_bio2D(1).name];
rec_bio2D=ncread(nc_file,'T',1,Inf);
display(['Reading "2D_bio_diagnostics" NetCDF Files']);
display(['Total Records = ', num2str(length(rec_bio2D))]);

% >>>>>>>> Time Record Loop <<<<<<<< % 
count=0;
for tt=1:length(rec_bio2D)
display(['Record = ',num2str(tt), '/',num2str(length(rec_bio2D))]);tic

% ------- Tile Loop ------------- % 
diag2D_field=[];
for dd=1:length(dir_names)
files_bio2D=dir([path1,dir_names(dd).name,'/*2D_bio*']);
nc_file=[path1,dir_names(dd).name,'/',files_bio2D(1).name];

% Get variable names
temp=ncinfo(nc_file); var_names={temp.Variables.Name};
var_names(1:5)=[];      % dimensions
%var_names=var_names([9,10,11);		% read files quicker for specific var

% Extract varibles
diag2D=[];
ncid=netcdf.open(nc_file,'NOWRITE');
for nn=1:length(var_names)
varid=netcdf.inqVarID(ncid,cell2mat(var_names(nn)));
temp=netcdf.getVar(ncid,varid,[0,0,0,tt-1],[txd,tyd,1,1]);
diag2D(:,:,nn)=double(temp);
end 
netcdf.close(ncid);

% Rotate variables
diag2D_2=[];
for pp=1:size(diag2D,3)
diag2D_2(:,:,pp)=diag2D(:,:,pp)';
end
diag2D=diag2D_2; clear diag2D_2;

% Save Tile 
diag2D_field(:,:,:,dd)=diag2D;
end     % Directory Loop 
% ------- Tile Loop ------------- % 

% --------- Organize into single array ---------- % 
diag2D_field2=[];
y1=1; y2=tyd;
for yy=1:num_ty
x1=1; x2=txd;
for xx=1:num_tx
ind=find(xlim_all(:,1)==xlim(xx,1) & xlim_all(:,2)==xlim(xx,2)...
& ylim_all(:,1)==ylim(yy,1) & ylim_all(:,2)==ylim(yy,2));

diag2D_field2(y1:y2,x1:x2,:)=diag2D_field(:,:,:,ind);

x1=x2+1; x2=x2+txd;
end
y1=y2+1; y2=y2+tyd;
end
diag2D_field=diag2D_field2; clear diag2D_field2;
% ------------------------------------------------- % 

% Apply Mask 
for nn=1:size(diag2D_field,3);
diag2D_field(:,:,nn)=diag2D_field(:,:,nn).*mask2d;
end

% Average Fields 
if count==0
diag2D_avg=diag2D_field;
count=count+1;
else
diag2D_avg=diag2D_avg+diag2D_field;
count=count+1;
end

% Save fields 
ind=find(ismember(var_names, 'chl_2Dfi'));    avgChl_all(:,:,tt)=diag2D_field(:,:,ind);
ind=find(ismember(var_names, 'surf_chl'));    surfChl_all(:,:,tt)=diag2D_field(:,:,ind);
ind=find(ismember(var_names, 'mld_2Dfi'));    mld_all(:,:,tt)=diag2D_field(:,:,ind);
%diag2D_all(:,:,:,tt)=diag2D_field;		% save all fields

toc
end % Time Loop

% Get Average 
diag2D_avg=diag2D_avg/count;

display('Done Reading "diag2D_fields"'); display('');  %pause
% ####################################################################### % 
% ####################################################################### % 
% #######################################################################
end % if diag2D_flag==1



