%% Module 3: Create a set of predictions at the pixel scale
% with NN 12-5

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


%% Load the ancillary variables
ancillary_names = ["aridity", "elevation", "latitude", "slope", "veg", ...
    "irrigation", "longitude", "burned", "snow", "solar", "temperature", "veg_delta"];

% A comes as 360 x 720 x 4
A_cube = nan(360, 720, 240, 12);

for i = 1:12
    fname = sprintf("data/grids05/ancillary/%s.mat", ancillary_names(i));
    load(fname, 'V');
    A_cube(:, :, :, i) = V;
end

A = reshape(A_cube, [360*720*240, 12]);

%summary_stats(samplerows(A, 0.001), "headers", ancillary_names)

%% Run the 4 mixture networks

% Second step is to read the EO data, one var at a time
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
    "DS_CSR" ;       %7   |  3 Storage datasets
    "DS_GSFC";       %8
    "DS_JPL";        %9
    "R_GRUN"         %10  |  1 Time series of runoff
];


data_folder = "data/grids05/eo";

% Networks 
load('output/nn_results12-5.mat', 'calnets');

C = nan(360 * 720 * 240, 10);
X = nan(360, 720, 240, 10);

for v = 1: 10
    fprintf("Calibrating variable #%g\n", v);
    fname = sprintf("%s/%s.mat", data_folder, var_names(v));
    load(fname, "D"); 
    % D is 360 x 720 x 480;
    D(:, :, 1:240) = [];
    X(:, :, :, v) = D;
    Xin = [D(:), A]';
    net = calnets{v}; 
    M = sim(net, Xin)';
    if v < 4 || v == 10
        M(M < 0) = 0;
    end
    C(:, v) = M;
end


%% Some QC this to verify that the input data looks right... 

X = reshape(X, [360 * 720 * 240, 10]);

catnames = ["Precip", "Evap", "Delta S", "Runoff"];

idx = randsample(62208000,50000);

disp("Calibrated")
summary_stats(C(idx, :), "headers", var_names)

disp("Inputs")
summary_stats(X(idx, :), "headers", var_names)


%% Plots
figure('units','normalized','outerposition',[0 0 1 1]);
for v = 1:10
    subplot(3, 4, subplot_column_indices(3, 4, v));
    x = X(idx, v);
    y = C(idx, v);
    plot(x, y, '.');
    myAxisEqual(gca);
    R2onPlot(gca, x, y, true);
    title(var_names(v), "Interpreter","none");
end
sgtitle("Calibration: x: input, y: calibrated, in 10,000 random pixels"); 


%% Mixture Model
clear calnets X D ANCILL A_cube net v V Xin x y i
load('output/nn_results12-5.mat', "mixnets");

M = nan(360*720*240, 3);

for m = 1:3
    fprintf("Running Mixture Model %g of 3 \n", m);
    switch m
        case 1
            Xin = [C(:, 1:3), A];
        case 2
            Xin = [C(:, 4:6), A];
        case 3
            Xin = [C(:, 7:9), A];
    end

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

end

M = [M, C(:, 10)];


%% Scatterplots to show how much the input data have changed

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

mapping = [1, 1, 1, 2, 2, 2, 3, 3, 3, 4];

for i = 1:10
    subplot(3, 4, subplot_column_indices(3, 4, i));
    x = C(idx, i);
    y = M(idx, mapping(i));
    plot(x, y, '.');
    myAxisEqual(gca);
    R2onPlot(gca, x, y, true);
    title(var_names(i), "Interpreter","none");
end


%% MAPS of the mixture model output for P, E, ΔS, and R 

M = reshape(M, [360, 720, 240, 4]);

figure('units','normalized','outerposition',[0 0 1 1]);
[lat, lon] = cdtgrid(0.5);
mo = 55;
dt = datetime(2000, mo, 1);
mo_str = ISO_month(dt);
for i = 1:4
    subplot(2, 2, i);
    imagescn(lon, lat, M(:, :, mo, i));
    coasts('k');
    mapzoom(8);
    title(catnames(i));
    colorbar;
end

sgtitle("One month of the Mixture Model output, NN #5, " + mo_str);


%% Map the Calibrated EO datasets
C = reshape(C, [360, 720, 240, 10]);

figure('units','normalized','outerposition',[0 0 1 1]);
for i = 1:10
    subplot(3, 4, i);
    imagescn(lon, lat, C(:, :, mo, i));
    coasts('k');
    mapzoom(8);
    title(var_names(i));
    colorbar;
end

sgtitle("One month of the calibrated variables, NN #5, " + mo_str);


%% SAVE 
descr = "NN predictions of the full water cycle from the mixture models. " + ...
    "M is 360 x 720 x 240 x 4";
metadata = filemeta(descr, src);
save("nn_grid_output12-5.mat", "metadata", "C", "M", "-v7.3");

