%% Module 2: Train a set of NNs to predict P, E, ΔS, and R 
% based on the OI solution
% NN 12-5: includes Calibration + Mixture

clc
clear
close all 
cd(fileparts(matlab.desktop.editor.getActiveFilename))
src = matlab.desktop.editor.getActiveFilename; 

% Set to `true` to do a fast training for testing
% Or `false` for production
bTesting = false;

% Number of neurons in the hidden layer(s)
net_size = 10;

n_months = 240;


% Which category, P=1, E=2, R=4
% This is a mapping of my EO variables to the targets (OI soln for P, E, S, R) 
categories = [1, 1, 1, ...
              2, 2, 2, ...
              3, 3, 3, ...
              4];

var_names = [
    "P_GPCP";        %1   |  3 Precip datasets
    "P_GPM_IMERG";   %2
    "P_MSWEP";       %3 
    "E_Gleam_a";     %4   |  4 ET datasets
    "E_Gleam_b";     %5
    "E_ERA5";        %6
    "S_CSR" ;        %7   |  3 Storage datasets
    "S_GSFC";        %8
    "S_JPL";         %9
    "R_GRUN"              %10  |  1 Time series of runoff
];

catnames = ["Precipitation", "Evapotranspiration", "Change in Storage", "Runoff"];

ancillary_names = ["arid", "elev", "lat", "slope", "veg", "irrig", "lng", ...
                   "burned", "snow", "solar", "temperature", "veg_delta"];

n_ancillaries = 12;

% Get the input and target data
[X, A, Y, basins, months] = get_nn_inputs12_5("train");

n_basins = length(basins);  % 80% of our basins for training, currently 1,358

%% Check PLOTS

% time series plot
basin = 288;
x = squeeze(X(:, basin, :));
y = squeeze(Y(:, basin, :));
a = squeeze(A(:, basin, :));

dt = datetime(2000, 1:240, 15);

figure('units','normalized','outerposition',[0 0 1 1]);
cols = {[1,2,3], [4,5,6], [7, 8, 9], 10};
for i = 1:4
    subplot(4, 1, i);
    plot(dt, x(:, cols{i}));
    hold on
    plot(dt, y(:, i), 'Color','r', 'LineWidth', 1);
    ylabel(catnames(i));
    title(catnames(i));
end


%% Time series plots of the ancillary vars in a single basin
for i = 1:n_ancillaries
    subplot(3, 4, i);
    plot(dt, a(:, i));
    ylabel(ancillary_names(i))
end


%% Summary Statistics
disp("EO")
summary_stats(x, "headers", var_names)
disp("OI")
summary_stats(y, "headers", catnames)
disp("Ancillaries")
summary_stats(a, "headers", ancillary_names)


%% Summary stats of ALL ancillaries
A_vectorized = reshape(A, [n_months * n_basins, n_ancillaries]);
summary_stats(A_vectorized, "headers", ancillary_names)
clear A_vectorized


%% Histogram of the Ancillaries, just as check.
figure('units','normalized','outerposition',[0 0 1 1]);

A_vectorized = reshape(A, n_months * n_basins, n_ancillaries);
for i = 1:n_ancillaries
    subplot(3, 4, i);
    histogram(A_vectorized(:, i));
    title(ancillary_names(i));
end

clear A_vectorized

%% Map check on the ancillaries
A_matrix = reshape(A, [n_months, n_basins, n_ancillaries]);
G = get_basin_data("train");

figure('units','normalized','outerposition',[0 0 1 1]);
for i = 1:n_ancillaries
    subplot(3, 4, i);
    a = A_matrix(55, :, i);
    scatter(G.lon, G.lat, 8, a, 'filled', 'o');
    colorbar
    title(ancillary_names(i));
    coasts 'k';
    mapzoom(8);
end
% Everything looks as I expect based on this series of checks.

%% Remove or FILTER OUT rows with a NaN in either X, A, or Y
% Vectorize for passing to the NN functions
X = reshape(X, [n_months * n_basins, 10]);
Y = reshape(Y, [n_months * n_basins, 4] );
A = reshape(A, [n_months * n_basins, n_ancillaries] );
months = months(:);
basins = basins(:);

