clc
clear all
close all

%-----------------------------------------
% EXPERIMENT DEFINITION
%-----------------------------------------

storm_name = 'Idai';
save_path = ['/discover/nobackup/jkolassa/SMAP_TC_project/Data_analysis/' storm_name '/'];

exp_dir = '/discover/nobackup/projects/smaptc/jkolassa/archive/Idai/'; % '/discover/nobackup/projects/smaptc/jkolassa/archive/Idai/';
exp_name = 'Idai_LDAS_DA_4D_C360';
exp_resolution = 'CF0360x6C_GLOBAL';
da_interval = 3;   % DA interval in hours

%lat_range = [12 35]; 
lat_range = [-24 -10];
%lon_range = [52 97]; 
lon_range = [24 52];

start_year = 2019;
start_month = 2;
start_day = 4;

%start_year = 2021;
%start_month = 9;
%start_day = 1;

end_year = 2019;
end_month = 3;
end_day = 22;

%end_year = 2021;
%end_month = 10;
%end_day = 4;

%---------------------------------------------
% DATA READ PREPARATION
%---------------------------------------------

leapYears = 1980:4:2022;
monthLength_nly = [31 28 31 30 31 30 31 31 30 31 30 31];
monthLength_ly = [31 29 31 30 31 30 31 31 30 31 30 31];
mat_index = 0;

%---------------------------------------------
% DATA READ
%---------------------------------------------

% read track data

track_file = ['/discover/nobackup/jkolassa/SMAP_TC_project/track_files/' storm_name '_track.ascii'];
track_data = dlmread(track_file);
track_datenum = datenum(start_year, track_data(:,1), track_data(:,2), track_data(:,3), 0, 0);

disp('read data')
% loop over years
for y = start_year:end_year

    y_str = num2str(y,'%04i');

    % define month length based on leap year status
    if sum(y==leapYears)>0
        monthLength = monthLength_ly;
    else
        monthLength = monthLength_nly;
    end

    % loop over months
    for m = start_month:end_month

        m_str = num2str(m,'%02i');

        % define start and end day for current month
        day_first = 1;
        day_last = monthLength(m);

        if (m == start_month)
           day_first = start_day;
        end

        if (m == end_month)
           day_last = end_day;
        end

        % loop over days
        for d = day_first:day_last

            mat_index = mat_index + 1;

            d_str = num2str(d,'%02i');
            t_index = 0;

            % loop over time steps
            for t = 1:3:22
                t_index = t_index + 1;

                t_str = num2str(t,'%02i');

                % define increment file name

                ana_file = [exp_dir exp_name '/output/' exp_resolution '/cat/ens_avg/Y' y_str '/M' m_str '/' exp_name ... 
                            '.inst3_2d_lndfcstana_Nx.' y_str m_str d_str '_' t_str '30z.nc4'];

                % on first time step read coordinate information and initialize matrices

                if (mat_index == 1) && (t == 1)

                   lat_vec = double(ncread(ana_file,'lats'));
                   lat_vec = lat_vec(:);
                   lon_vec = double(ncread(ana_file,'lons'));
                   lon_vec = lon_vec(:);

                   datenum_vec = NaN*ones(sum(monthLength(start_month:end_month)),1);

                   sfmc_ttot24_mat = NaN*ones(length(lon_vec),sum(monthLength(start_month:end_month)));
                   rzmc_ttot24_incr_mat = NaN*ones(length(lon_vec),sum(monthLength(start_month:end_month)));

                   sfmc_3hour_incr_mat = NaN*ones(length(lon_vec),24/da_interval);
                   rzmc_3hour_incr_mat = NaN*ones(length(lon_vec),24/da_interval);
                end

                sfmc_ana_mat = double(ncread(ana_file,'SFMC_ANA'));
                sfmc_ana_mat = sfmc_ana_mat(:);
                sfmc_fcst_mat = double(ncread(ana_file,'SFMC_FCST'));
                sfmc_fcst_mat = sfmc_fcst_mat(:);
                rzmc_ana_mat = double(ncread(ana_file,'RZMC_ANA'));
                rzmc_ana_mat = rzmc_ana_mat(:);
                rzmc_fcst_mat = double(ncread(ana_file,'RZMC_FCST'));
                rzmc_fcst_mat = rzmc_fcst_mat(:);

                sfmc_3hour_incr_mat(:,t_index) = sfmc_ana_mat - sfmc_fcst_mat;
                rzmc_3hour_incr_mat(:,t_index) = rzmc_ana_mat - rzmc_fcst_mat;


            end % t

            % log date

            datenum_vec(mat_index,1) = datenum(y,m,d);

            % compute daily total increments

            sfmc_ttot24_incr_mat(:,mat_index) = nansum(sfmc_3hour_incr_mat,2);
            rzmc_ttot24_incr_mat(:,mat_index) = nansum(rzmc_3hour_incr_mat,2);

            % reset hourly matrices

            sfmc_3hour_incr_mat = NaN*ones(length(lon_vec),24/da_interval);
            rzmc_3hour_incr_mat = NaN*ones(length(lon_vec),24/da_interval);

         end % d
      end %m
