function G = get_basin_data(category, partition_number)
% Returns some metadata about my 1696 basins for expt 12-4
% Usage: 
%    G = get_basin_data("train")
%    G = get_basin_data("validate")
%    G = get_basin_data("all")
% Returns G, a structure array with G.lat, G.on, and G.areas
% containing the centroid coordinates and the area in km² for my basins

if nargin < 2
    partition_number = 0;
end

% Get the metadata
fname = "data/basin_masks.mat";
load(fname, "centr_lat", "centr_lon", "areas");

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

if category == "train"
    indices = training_basins;
    n_basins = n_train;
elseif category == "validate"
    indices = validation_basins;
    n_basins = n_validate; 
elseif category == "all"
    n_basins = n_train + n_validate;
    indices = 1:n_basins; 
end


G.lat = centr_lat(indices);
G.lon = centr_lon(indices);
G.area = areas(indices);
G.n_basins = n_basins;

end