validEntries = ~isnan([X, A, Y]);      % creates an n x m matrix where enties are 1 if not NaN, or 0 if NaN
validRows = all(validEntries, 2);      % creates an n x 1 matrix where entries do an AND across rows. 
                                       % i.e.: returns 1 (true) if all true or 0 (false) if even a single entry in the row is false                        
X_filtered = X(validRows, :);
Y_filtered = Y(validRows, :);
A_filtered = A(validRows, :);
months_filtered = months(validRows);
basins_filtered = basins(validRows);

[n_samples, ~] = size(X); 
[n_filtered, ~] = size(X_filtered);

months_remaining = uniquevals(months_filtered);
basins_remaining = uniquevals(basins_filtered);


%% NN
% Cell arrays to save the networks
calnets = {};
mixnets = {};

% Create the NEURAL NETWORKs

% Let C contain the output of each of the calibration networks
C = nan(10, n_filtered);

close all
figure('units','normalized','outerposition',[0 0 0.75 1]);

for v = 1:10

    fprintf('Training Calibration NN %g of 10 \n', v);

    % For each simulation, we need to know which category {1: P, 2: E, 3:S, 4:R}
    % the variable is in so that we can select the correct target.
    var_category = categories(v);
        
    % Inputs: Single variable plus Ancillaries
    % use the variable v to select the correct column
    Xin = [ X_filtered(:, v), A_filtered ]';

    % Targets: Create the Output layer.
    T = Y_filtered(:, var_category)';

    % For debugging, make a plot
    subplot(3, 4, v);
    x = T';
    y = Xin(1, :)';
    scatter(x, y, 2, '.', 'MarkerFaceAlpha', 0.1, 'MarkerEdgeAlpha', 0.1);
    R2onPlot(gca, x, y, true);
    title(var_names(v), "Interpreter", "none");
    myAxisEqual(gca);

    % Create neural network
    net = feedforwardnet(net_size);
    net.divideFcn = 'dividerand';
    net.trainParam.showWindow = 0;

    if bTesting
        net.trainParam.epochs = 10;
    end

    % Train the network
    net = train(net, Xin, T, 'useParallel', 'yes');

    % Get the predictions of the calibration mode
    N = sim(net, Xin);

    % Force P and R to be positive
    if v < 4 || v == 10
        N(N < 0) = 0; 
    end

    % Save the trained network
    calnets{v} = net;

    % Put the results of the simulation into our matrix C, for "Calibrated"
    C(v, :) = N;
end

sgtitle("Before calibration: Observations on X axis, OI soln. on Y axis");
fname = "plots/scatter_obs.jpg";
exportgraphics(gcf, fname, 'Resolution', 150);


%% Scatterplots of NN calibrated (y) vs. target OI (x) for all 10 variables over the training basins

figure('units','normalized','outerposition',[0 0 0.75 1]);
for v = 1: 10
    y = C(v, :)';
    target_col = categories(v);
    x = Y_filtered(:, target_col);
    ax = subplot(3, 4, v);
    scatter(x, y, 2, '.', 'MarkerFaceAlpha', 0.1, 'MarkerEdgeAlpha', 0.1);
    title(var_names(v), "Interpreter", "none");
    myAxisEqual(ax);
    one2one(ax)
    R2onPlot(ax, x, y, true);
end

figmeta(src);
sgtitle("Post-Calibration Calibrated Obs. on the X axis, OI target on the Y axis");
fname = "plots/scatter_cal_vs_oi.jpg";
exportgraphics(gcf, fname, 'Resolution', 150);


%% PART 2: MIXTURE MODEL
% Run 3 times, each time with a different output: (1: P, 2: E, 3: R)
disp("Mixture model");

figure('units','normalized','outerposition',[0 0 0.75 1]);