end % y

% remove excess allocated slots

datenum_vec((mat_index+1):end) = [];
sfmc_ttot24_incr_mat(:,(mat_index+1):end) = [];
rzmc_ttot24_incr_mat(:,(mat_index+1):end) = [];

sfmc_ttot24_incr_mat(sfmc_ttot24_incr_mat==0) = NaN;
rzmc_ttot24_incr_mat(rzmc_ttot24_incr_mat==0) = NaN;

% define plot domain
load coastlines;
[latcells, loncells] = polysplit(coastlat, coastlon);
latlim = lat_range;
lonlim = lon_range;

lat_index = (lat_vec>=lat_range(1)) & (lat_vec<=lat_range(2));
lon_index = (lon_vec>=lon_range(1)) & (lon_vec<=lon_range(2));
data_index = lat_index & lon_index;

% define colorbars and limits
disp('define colorbars')
tmp_mat = sfmc_ttot24_incr_mat(data_index,:);
tmp_mat(tmp_mat==0) = NaN;

s1 = quantile(tmp_mat(:),0.99);
s2 = quantile(tmp_mat(:),0.01);
caxis_min_sfmc = -1*max(abs(s1),abs(s2));
caxis_min_sfmc = round(caxis_min_sfmc,2);
caxis_max_sfmc = max(abs(s1),abs(s2));
caxis_max_sfmc = round(caxis_max_sfmc,2);

c_length = caxis_max_sfmc/0.01;
c_length = c_length*2;

cmap_mat = cmocean('balance',c_length);
cmap_mat = flipud(cmap_mat);
cmap_mat(length(cmap_mat)/2:length(cmap_mat)/2+1,:) = ones(2,3);
cmap_sfmc = cmap_mat;

tmp_mat = rzmc_ttot24_incr_mat(data_index,:);
tmp_mat(tmp_mat==0) = NaN;

s1 = quantile(tmp_mat(:),0.99);
s2 = quantile(tmp_mat(:),0.01);
caxis_min_rzmc = -1*max(abs(s1),abs(s2));
caxis_min_rzmc = round(caxis_min_rzmc,2);
caxis_max_rzmc = max(abs(s1),abs(s2));
caxis_max_rzmc = round(caxis_max_rzmc,2);

c_length = caxis_max_rzmc/0.001;
c_length = c_length*2;

cmap_mat = cmocean('balance',c_length);
cmap_mat = flipud(cmap_mat);
cmap_mat(length(cmap_mat)/2:length(cmap_mat)/2+1,:) = ones(2,3);
cmap_rzmc = cmap_mat;

