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

namespace LabelerCore.Model.CrossValidation
{
    //Stratified K-Folds cross validation iterator
    //Based on the Stratified  KFold from the scikit-learn library
    //see
    //https://github.com/scikit-learn/scikit-learn/blob/51a765a/sklearn/cross_validation.py#L447
    //Here is straghtforward port of the python code
    public class StratifiedKFold : IEnumerable
    {
        private int[] _testFolds;
        private int _nFolds;

        public StratifiedKFold(int[] y, int nFolds, bool shuffle = false, int? seed = null)
        {
            if (nFolds < 2 || nFolds > y.Length)
            {
                throw new
                ArgumentException("Number of folds should be in [2; y.Length]");
            }
            this._nFolds = nFolds;
            //unique_labels, y_inversed = np.unique(y, return_inverse=True)
            List<int> uniqueLabels = new HashSet<int>(y).ToList();
            uniqueLabels.Sort();
            List<int> yInversed = y.Select(x => uniqueLabels.IndexOf(x)).ToList();
            // label_counts = bincount(y_inversed)
            int[] labelCounts = new int[uniqueLabels.Count];
            foreach (int v in yInversed)
                labelCounts[v] += 1;

            //per_label_cvs = [
            //    KFold(max(c, self.n_folds), self.n_folds, shuffle=self.shuffle,
            //            random_state=rng) for c in label_counts]
            List<KFold> perLabelCvs = new List<KFold>();
            foreach (int count in labelCounts)
                perLabelCvs.Add(new KFold(System.Math.Max(count, nFolds), nFolds, shuffle, seed));

            //zip(*per_label_cvs)
            var testSplits = new List<List<List<int>>>();
            for (int n = 0; n < perLabelCvs[0].Length; n++)
            {
                testSplits.Add(new List<List<int>>());
                for (int i = 0; i < perLabelCvs.Count; i++)
                {
                    int k = 0;
                    foreach (Tuple<List<int>, List<int>> trainTest in perLabelCvs[i])
                    {
                        if (n == k)
                            testSplits[n].Add(trainTest.Item2);
                        k++;
                    }
                }
            }
            int[] testFolds = new int[y.Length];
            //for test_fold_idx, per_label_splits in enumerate(zip(*per_label_cvs)):
            //note : testSplits = zip(*perLabelCvs)
            int testFoldIdx = 0;
            foreach (List<List<int>> perLabelSplit in testSplits)
            {
                //for label, (_, test_split) in zip(unique_labels, per_label_splits):
                int labelIdx = 0;
                foreach (List<int> testSplit in perLabelSplit)
                {
                    int label = uniqueLabels[labelIdx];
                    // label_test_folds = test_folds[y == label]
                    var labelTestFolds = y.Select((x, idx) => new KeyValuePair<double, int>(x, idx))
                                          .Where(x => x.Key == label)
                                          .Select(x => testFolds[x.Value])
                                          .ToList();
                    //test_split = test_split[test_split < len(label_test_folds)]
                    var testSplit_ = testSplit.Where(x => x < labelTestFolds.Count)
                                                 .ToList();
                    //label_test_folds[test_split] = test_fold_idx
                    for (int i = 0; i < testSplit_.Count; i++)
                        labelTestFolds[testSplit[i]] = testFoldIdx;
                    //test_folds[y == label] = label_test_folds
                    int k = 0;
                    for (int i = 0; i < y.Length; i++)
                    {
                        if (y[i] == label)
                        {
                            testFolds[i] = labelTestFolds[k];
                            k += 1;
                        }
                    }
                    labelIdx++;
                }
                testFoldIdx++;

            }
            this._testFolds = testFolds;
        }

        //Yields tuple<train list indices, test list indices>
        public IEnumerator GetEnumerator()
        {
            for (int i = 0; i < this._nFolds; i++)
            {
                //for i in range(self.n_folds):
                //    yield self.test_folds == i
                List<int> test = this._testFolds
                                         .Select((x, j) => new KeyValuePair<double, int>(x, j))
                                         .Where(x => x.Key == i)
                                         .Select(x => x.Value)
                                         .ToList();

                List<int> train = Enumerable.Range(0, this._testFolds.Length)
                                            .Where(x => !test.Contains(x))
                                            .ToList();

                yield return new Tuple<List<int>, List<int>>(train, test);
            }

        }

    }
}
