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/';
exp_name = 'Idai_LDAS_DA_4D_C360';
exp_resolution = 'CF0360x6C_GLOBAL';
da_interval = 3;   % DA interval in hours

lat_range = [-24 -10];
lon_range = [24 52];

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

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

%---------------------------------------------
% 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

                exp_incr_file = [exp_dir exp_name '/output/' exp_resolution '/cat/ens_avg/Y' y_str '/M' m_str '/' exp_name ... 
                                '.catch_progn_incr.' 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(exp_incr_file,'lat'));
                   lon_vec = double(ncread(exp_incr_file,'lon'));

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

                   exp_srfexc_ttot24_incr_mat = NaN*ones(length(lon_vec),sum(monthLength(start_month:end_month)));
                   exp_rzexc_ttot24_incr_mat = NaN*ones(length(lon_vec),sum(monthLength(start_month:end_month)));

                   exp_srfexc_3hour_incr_mat = NaN*ones(length(lon_vec),24/da_interval);
                   exp_rzexc_3hour_incr_mat = NaN*ones(length(lon_vec),24/da_interval);
                end

                exp_srfexc_3hour_incr_mat(:,t_index) = double(ncread(exp_incr_file,'SRFEXC_INCR'));
                exp_rzexc_3hour_incr_mat(:,t_index) = double(ncread(exp_incr_file,'RZEXC_INCR'));
 
            end % t

            % log date

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

            % compute daily total increments

            exp_srfexc_ttot24_incr_mat(:,mat_index) = nansum(exp_srfexc_3hour_incr_mat,2);
            exp_rzexc_ttot24_incr_mat(:,mat_index) = nansum(exp_rzexc_3hour_incr_mat,2);

            % reset hourly matrices

            exp_srfexc_3hour_incr_mat = NaN*ones(length(lon_vec),24/da_interval);
            exp_rzexc_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) = [];
exp_srfexc_ttot24_incr_mat(:,(mat_index+1):end) = [];
exp_rzexc_ttot24_incr_mat(:,(mat_index+1):end) = [];

% create animation and plot of differences

%% map vectors to matrices
%
%unique_lat  = unique(lat_vec);
%unique_lon = unique(lon_vec);
%
%[lon_mat,lat_mat] = meshgrid(unique_lon,unique_lat);
%
%exp_srfexc_ttot24_incr_plot_mat = NaN*ones(size(lon_mat,1), size(lon_mat,2), size(exp_srfexc_ttot24_incr_mat,2)); 
%
%for i = 1:length(lon_vec)
%
%    lat_index = unique_lat == lat_vec(i);
%    lon_index = unique_lon == lon_vec(i);
%
%    exp_srfexc_ttot24_incr_plot_mat(lat_index,lon_index,:) = exp_srfexc_ttot24_incr_mat(i,:);
%
%end

% 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 = exp_srfexc_ttot24_incr_mat(data_index,:);
tmp_mat(tmp_mat==0) = NaN;

s1 = quantile(tmp_mat(:),0.99);
s2 = quantile(tmp_mat(:),0.01);
%s1 = max(tmp_mat(:));
%s2 = min(tmp_mat(:));
caxis_min_srfexc = -1*max(abs(s1),abs(s2));
caxis_min_srfexc = round(caxis_min_srfexc);
caxis_max_srfexc = max(abs(s1),abs(s2));
caxis_max_srfexc = round(caxis_max_srfexc);

c_length = caxis_max_srfexc/0.5;
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,:) = 0.85*ones(2,3);
cmap_srfexc = cmap_mat;

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


s1 = quantile(tmp_mat(:),0.99);
s2 = quantile(tmp_mat(:),0.01);
%s1 = max(tmp_mat(:));
%s2 = min(tmp_mat(:));
caxis_min_rzexc = -1*max(abs(s1),abs(s2));
caxis_min_rzexc = round(caxis_min_rzexc);
caxis_max_rzexc = max(abs(s1),abs(s2));
caxis_max_rzexc = round(caxis_max_rzexc);

caxis_min_rzexc = -12;
caxis_max_rzexc = 12;

c_length = caxis_max_rzexc/2;
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,:) = 0.85*ones(2,3);
cmap_rzexc = cmap_mat;

disp('plot data')
v_sm = VideoWriter([save_path storm_name '_SRFEXC_incr_ttot24']);
v_sm.FrameRate = 4;
open(v_sm);

v_rz = VideoWriter([save_path storm_name '_RZEXC_incr_ttot24']);
v_rz.FrameRate = 4;
open(v_rz);


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)
    setm(ax,'MLineLocation',5,'PLineLocation',5,...
         'ParallelLabel','on','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,exp_srfexc_ttot24_incr_mat(:,i),'s','filled')
    caxis([caxis_min_srfexc caxis_max_srfexc]);
    set(gcf,'Colormap',cmap_srfexc);
    plotm(coastlat, coastlon,'k')
    hold on
    plotm(track_data(1:date_index,5),track_data(1:date_index,4),'-m','LineWidth',2)
    hold off
    set(gca,'XLim',lonlim);
    set(gca,'YLim',latlim);
    tightmap
    title(['SRFEXC increment [kg m^{-2}] ' datestr(datenum_vec(i),'yyyy-mm-dd')])
    h = colorbar;
    set(h,'Location','southoutside','Orientation','horizontal','FontSize',14);
    set(gcf,'paperunits','centimeters')
    set(gcf,'paperposition',[1 1 17 15])
    saveas(gcf,[save_path storm_name '_incr_SRFEXC_' datestr(datenum_vec(i),'yyyy-mm-dd')],'png')

    frame = getframe(gcf);
    writeVideo(v_sm,frame);

    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)
    setm(ax,'MLineLocation',5,'PLineLocation',5,...
         'ParallelLabel','on','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,exp_rzexc_ttot24_incr_mat(:,i),'s','filled')
    caxis([caxis_min_rzexc caxis_max_rzexc]);
    set(gcf,'Colormap',cmap_rzexc);
    plotm(coastlat, coastlon,'k')
    hold on
    plotm(track_data(1:date_index,5),track_data(1:date_index,4),'-m','LineWidth',2)
    hold off
    set(gca,'XLim',lonlim);
    set(gca,'YLim',latlim);
    tightmap
    title(['RZEXC increment [kg m^{-2}] ' datestr(datenum_vec(i),'yyyy-mm-dd')])
    h = colorbar;
    set(h,'Location','southoutside','Orientation','horizontal','FontSize',14);
    set(gcf,'paperunits','centimeters')
    set(gcf,'paperposition',[1 1 17 15])
    saveas(gcf,[save_path storm_name '_incr_RZEXC_' datestr(datenum_vec(i),'yyyy-mm-dd')],'png')

    frame = getframe(gcf);
    writeVideo(v_rz,frame);

end

close(v_sm);
close(v_rz);