for i = 1:length(datenum_vec)

    datenum_dist = sqrt((track_datenum-datenum_vec(i)).^2);
    date_index = find(datenum_dist==min(datenum_dist));

    figure
    ax = axesm('miller','MapLatLimit',lat_range,'MapLonLimit',lon_range,...
               'FLatLimit',lat_range,'FLonLimit',lon_range);
    framem
    gridm on;
    setm(ax,'MLineLocation',5,'PLineLocation',5,'MeridianLabel','on',...
         'ParallelLabel','on','MLabelLocation',5,'PLabelLocation',5,...
          'MLabelParallel','south','FontSize',12)
    geoshow('landareas.shp', 'FaceColor', [0.7 0.7 0.7])
   % pcolorm(lat_mat,lon_mat,exp_srfexc_ttot24_incr_plot_mat(:,:,i))
    scatterm(ax,lat_vec,lon_vec,10,sfmc_ttot24_incr_mat(:,i),'s','filled')
    caxis([caxis_min_sfmc caxis_max_sfmc]);
    set(gcf,'Colormap',cmap_sfmc);
    plotm(coastlat, coastlon,'k')
    hold on
    %plotm(track_data(1:date_index,5),track_data(1:date_index,4),'-m','LineWidth',2)
    tmp_lon = track_data(1:date_index,5);
    tmp_lat = track_data(1:date_index,4);
    scatterm(ax,tmp_lat,tmp_lon,10,"filled","m")
    hold off
    set(gca,'XLim',lonlim);
    set(gca,'YLim',latlim);
    tightmap
    title(['SFMC increment [m^{3} m^{-3}] ' datestr(datenum_vec(i),'yyyy-mm-dd')])
    h = colorbar;
    set(h,'Location','eastoutside','Orientation','vertical','FontSize',14);
    set(gcf,'paperunits','centimeters')
    set(gcf,'paperposition',[1 1 17 15])
    saveas(gcf,[save_path storm_name '_incr_SFMC_' datestr(datenum_vec(i),'yyyy-mm-dd')],'png')

    figure
    ax = axesm('miller','MapLatLimit',lat_range,'MapLonLimit',lon_range,...
               'FLatLimit',lat_range,'FLonLimit',lon_range);
    framem
    gridm on;
    setm(ax,'MLineLocation',5,'PLineLocation',5,'MeridianLabel','on',...
         'ParallelLabel','on','MLabelLocation',5,'PLabelLocation',5,...
          'MLabelParallel','south','FontSize',12)
    geoshow('landareas.shp', 'FaceColor', [0.7 0.7 0.7])
   % pcolorm(lat_mat,lon_mat,exp_srfexc_ttot24_incr_plot_mat(:,:,i))
    scatterm(ax,lat_vec,lon_vec,10,rzmc_ttot24_incr_mat(:,i),'s','filled')
    caxis([caxis_min_rzmc caxis_max_rzmc]);
    set(gcf,'Colormap',cmap_rzmc);
    plotm(coastlat, coastlon,'k')
    hold on
    %plotm(track_data(1:date_index,5),track_data(1:date_index,4),'-m','LineWidth',2)
    tmp_lon = track_data(1:date_index,5);
    tmp_lat = track_data(1:date_index,4);
    scatterm(ax,tmp_lat,tmp_lon,10,"filled","m")
    hold off
    set(gca,'XLim',lonlim);
    set(gca,'YLim',latlim);
    tightmap
    title(['RZMC increment [m^{3} m^{-3}] ' datestr(datenum_vec(i),'yyyy-mm-dd')])
    h = colorbar;
    set(h,'Location','eastoutside','Orientation','vertical','FontSize',14);
    set(gcf,'paperunits','centimeters')
    set(gcf,'paperposition',[1 1 17 15])
    saveas(gcf,[save_path storm_name '_incr_RZMC_' datestr(datenum_vec(i),'yyyy-mm-dd')],'png')

end
