%% Optimize equation: iGrav = series1 - series2...
% Prior running this script, prepare the input time series using
% \Scripts\DataProcessing\Prepare_Optimization_Inputs.m
%
% The optimization is based on solving simple equation containing time
% series (on right side of the equation) that are:
%	multiplied and added/subtracted = precipitaion, evapotranspiration,
%	runoff, global hydrological effect (other optional time series such as
%	atmospheric effect)
%
% Requirements:
%   Optimization Toolbox
%
%                                                   M. Mikolaj, GFZ Potsdam

addpath(fullfile('..','MatlabLibrary','hydroGravityLib'));
addpath(fullfile('..','MatlabLibrary','aux_fce'));
clear 
close all
clc
%% Main Settings
% Set all required parameters for right side of the equation:
%   name:       name of the time series
%   input_file: input file name
%   column:     column of the time series in the input_file
%   ini_val:    initial value for optimization
%   max_val:    maximim value for optimization
%   min_val:    minimum value for optimization
%
% Use temporary/auxiliary 'input_seriesX' for simple filling of input
% parameters of 'series1' (just to spare some space in the code)
in_series1 = fullfile('..','..','Data','Optimization','Optimization_Inputs.tsf'); 
%
% Set:      name  ,input file,column,ini_val,max_val,min_val
series1 = {'Glob' ,in_series1,    3 ,   1.00,  1.282,  0.489;...%'Atmo' ,in_series1,    2,    0.10,   0.900, -0.900;...
           'Rain' ,in_series1,    5 ,   1.00,  1.100,  0.900;...
           'Evap' ,in_series1,    6 ,  -1.00,  0.000, -1.000;...
           'Disc' ,in_series1,    4 ,  -1.00, -0.900, -1.100};
% 
column_gravity = 1; % column with gravity residuals corrected for all effect except global hydrology
%
% Set paramter convertig gravity into water column. 
response = 0.478; % nm/s^2/mm
%
% Use following time interval only 
start = [2015 06 19];
stop =  [2016 06 29];
%
% Set file with Lysimeter ET time series
% Lysimeter ET/Rain
lysi.in_file = fullfile('..','..','Data','HydroMeteo','IGETS-AUX-we006-ALL.tsf');
lysi.in_evap = 3; % column with ET values
lysi.in_rain = 2; % column with rain values

% Set output file name
output_file = fullfile('..','..','Data','Optimization','Optimization_Outputs.tsf'); 
%Output figure
output_figure = fullfile('..','..','Manuscript','Figures','Figure10.jpg');

%% Output/Plot settings
% Plot font size
font_size = 11;
% XTick label
xtl = sort([datenum(2015,1:3:12,1),datenum(2016,1:3:12,1)]);

%% Load data
% Load gravity
[time.gravity,data.gravity] = loadtsf(in_series1);
% Remove time interval out of requested range
data.gravity = data.gravity(time.gravity>datenum(start) & time.gravity<datenum(stop),column_gravity);
time.gravity = time.gravity(time.gravity>datenum(start) & time.gravity<datenum(stop),:);
% Subtract the first value to make sure all time series start with zero
% (in case of series2, such correction will be applied within the
% optimization function)
data.gravity = data.gravity - data.gravity(1);

% Load series1
for i = 1:size(series1,1)
    % Load time series
    [temp_time,temp_data] = loadtsf(series1{i,2});
    data.series1(:,i) = temp_data(temp_time>datenum(start) & temp_time<datenum(stop),series1{i,3});
end

% Load Lysimeter data
[lysi.time,temp] = loadtsf(lysi.in_file);
lysi.data = temp(:,[lysi.in_evap,lysi.in_rain]);
lysi.data = lysi.data(lysi.time>datenum(start) & lysi.time<datenum(stop),:);
lysi.time = lysi.time(lysi.time>datenum(start) & lysi.time<datenum(stop),:);
[lysi.time_day,temp] = data2daily(lysi.time,lysi.data,2,0);
lysi.time_day = datenum(lysi.time_day) + 12/24; % add 12 hours to point to the noon. Independent of 'time_shift' parameter
lysi.evap_day = temp(:,1);
lysi.rain_day = temp(:,2);

