clear;close all;clc; 

addpath('/Users/xudo627/donghui/CODE/Setup-E3SM-Mac/matlab-scripts-for-mosart/');
addpath('/Users/xudo627/donghui/mylib/m/');
addpath('/Users/xudo627/donghui/CODE/dengwirda-inpoly-355c57c/');
addpath('/Users/xudo627/developments/getPanoply_cMap/');
addpath('functions/');

re = 6.37122e6;% Earth radius
cmap1 = getPanoply_cMap('NEO_ceres_net');
cmap2 = getPanoply_cMap('NEO_omi_ozone_to3');

datadir = '/Volumes/LaCie/PROJECTS/Pantanal/';
yrs = 1979;
yre = 2010;

update = 0;

[lon_gsim,lat_gsim,Sh,yr,mo,da,mu_gsim,sd,cv] = get_GSIM_discharge('US_0000919',1);

forcing = 'CRUNCEPv7';%'GSWP3v1';%'CRUNCEPv7';
ex      = 'run3'; % sub, sur2
tag = ['SUS_outputs_' ex '_CLM' forcing];
area = areaint(Sh.Y,Sh.X)*4*pi.*(re^2); 

xv   = ncread('inputdata_run/domain_lnd_Susquehanna_run3_template.nc','xv');
yv   = ncread('inputdata_run/domain_lnd_Susquehanna_run3_template.nc','yv');
xc   = ncread('inputdata_run/domain_lnd_Susquehanna_run3_template.nc','xc');
yc   = ncread('inputdata_run/domain_lnd_Susquehanna_run3_template.nc','yc');
aream2 = ncread('inputdata_run/MOSART_Susquehanna_run3_template_c230518.nc','area');

numc = length(xc);

yr1 = 1979;
yr2 = 2010;
if exist('data/SUS_outputs_default_CLMCRUNCEPv7.mat','file')
    def = load('data/SUS_outputs_default_CLMCRUNCEPv7.mat');
else
    % Read MOSART
    [~, q, fdfall, fdf, R2, NSE, KGE, ind1, ind2, imax]                    = ...
             read_mosart(datadir, 'SUS_outputs_default_CLMCRUNCEPv7', yr1, yr2, mu_gsim, lon_gsim, lat_gsim, yr, area,1);
    for i = 1 : 1
        fdf(i,:) = nanmean(fdfall((i-1)*numc+1:i*numc,:),1);
    end
    % Read ELM
    [fh2osfc,qrunoff,qdrai,fsat,fh2osfcall,qrunoffall,qdraiall,zwt,zwtall,twsall,et,etall,rain] = ...
             read_elm(datadir, 'SUS_outputs_default_CLMCRUNCEPv7', 1, numc, ind2);  
    save('data/SUS_outputs_default_CLMCRUNCEPv7.mat','q','fdfall','fdf','R2', ...
         'NSE','KGE','fh2osfc','qrunoff','qdrai','fsat','fh2osfcall','qrunoffall', ...
         'qdraiall','zwt','zwtall','twsall','et','etall','rain');
end

tag2 = ['SUS_outputs_hydraulic_CLM' forcing];

if exist(['data/' tag2 '.mat'],'file') && ~update
    step2 = load(['data/' tag2 '.mat']);
else
    % Read MOSART
    [ntot, q, fdfall, fdf, R2, NSE, KGE, ind1, ind2, imax]                    = ...
             read_mosart(datadir, tag2, yr1, yr2, mu_gsim, lon_gsim, lat_gsim, yr, area,1000);
    for i = 1 : ntot
        fdf(i,:) = nanmean(fdfall((i-1)*numc+1:i*numc,:),1);
    end
    save(['data/' tag2 '.mat'], 'q','ntot','fdf','fdfall','R2','NSE','KGE','ind1','ind2','imax');
end

if exist(['data/' tag '.mat'],'file') && ~update
    load(['data/' tag '.mat']);
