## Copyright (C) 2014 Nir Krakauer
##
## This file is part of the statistics package for GNU Octave.
##
## This program is free software; you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation; either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program; If not, see <http://www.gnu.org/licenses/>.

## -*- texinfo -*-
## @deftypefn  {statistics} {@var{C} =} cvpartition (@var{X}, [@var{partition_type}, [@var{k}]])
##
## Create a partition object for cross validation.
##
## @var{X} may be a positive integer, interpreted as the number of values
## @var{n} to partition, or a vector of length @var{n} containing class
## designations for the elements, in which case the partitioning types
## @var{KFold} and @var{HoldOut} attempt to ensure each partition represents the
## classes proportionately.
##
## @var{partition_type} must be one of the following:
##
## @table @asis
## @item @samp{KFold}
## Divide set into @var{k} equal-size subsets (this is the default, with
## @var{k}=10).
## @item @samp{HoldOut}
## Divide set into two subsets, "training" and "validation".  If @var{k} is a
## fraction, that is the fraction of values put in the validation subset; if it
## is a positive integer, that is the number of values in the validation subset
## (by default @var{k}=0.1).
## @item @samp{LeaveOut}
## Leave-one-out partition (each element is placed in its own subset).
## @item @samp{resubstitution}
## Training and validation subsets that both contain all the original elements.
## @item @samp{Given}
## Subset indices are as given in @var{X}.
## @end table
##
## The following fields are defined for the @samp{cvpartition} class:
##
## @table @asis
## @item @samp{classes}
## Class designations for the elements.
## @item @samp{inds}
## Subset indices for the elements.
## @item @samp{n_classes}
## Number of different classes.
## @item @samp{NumObservations}
## @var{n}, number of elements in data set.
## @item @samp{NumTestSets}
## Number of testing subsets.
## @item @samp{TestSize}
## Number of elements in (each) testing subset.
## @item @samp{TrainSize}
## Number of elements in (each) training subset.
## @item @samp{Type}
## Partition type.
## @end table
##
## @seealso{crossval, @@cvpartition/display}
## @end deftypefn