%% Prepare data=bounds for optimization
% Copy original parameter specifications. Some of them will be
% converted from gravity to water column
param_series1.ini_val = cell2mat(series1(1:size(series1,1),4));
param_series1.max_val = cell2mat(series1(1:size(series1,1),5));
param_series1.min_val = cell2mat(series1(1:size(series1,1),6));
% Convert series1 parameters
for i = 1:length(param_series1.ini_val)
    % Convert specified ranges to water column (nm/s^2 -> mm)
    % using gravity response. These time series stay in original units! 
    % Only ranges are modified!! Applies to all gravity time series
    if strcmp(series1{i,1},'Glob') || strcmp(series1{i,1},'NTOL') || strcmp(series1{i,1},'Atmo') ||...
       strcmp(series1{i,1},'Drift') || strcmp(series1{i,1}(1:4),'Step') || strcmp(series1{i,1},'Const')
        % Convert ranges
        param_series1.ini_val(i) = param_series1.ini_val(i)/response;
        param_series1.max_val(i) = param_series1.max_val(i)/response;
        param_series1.min_val(i) = param_series1.min_val(i)/response;
    end
    % Shift so time series start at zero. 
    data.series1(:,i) = data.series1(:,i) - data.series1(1,i);
end

%% Opimization
% Call optimization function. This function will automatically remove NaNs
% and sum all time series muplitplied by parameters to obtain the TWS as
% given by gravimeter.
[param_opt,opt_tws,opt_res,NS] = Optimization_iGrav_TWS_annual_function(data.gravity/response,data.series1,param_series1);

%% Display Results & compute output time series
% Clear command prompt before writting results
clc
% Define output time series
plot_tws_opt(1:length(time.gravity),1:2) = 0;
plot_grav_opt = data.gravity;
% Print results in command line
fprintf('Results after optimization. Analysed interval: %04d/%02d/%02d - %04d/%02d/%02d\n',...
        start(1),start(2),start(3),stop(1),stop(2),stop(3));
% Aux paramter
c = 1;
% Write results to command prompt. One after another.
for i = 1:length(param_opt)
    % Get parameter names and ranges
    if i <= size(series1,1)
        param_names{i} = series1{i,1};
        param_range{i,1} = series1{i,5};
        param_range{i,2} = series1{i,6};
        % Convert ranges back to original units
        if strcmp(series1{i,1},'Glob') || strcmp(series1{i,1},'NTOL') || strcmp(series1{i,1},'Atmo') ||...
               strcmp(series1{i,1},'Drift') || strcmp(series1{i,1}(1:4),'Step') || strcmp(series1{i,1},'Const')
            param_opt(i) = param_opt(i)*response;
            % Subtract from gravity (=> gravity residuals)
            plot_grav_opt = plot_grav_opt - data.series1(:,i)*param_opt(i);
        else
            % Add up all TWS series
            plot_tws_opt(:,c) = data.series1(:,i)*param_opt(i);
            param_names_ws{c} = series1{i,1};
            c = c + 1;
        end
        % In addiotion get Precipitation and Evapotranspiration time series
        % for comparison to Lysimeter
        if strcmp(series1{i,1},'Evap') % get ET data for comparison to lysimeter
            plot_evap_opt = -data.series1(:,i)*param_opt(i);
            plot_evap_ref = data.series1(:,i);
        end
        if strcmp(series1{i,1},'Rain') % get Rain data for comparison to lysimeter
            plot_rain_opt = data.series1(:,i)*param_opt(i);
            plot_rain_ref = data.series1(:,i);
        end
    end
    % Write results to command prompt
    fprintf('%s\t=\t%9.2f (min = %9.2f, max = %9.2f)\n',param_names{i},param_opt(i),param_range{i,2},param_range{i,1});
end

%% Figure for manuscript
figure('Position',[200 400 800, 260],'paperpositionmode','auto')
% Optimized parameter
plot(time.gravity,[plot_grav_opt/response,sum(plot_tws_opt,2)],'linewidth',1);
[~,~,~,leg] = mm_setaxes(gca,'xlim',[datenum(start)-10,datenum(stop)+10],'fontsize',font_size,...
                    'legend',{'\rmd\itS\rm/\rmd\itt \rm{(left)}','P-E-R \rm{(right)}'},'ylabel','mm','grid','on',...
                    'xtick',xtl,'dateformat','dd.mm.yyyy','ylim',[-180 80]);
set(leg,'location','northwest','EdgeColor','none','Color','none','FontAngle','italic');
print(gcf,output_figure,'-djpeg','-r300');

%% Write statistics of the fit
% Compute correlation
cor = mmcorr(plot_grav_opt/response,sum(plot_tws_opt,2));
% Compute STD and mean
[out_std,out_mean] = mm_statnan(plot_grav_opt/response-sum(plot_tws_opt,2));
[~,~,~,~,tws_range] = mm_statnan(sum(plot_tws_opt,2));
% Write to command prompt
fprintf('\nFit statistics:\n');
fprintf('Kling-Gupta efficiency = %9.2f \n',NS);
fprintf('Correlation (Grav-GHE) vs (P-E-R) = %9.2f \n',cor);
fprintf('Standard deviation (Grav-GHE) - (P-E-R) = %9.2f mm\n',out_std);
fprintf('Mean value of (Grav-GHE) - (P-E-R) = %9.2f mm\n',out_mean);
fprintf('Range of P-E-R = %9.2f mm\n',tws_range);