else
    % Read MOSART
    [ntot, q, fdfall, fdf, R2, NSE, KGE, ind1, ind2, imax]                    = ...
             read_mosart(datadir, tag, yr1, yr2, mu_gsim, lon_gsim, lat_gsim, yr, area);
    for i = 1 : ntot
        fdf(i,:) = nanmean(fdfall((i-1)*numc+1:i*numc,:),1);
    end
    % Read ELM
    [fh2osfc,qrunoff,qdrai,fsat,fh2osfcall,qrunoffall,qdraiall,zwt,zwtall,twsall,et,etall,rain] = ...
             read_elm(datadir, tag, ntot, numc, ind2);  
    save(['data/' tag '.mat'], '-regexp', '^(?!(show_domain|show_sub_basin||show_inund|show_default_BFI|show_annual)$).');
end

load('SUS_BFI3.mat');
numc = length(BFI);


SWF = load('~/DATA/GLAD/Susquehanna_swf.mat');
SWF05 = NaN(length(xc),1);
tmpswf = nanmean(SWF.swf2,3);
for i = 1 : length(xc)
    in = inpoly2([SWF.lon(:) SWF.lat(:)],[xv(:,i) yv(:,i)]);
    SWF05(i) = nanmean(tmpswf(in))/100;
end
clear tmpswf;

rng(1);
xtrain = NaN(ntot,6);
xtrain(:,1)= 0.1 + (5-0.1)*rand(ntot,1);
xtrain(:,2)= 0.1 + (5-0.1)*rand(ntot,1);
xtrain(:,3)= 0.001 + (1-0.001)*rand(ntot,1); % micro_sigma
xtrain(:,4)= 0.001 + (0.5-0.001)*rand(ntot,1); % fc   
xtrain(:,5)= ones(ntot,1).*7.2;  
xtrain(:,6)= ones(ntot,1).*0.26;  