for m = 1:3

    fprintf('Training Mixture NN %g of 3 \n', m);
    
    % Get the name of the category of variable we are simulating (Precip, Evap, Storage)
    catname = catnames(m);

    % Inputs are the calibrated EO variables of the same type (P1, P2, P3)
    % plus the ancillaries. For this version, we are NOT mixing unlike vars. 
    switch m
        case 1
            Xin = [C(1:3, :); A_filtered'];
        case 2
            Xin = [C(4:6, :); A_filtered'];
        case 3 
            Xin = [C(7:9, :); A_filtered'];
    end
    
    target_col = m;
    T = Y_filtered(:, target_col)';

    % Create network
    net = feedforwardnet(net_size);
    net.divideFcn = 'dividerand';
    net.trainParam.showWindow = 0;

    if bTesting
        net.trainParam.epochs = 10;
    end

    % Train the network
    net = train(net, Xin, T, 'useParallel', 'yes');

    % Save the network
    mixnets{m} = net;
    
    % Get the values of the mixture net output for the plot
    M = sim(net, Xin);
    subplot(2, 2, m);
    x = T';
    y = M';
    scatter(x, y, 2, '.', 'MarkerFaceAlpha', 0.1, 'MarkerEdgeAlpha', 0.1);
    title(catname);
    myAxisEqual(gca);
    R2onPlot(gca, x, y, true);
end 

figmeta(src);
sgtitle("12-5 Mixture Model Output on Y axis, OI soln. (target) on X axis over training basins (n = 1,358)")
fname = "plots/scatter_mix_vs_oi.jpg";
exportgraphics(gcf, fname, 'Resolution', 150);

% We should now have a full set of trained networks.

%% Run a SIMULATION with the trained networks
% over the validation basins

clear x y X Y Y_filtered X_filtered A_filtered Xin XM YM

figure('units','normalized','outerposition',[0 0 1 1]);

[X, A, Y] = get_nn_inputs12_5("validate");

% revise the number of observations; there are 340 validation basins
n_basins = 340;

% Let M be the output of the cascade of models (cols for P, E, R)
n = n_months * n_basins;

C = nan(10, n); % output of the calibration models
M = nan(4, n); % output of the mixture model

% Calibration networks
A_vectorized = reshape(A, [n_months * n_basins, n_ancillaries]);
for v = 1:10
    %Xin = [X(:, v), A];
    Xin = reshape(X(:, :, v), [n_months * n_basins, 1]);
    Xin = [Xin, A_vectorized]; 
    net = calnets{v};
    P = sim(net, Xin');
    if v < 4 || v == 10
        P(P < 0) = 0;
    end
    C(v, :) = P;

    subplot(3, 4, v);
    target_col = categories(v);
    x = Y(:, :, target_col);
    x = x(:);
    y = Xin(:, 1);
    scatter(x, y, 2, '.', 'MarkerFaceAlpha', 0.1, 'MarkerEdgeAlpha', 0.1);
    R2onPlot(gca, x, y, true);
    title(var_names(v), "Interpreter", "none");
    myAxisEqual(gca);

end

figmeta(src); 
sgtitle("12-5 EO inputs (x) vs. NN calibrated (Y), for the validation basins (n = 340)");
fname = "plots/scatter_val_cal_vs_eo.jpg";
exportgraphics(gcf, fname, 'Resolution', 150);

%% Mixture Networks

figure('units','normalized','outerposition',[0 0 0.75 1]);

for m = 1:4
    switch m
        case 1
            Xin = [C(1:3, :); A_vectorized'];
        case 2
            Xin = [C(4:6, :); A_vectorized'];
        case 3
            Xin = [C(7:9, :); A_vectorized'];
    end

    if m < 4
        net = mixnets{m};
        P = sim(net, Xin);
        if m == 1
            P(P < 0) = 0;
        end
        M(m, :) = P;
    else
        M(4, :) = C(10, :);
    end

    % plots - let's just see how close we get to the target
    subplot(2, 2, m);
    x = Y(:, :, m);
    x = x(:);
    y = M(m, :)';
    scatter(x, y, 2, '.', 'MarkerFaceAlpha', 0.1, 'MarkerEdgeAlpha', 0.1);
    R2onPlot(gca, x, y, true);
    title(catnames(m));
    myAxisEqual(gca);
    
end

figmeta(src);
sgtitle('12-5 Validation, mix NN (y) vs. OI targets (X)');
fname = "plots/scatter_val_mix_vs_oi.jpg";
exportgraphics(gcf, fname, 'Resolution', 150);


%% Save
% Use a separte set of scripts for analysis-- calculate the Imbalance, plots, summaries, etc. 

fname = "output/nn_results12-5.mat";
descr = "Trained network and predictions over validation basins for NN12-5.";
metadata = filemeta(descr, src);
save(fname, "calnets", "mixnets", "C", "M", "metadata", "-v7.3");
