using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;

namespace LabelerCore.Model.CrossValidation
{
    //K-Folds cross validation iterator
    //Based on the KFold from the scikit-learn library
    //see: https://github.com/scikit-learn/scikit-learn/blob/51a765a/sklearn/cross_validation.py#L254

    public class KFold : IEnumerable
    {
        private int _n;
        private int _nFolds;
        private List<int> _indices;

        public KFold(int n, int nFolds, bool shuffle=false, int? seed=null)
        {
            if (nFolds < 2 || nFolds > n)
            {
                throw new
                ArgumentException("Number of folds should be in [2; n]");
            }
            _n = n;
            _nFolds = nFolds;
            _indices = Enumerable.Range(0, n).ToList();
            if (shuffle)
            {
                Random rand = seed == null ? new Random(): new Random(seed.Value);
                _indices = _indices.OrderBy(x => rand.Next()).ToList();
            }
        }

        public int Length
        {
            get
            {
                return this._nFolds;
            }
        }

        //Yields tuple<train list indices, test list indices>
        public IEnumerator GetEnumerator()
        {
            int[] foldSizes = new int[_nFolds];
            for (int i = 0; i < foldSizes.Length; i++)
                foldSizes[i] = _n / _nFolds;
            for (int i = 0; i < _n % _nFolds; i++)
                foldSizes[i] += 1;
            int current = 0;
            foreach (int foldSize in foldSizes)
            {
                int start = current;
                int stop = current + foldSize;
                List<int> test = _indices.Skip(start).Take(stop - start).ToList();
                List<int> train = _indices.Where(x => !test.Contains(x)).ToList();
                yield return new Tuple<List<int>, List<int>>(train, test);
                current = stop;
            }
        }
    }


}