BFI_sims = nanmean(qdraiall./qrunoffall,2);
BFI_def = nanmean(def.qdraiall./def.qrunoffall,2);
inund_def = nanmean(def.fdfall(:,20*12+1:end),2) + nanmean(def.fh2osfcall(:,20*12+1:end),2);
inund = fdfall(:,20*12+1:end) + fh2osfcall(:,20*12+1:end);
inund = nanmean(inund,2);
sensobs = sens_slope(nanmean(reshape(mu_gsim(yr >= yr1 & yr <= yr2),[12 32]),1)');
[~, ~, ~, pobs] = mk_test(nanmean(reshape(mu_gsim(yr >= yr1 & yr <= yr2),[12 32]),1)');

BFI_cal = NaN(numc,1);
SWF_cal = NaN(numc,1);
zwt_cal = NaN(numc,1);
run_cal = NaN(numc,384);
fdrain      = NaN(numc,1);
fover       = NaN(numc,1);
micro_sigma = NaN(numc,1);
fc          = NaN(numc,1);

w1 = 1;%1;
w2 = 10;%10;
w3 = 1.5;%1.5;

for i = 1 : numc
    disp(i);
    tmp1 = BFI_sims(i : numc : end);
    tmp2 = inund(i : numc :end);
    tmp3 = qrunoffall(i:numc:end,:);
    tmp4 = zwtall(i:numc:end,:);

    sens = NaN(ntot,1);
    for j = 1 : ntot
        tmp = nanmean(reshape(tmp3(j,:)',[12,32]),1)./1000 .* aream2(i);
        sens(j) = sens_slope(tmp(1:32)').*numc;
    end

    ind = find(xtrain(:,1) <= 0.2);
    tmp1(ind) = NaN;
    tmp2(ind) = NaN;
    tmp3(ind,:) = NaN;
    tmp4(ind,:) = NaN;
 
    obj =  w1*abs(log(tmp2./SWF05(i))) + ...
           w2*abs(log(tmp1./BFI(i)))   + ...
           w3*abs(log(sens./sensobs));
    ind = find(obj == min(obj)); ind = ind(1);
    
    BFI_cal(i)   = tmp1(ind);
    SWF_cal(i)   = tmp2(ind);
    run_cal(i,:) = tmp3(ind,:); 

    fdrain(i)      = xtrain(ind,1);
    fover(i)       = xtrain(ind,2);
    micro_sigma(i) = xtrain(ind,3);
    fc(i)          = xtrain(ind,4);
end

save('SUS_calibrated_hydrology.mat','fdrain','fover','micro_sigma','fc');

figure; set(gcf,'Position',[10 10 1200 1000]);
ax1 = subplot(3,3,1);
patch(xv,yv,BFI,'LineStyle','none'); hold on; clim([0.4 0.6]);
add_title(gca,'(a) GSCD BFI'); colormap(gca,cmap1);
ax2 = subplot(3,3,2);
patch(xv,yv,BFI_cal,'LineStyle','none'); hold on; clim([0.4 0.6]); 
set(gca,'XTick',[],'YTick',[]);
add_title(gca,'(b) Calibrated BFI'); colormap(gca,cmap1);
ax3 = subplot(3,3,3);
patch(xv,yv,BFI_def,'LineStyle','none'); hold on; clim([0.4 0.6]); cb1 = colorbar('west');
set(gca,'XTick',[],'YTick',[]);
add_title(gca,'(c) Default BFI'); colormap(gca,cmap1);

ax4 = subplot(3,3,4);
patch(xv,yv,SWF05,'LineStyle','none'); hold on; clim([0.0 0.05]);
add_title(gca,'(d) GLAD SWF'); colormap(gca,cmap2);
ax5 = subplot(3,3,5);
patch(xv,yv,SWF_cal,'LineStyle','none'); hold on; clim([0.0 0.05]);
add_title(gca,'(e) Calibrated SWF'); colormap(gca,cmap2);
set(gca,'XTick',[],'YTick',[]);
ax6 = subplot(3,3,6);
patch(xv,yv,inund_def,'LineStyle','none'); hold on; clim([0.0 0.05]); cb2 =  colorbar('west'); 
add_title(gca,'(f) Default SWF'); colormap(gca,cmap2);
set(gca,'XTick',[],'YTick',[]);

ax7 = subplot(3,3,7:9);
k = 1;
for i = yr1 : yr2
    for j = 1 : 12
        tmon(k,1) = datenum(i,j,15,0,0,0);
        k = k + 1;
    end
end

h(1) = plot(tmon,mu_gsim(yr >= yr1 & yr <= yr2),'k--','LineWidth',2); hold on; grid on;
h(2) = plot(tmon,def.q,'r-','LineWidth',2);
h(3) = plot(tmon,step2.q(find(step2.R2 == max(step2.R2)),:),'b-','LineWidth',2);
% h(4) = plot(tmon,q(find(R2 == max(R2)),:),'g-','LineWidth',2);
xlim([tmon(180) tmon(260)]);
datetick('x','mmmyy','keepticks');
set(gca,'FontSize',14);
% h(1) = plot(yr1:yr2,nanmean(reshape(mu_gsim(yr >= yr1 & yr <= yr2),[12 32]),1),'k--','LineWidth',3); hold on;
% h(2) = plot(yr1:yr2,nanmean(reshape(nanmean(run_cal,1)',[12 32]),1)./1000.*sum(aream2),'b-','LineWidth',3);
% h(3) = plot(yr1:yr2,nanmean(reshape(def.q',[12 32]),1),'r-','LineWidth',3); grid on;
add_title(gca,'(g) Monthly discharge');
% ylim([0 2000]); set(gca,'FontSize',14);
ylabel('[m^{3}/s]','FontSize',15,'FontWeight','bold');
leg = legend(h,{'Obervvation','Default','Two-steps calibration'}); leg.FontSize = 15; leg.FontWeight = 'bold';
leg.Location = 'northwest';
ax2.Position(1) = ax2.Position(1) - 0.04;
ax3.Position(1) = ax3.Position(1) - 0.08;
ax5.Position(1) = ax5.Position(1) - 0.04;
ax6.Position(1) = ax6.Position(1) - 0.08;
ax7.Position(3) = ax7.Position(3) - 0.08;

cb1.Position = [ax3.Position(1) + ax3.Position(3) + 0.01 ax3.Position(2) 0.02 ax3.Position(4)];
cb2.Position = [ax6.Position(1) + ax6.Position(3) + 0.01 ax6.Position(2) 0.02 ax6.Position(4)];

cb1.FontSize = 13;
cb2.FontSize = 13;

sensobs = round(sensobs,2); pobs = round(pobs,2);

ARdef  = nanmean(def.q) / nanmean(mu_gsim(yr >= yr1 & yr <= yr2));
rSDdef = nanstd(def.q) / nanstd(mu_gsim(yr >= yr1 & yr <= yr2));
ARcal  = nanmean(step2.q(find(step2.R2 == max(step2.R2)),:)) / nanmean(mu_gsim(yr >= yr1 & yr <= yr2));
rSDcal = nanstd(step2.q(find(step2.R2 == max(step2.R2)),:)) / nanstd(mu_gsim(yr >= yr1 & yr <= yr2));

ARhyd  = nanmean(q(find(R2 == max(R2)),:)) / nanmean(mu_gsim(yr >= yr1 & yr <= yr2));
rSDhyd = nanstd(q(find(R2 == max(R2)),:)) / nanstd(mu_gsim(yr >= yr1 & yr <= yr2));

rhodef = round(sqrt(def.R2),2);        ARdef = round(ARdef,2); rSDdef = round(rSDdef,2);
rhocal = round(sqrt(max(step2.R2)),2); ARcal = round(ARcal,2); rSDcal = round(rSDcal,2);

strs1 = ['\rho = ' num2str(rhodef) ', AR = ' num2str(ARdef) ', rSD = ' num2str(rSDdef)];
strs2 = ['\rho = ' num2str(rhocal) ', AR = ' num2str(ARcal) ', rSD = ' num2str(rSDcal)];

t1   = annotation('textbox',[0.525 0.225 0.3 0.1],'String',strs1,       ... 
                'FitBoxToText','on','EdgeColor','none','Color','r',      ...
                'FontSize',14,'FontWeight','bold','fontname','Menlo'); 
t2   = annotation('textbox',[0.525 0.200 0.3 0.1],'String',strs2,       ... 
                'FitBoxToText','on','EdgeColor','none','Color','b',      ...
                'FontSize',14,'FontWeight','bold','fontname','Menlo'); 

senscal = sens_slope(nanmean(reshape(step2.q(find(step2.R2 == max(step2.R2)),:)',[12 32]),1)');
[~, ~, ~, pcal] = mk_test(nanmean(reshape(step2.q(find(step2.R2 == max(step2.R2)),:)',[12 32]),1)');
senscal = round(senscal,2); pcal = round(pcal,2);


sensdef = sens_slope(nanmean(reshape(def.q',[12 32]),1)');
[~, ~, ~, pdef] = mk_test(nanmean(reshape(def.q,[12 32]),1)');
sensdef = round(sensdef,2); pdef = round(pdef,2);


% plot(nanmean(reshape(q(NSE == max(NSE),:)',[12 32]),1),'r-','LineWidth',2);
% figure;
% for i = 1 : ntot
%     plot(1:12,nanmean(reshape(q(i,:)',[12 32]),2),'-'); hold on; grid on;
% end
% plot(1:12,nanmean(reshape(mu_gsim(yr >= yr1 & yr <= yr2),[12 32]),2),'k-','LineWidth',4); 
% plot(1:12,nanmean(reshape(q(NSE == max(NSE),:)',[12 32]),2),'r-','LineWidth',4);
% figure;
% plot(mu_gsim(yr >= yr1 & yr <= yr2),'k-','LineWidth',2); hold on;
% plot(q(R2 == max(R2),:),'b-','LineWidth',2);