function C = cvpartition (X, partition_type = "KFold", k = [])

  if (nargin < 1 || nargin > 3 || !isvector(X))
    print_usage ();
  endif

  if (isscalar (X))
    n = X;
    n_classes = 1;
  else
    n = numel (X);
  endif

  switch (tolower (partition_type))
    case {"kfold", "holdout", "leaveout", "resubstitution", "given"}
    otherwise
      warning ("cvpartition: unrecognized type, using KFold.")
      partition_type = "KFold";
  endswitch

  switch (tolower (partition_type))
    case {"kfold", "holdout", "given"}
      if (! isscalar (X))
        [y, ~, j] = unique (X(:));
        n_per_class = accumarray (j, 1);
        n_classes = numel (n_per_class);
      endif
  endswitch

  C = struct ("classes", [], "inds", [], "n_classes", [], "NumObservations", ...
            [], "NumTestSets", [], "TestSize", [], "TrainSize", [], "Type", []);
  ## The non-Matlab fields classes, inds, n_classes
  ## are only useful for some methods

  switch (tolower (partition_type))
    case "kfold"
      if (isempty (k))
        k = 10;
      endif
      if (n_classes == 1)
        inds = floor((0:(n-1))' * (k / n)) + 1;
      else
        inds = nan(n, 1);
        for i = 1:n_classes
          ## Alternate ordering over classes so that
          ## the subsets are more nearly the same size
          if (mod (i, 2))
            inds(j == i) = floor((0:(n_per_class(i)-1))' * ...
                                 (k / n_per_class(i))) + 1;
          else
            inds(j == i) = floor(((n_per_class(i)-1):-1:0)' * ...
                                 (k / n_per_class(i))) + 1;
          endif
        endfor
      endif
      C.inds = inds;
      C.NumTestSets = k;
      [~, ~, jj] = unique (inds);
      n_per_subset = accumarray (jj, 1);
      C.TrainSize = n - n_per_subset;
      C.TestSize = n_per_subset;
    case "given"
      C.inds = j;
      C.NumTestSets = n_classes;
      C.TrainSize = n - n_per_class;
      C.TestSize = n_per_class;
    case "holdout"
      if (isempty (k))
        k = 0.1;
      endif
      if (k < 1)
        f = k;              # target fraction to sample
        k = round (k * n);  # number of samples
      else
        f = k / n;
      endif
      inds = zeros (n, 1, "logical");
      if (n_classes == 1)
        inds(randsample(n, k)) = true;  # indices for test set
      else        # sample from each class
        k_check = 0;
        for i = 1:n_classes
          ki = round(f*n_per_class(i));
          inds(find(j == i)(randsample(n_per_class(i), ki))) = true;
          k_check += ki;
        endfor
        if (k_check < k)       # add random elements to test set to make it k
          inds(find(!inds)(randsample(n - k_check, k - k_check))) = true;
        elseif (k_check > k)   # remove random elements from test set
          inds(find(inds)(randsample(k_check, k_check - k))) = false;
        endif
        C.classes = j;
      endif
      C.n_classes = n_classes;
      C.TrainSize = n - k;
      C.TestSize = k;
      C.NumTestSets = 1;
      C.inds = inds;
    case "leaveout"
      C.TrainSize = ones (n, 1);
      C.TestSize = (n-1)  * ones (n, 1);
      C.NumTestSets = n;
    case "resubstitution"
      C.TrainSize = C.TestSize = n;
      C.NumTestSets = 1;
  endswitch

  C.NumObservations = n;
  C.Type = tolower (partition_type);

  C = class (C, "cvpartition");

endfunction


%!demo
%! ## Partition with Fisher iris dataset (n = 150)
%! ## Stratified by species
%! load fisheriris
%! y = species;
%! ## 10-fold cross-validation partition
%! c = cvpartition (species, 'KFold', 10)
%! ## leave-10-out partition
%! c1 = cvpartition (species, 'HoldOut', 10)
%! idx1 = test (c, 2);
%! idx2 = training (c, 2);
%! ## another leave-10-out partition
%! c2 = repartition (c1)

%!test
%! C = cvpartition (ones (10, 1));
%! assert (isa (C, "cvpartition"), true);
%!test
%! C = cvpartition (ones (10, 1), "KFold", 5);
%! assert (get (C, "NumObservations"), 10);
%! assert (get (C, "NumTestSets"), 5);
%! assert (get (C, "TrainSize"), ones(5,1) * 8);
%! assert (get (C, "TestSize"), ones (5,1) * 2);
%! assert (get (C, "inds"), [1 1 2 2 3 3 4 4 5 5]');
%! assert (get (C, "Type"), "kfold");
%!test
%! C = cvpartition (ones (10, 1), "KFold", 2);
%! assert (get (C, "NumObservations"), 10);
%! assert (get (C, "NumTestSets"), 2);
%! assert (get (C, "TrainSize"), [5; 5]);
%! assert (get (C, "TestSize"), [5; 5]);
%! assert (get (C, "inds"), [1 1 1 1 1 2 2 2 2 2]');
%! assert (get (C, "Type"), "kfold");
%!test
%! C = cvpartition (ones (10, 1), "HoldOut", 5);
%! assert (get (C, "NumObservations"), 10);
%! assert (get (C, "NumTestSets"), 1);
%! assert (get (C, "TrainSize"), 5);
%! assert (get (C, "TestSize"), 5);
%! assert (class (get (C, "inds")), "logical");
%! assert (length (get (C, "inds")), 10);
%! assert (get (C, "Type"), "holdout");
%!test
%! C = cvpartition ([1 2 3 4 5 6 7 8 9 10], "LeaveOut", 5);
%! assert (get (C, "NumObservations"), 10);
%! assert (get (C, "NumTestSets"), 10);
%! assert (get (C, "TrainSize"), ones (10, 1));
%! assert (get (C, "TestSize"), ones (10, 1) * 9);
%! assert (get (C, "inds"), []);
%! assert (get (C, "Type"), "leaveout");
%!test
%! C = cvpartition ([1 2 3 4 5 6 7 8 9 10], "resubstitution", 5);
%! assert (get (C, "NumObservations"), 10);
%! assert (get (C, "NumTestSets"), 1);
%! assert (get (C, "TrainSize"), 10);
%! assert (get (C, "TestSize"), 10);
%! assert (get (C, "inds"), []);
%! assert (get (C, "Type"), "resubstitution");
%!test
%! C = cvpartition ([1 2 3 4 5 6 7 8 9 10], "Given", 2);
%! assert (get (C, "NumObservations"), 10);
%! assert (get (C, "NumTestSets"), 10);
%! assert (get (C, "TrainSize"), ones (10, 1) * 9);
%! assert (get (C, "TestSize"), ones (10, 1));
%! assert (get (C, "inds"), [1:10]');
%! assert (get (C, "Type"), "given");

%!warning<cvpartition: unrecognized type, using KFold.> ...
%! C = cvpartition ([1 2 3 4 5 6 7 8 9 10], "some", 2);

