function [X, A, Y, basins_matrix, months_matrix] = get_nn_inputs12_5(category, partition_number)
% Returns a set of matrices that are my inputs and targets for expt 12-5 neural networks. 
% VERSION 5 
%   * Longer list of ANCILLARIES
%   * The variables are NOT NORMALIZED
%
% Usage:
%  [X, A, Y, basins, months] = get_nn_inputs12_5(category, partition_number)
% 
%   category is one of "train", "validate", or "all"
%   partition_number is either excluded, or an integer from 1 - 20
%
% Returns:
%   X (n_months × n_basins × 10) - "raw" EO data, basin means
%   A (n_months × n_basins × 6 ) - ancillary variables, native units, NOT normalized
%   Y (n_months × n_basins × 4 ) - OI solution
%   basins(n_months × n_basins) - the basinIDs (1...1698) corresponding to the observations in X, A, Y
%   months(n_months × n_basins) - the months (1...240) corresponding to the observations in X, A, Y
%
% where n_basins depends on whether we choose category = "train", "validate" or "all"
%
% Columns in X: 
% (1) P_GPCP, (2) P_GPM_IMERG, (3) P_MSWEP, 
% (4) E_Gleam_a, (5) E_Gleam_b, (6) E_ERA5
% (7) S_CSR, (8) S_GSFC, (9) S_JPL, 
% (10) R

if nargin < 1
    error("First argument should be 'train', 'validate', or 'all'");
end

if nargin < 2
    partition_number = 0;
end

n_months = 240; 
n_basins = 1698;  % Note we may adjust this when we load the partition data

% Create a couple of "helper vectors" for debugging & verification
% These just help to remind me where an observation belongs in terms of time and place
basins_matrix = repelem(1:n_basins, n_months, 1);   % basins is the numeric ID for each basin 1, 2, ... , 2056
months_matrix = repelem(1:n_months, n_basins, 1)';  % months is 1, 2, ... 240 (not 1-12!)

% Vectorize our helper variables
% basins_vector = basins_matrix(:);
% months_vector = months_matrix(:);
% gages_vector  = gages_matrix(:);

%% Create matrix of PREDICTORS or INPUTS, n x 10
load('data/eo_basin_means.mat', 'V');
% This puts V, 240 x 1698 x 10 into the workspace

% Let X be the Inputs to our NN, or the "Predictors"
% reshape it into an n x 10 matrix
% Columns represent: P1, P2, P3, P4, E1, E2, E3, S1, S2, S3, R
X = V(241:480, :, :);

% Get the NN TARGETS,  n x  4, saved from the OI routine run previously. 
OI_fname = "data/OI_soln.mat";  % Which OI solution are we using? 
load(OI_fname, 'XOI')
Y = XOI;
clear XOI V


%% Ancillaries

% Load the first set of ANCILLARY variable basin means
load("data/ancillaries12-5.mat", 'A');

% Recall that A is 240 x 1698 x 12
% and that is exactly the way we want it

% Drop the 20% of validation basin data from this matrix. 
if partition_number == 0
    load("data\partitions.mat");  
    % 'training_basins', 'validation_basins', 'n_train', 'n_validate');
else
    fname = sprintf("data/partitions20/partitions%02d.mat", partition_number);
    load(fname);
end

if category == "train"
    X = X(:, training_basins, :);
    A = A(:, training_basins, :);
    Y = Y(:, training_basins, :);
    basins_matrix = basins_matrix(:, training_basins);
    months_matrix = months_matrix(:, training_basins);
    n_basins = n_train; % Update the number of basins (from 1698 to 1358)

elseif category == "validate"
    X = X(:, validation_basins, :);
    A = A(:, validation_basins, :);
    Y = Y(:, validation_basins, :);
    basins_matrix = basins_matrix(:, validation_basins);
    months_matrix = months_matrix(:, validation_basins);
    n_basins = n_validate; % Update the number of basins (from 1696 to 340)
elseif category == "all"
    % No adjustments needed.
    disp("Returning all")
else
    error("Invalid category. Enter 'train', 'validate', or 'all'");
end



end