%% Compare to lysimeter
% Compute daily ET from estimated time series
temp = [0;diff(plot_evap_opt)];
[~,opt_evap] = data2daily(time.gravity,temp,2,0);
temp = [0;diff(plot_evap_ref)];
[~,ref_evap] = data2daily(time.gravity,temp,2,0);
% Remove common NaNs
r = find(~isnan(opt_evap+lysi.evap_day+ref_evap));
sum_lysi_evap = sum(lysi.evap_day(r));
sum_opt_evap = sum(opt_evap(r));
sum_ref_evap = sum(ref_evap(r));
% Write to command prompt
fprintf('\nLysimeter ET comparison (within lysimeter time period):\n');
fprintf('Optimized/Reference  = %9.2f \n',sum_opt_evap/sum_ref_evap);
fprintf('Lysimeter/Reference  = %9.2f \n',sum_lysi_evap/sum_ref_evap);
fprintf('Annual difference (lysimeter) = %9.2f%% \n',(sum_lysi_evap-sum_opt_evap)/sum_ref_evap*100);

% Compute daily Rain from estimated time series
temp = [0;diff(plot_rain_opt)];
[~,opt_rain] = data2daily(time.gravity,temp,2,0);
temp = [0;diff(plot_rain_ref)];
[~,ref_rain] = data2daily(time.gravity,temp,2,0);
% Remove common NaNs
r = find(~isnan(opt_rain+lysi.rain_day+ref_rain));
sum_lysi_rain = sum(lysi.rain_day(r));
sum_opt_rain = sum(opt_rain(r));
sum_ref_rain = sum(ref_rain(r));
% Write to command prompt
fprintf('\nLysimeter Precipitation comparison (within lysimeter time period):\n');
fprintf('Optimized/Reference  = %9.2f \n',sum_opt_rain/sum_ref_rain);
fprintf('Lysimeter/Reference  = %9.2f \n',sum_lysi_rain/sum_ref_rain);
fprintf('Annual difference (lysimeter) = %9.2f%% \n',(sum_lysi_rain-sum_opt_rain)/sum_ref_rain*100);

%% Write outputs to a file
% File header (channel names and units)
header = {'Wettzell','iGrav','ResidualsIN','nm/s^2';...
          'Wettzell','iGrav','ResidualsCor','nm/s^2';...
          'Wettzell','iGrav','ResidualsIN','mm';...
          'Wettzell','iGrav','ResidualsOptimized','mm';...
          'Wettzell','iGrav','Precipitation','mm';...
          'Wettzell','iGrav','ET','mm';...
          'Wettzell','iGrav','Runoff','mm'
          'Wettzell','iGrav','TWS','mm'};
% time vector  +  output time series
data_write = [datevec(time.gravity),data.gravity,...
              plot_grav_opt,...
              data.gravity/response,...
              opt_tws,...
              abs(plot_tws_opt),...
              sum(plot_tws_opt,2)];
fprintf('\n Total sums:\n'); 
total_sums = [plot_grav_opt/response,plot_tws_opt];
% Compute difference between staring point and point after one (365days)
% year.
total_diff = interp1(time.gravity,total_sums,[time.gravity(1),time.gravity(1)+365]);
fprintf('Annual Precip: %.0f mm\n',total_diff(end,2)-total_diff(1,2));
fprintf('Annual Evapor: %.0f mm\n',total_diff(end,3)-total_diff(1,3));
fprintf('Annual Runoff: %.0f mm\n',total_diff(end,4)-total_diff(1,4));
fprintf('Annual deltaS: %.0f mm\n',total_diff(end,1)-total_diff(1,1));

% Comment containing all important input parameters
comment = {'Input parameters: name, input file, column, ini_val, max_val, min_val, optim_value'};
comment(2) = {sprintf('Gravity time series, %s, %02d, NaN, NaN, NaN, NaN',in_series1,column_gravity)};
for i = 1:size(series1,1)
    comment(i+2) = {sprintf('%s, %s, %02d, %6.4f, %6.4f, %6.4f, %6.4f',...
                    series1{i,1},series1{i,2},series1{i,3},series1{i,4},series1{i,5},series1{i,6},param_opt(i))};
end
comment(i+2) = {sprintf('Gravity response: %6.4f nm/s^2/mm. NSE/KGE = %6.3f',response,NS)};
comment(i+3) = {sprintf('Time interval analysed: %02d/%02d/%04d - %02d/%02d/%04d',...
                start(3),start(2),start(1),...
                stop(3),stop(2),stop(1))};
% Write
writetsf(data_write,header,output_file,2,comment);
          
rmpath(fullfile('..','MatlabLibrary','hydroGravityLib'));
rmpath(fullfile('..','MatlabLibrary','aux_fce'));