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

namespace LabelerCore.Model.Classifiers
{
    public class KMeans: IClassifier
    {
        //params
        private int _nClusters;
        private int _nInit;
        private int _maxIter;
        private double _inertia;
        private double _tolerance;

        //distance
        //distance type to use
        protected Math.Metrics.Metric _metric;
        protected Math.Metrics.Distance _distance;

        //n_clusters x m_features
        private double[][] _centroids;

        //optimization objective
        public double Inertia { get { return _inertia; } }

        
        public delegate void ProgressChangedHandler(double progress);
        public event ProgressChangedHandler ProgressChanged;

        private int _iter;


        public KMeans(int nClusters = 3, int nInit = 10, int maxIter = 300, double tolerance=1E-6,
                      Math.Metrics.Metric metric = Math.Metrics.Metric.SquaredEuclidean)
        {
            _nClusters = nClusters;
            _nInit = nInit;
            _maxIter = maxIter;
            _metric = metric;
            _tolerance = tolerance;
            _distance = Math.Metrics.SelectDistance(metric);
        }

        //return clusters, key - cluster _idx, value List of samples
        private Dictionary<int, List<double[]>> _assignClusters(double[][] X)
        {
            Dictionary<int, List<double[]>> clusters = new Dictionary<int, List<double[]>>(_nClusters);
            for (int i = 0; i < _nClusters; i++)
                clusters.Add(i, new List<double[]>());

            int[] predicted = Predict(X);
            for (int i = 0; i < predicted.Length; i++)
                clusters[predicted[i]].Add(X[i]);
            return clusters;
        }

        //sets and returns _centroids, sets _inertia
        private double[][] _fitOnce(double[][] X, int[] y = null)
        {
            double[] min = Math.Stats.Min(X, axis: 0);
            double[] max = Math.Stats.Max(X, axis: 0);

            Dictionary<int, List<double[]>> clusters;

            //compute 
            //randomly initialize cluster centroids
            _centroids = new double[_nClusters][];
            Random rand = new Random();
            for (int i = 0; i < _nClusters; i++)
            {
                _centroids[i] = new double[X[0].Length];
                for (int j = 0; j < X[0].Length; j++)
                    _centroids[i][j] = rand.NextDouble() * (max[j] - min[j]) + min[j];
            }

            double inertia = double.PositiveInfinity;
            _inertia = double.PositiveInfinity;
            for (int it = 0; it < _maxIter; it++)
            {
                //step 1 - assign clusters
                clusters = _assignClusters(X);
                inertia = _cmpInertia(clusters);
                if (System.Math.Abs(inertia - _inertia) < _tolerance)
                    break;
                _inertia = inertia;

                //step 2 - recompute centroids
                for (int i = 0; i < _nClusters; i++)
                {
                    if (clusters[i].Count > 0)
                        _centroids[i] = Math.Stats.Mean(clusters[i].ToArray(), axis: 0);
                    else
                    {
                        for (int j = 0; j < X[0].Length; j++)
                            _centroids[i][j] = rand.NextDouble() * (max[j] - min[j]) + min[j];
                    }
                }
                _iter++;
                if (ProgressChanged != null)
                    ProgressChanged(_iter * 100 / (_maxIter * _nInit));
            }

            //assign clusters
            clusters = _assignClusters(X);

            return _centroids;
        }

        private double _cmpInertia(Dictionary<int, List<double[]>> clusters)
        {
            //compute inertia (sum of all clusters inertia)
            double inertia = 0.0;
            foreach (int i in clusters.Keys)
                if (clusters[i].Count > 0)
                    inertia += Math.Stats.SumOfDists(clusters[i].ToArray(), centroid: _centroids[i]);
            return inertia;
        }


        public void Fit(double[][] X, int[] y = null)
        {
            List<Tuple<double, double[][]>> centroidsLs = new List<Tuple<double, double[][]>>(_nInit);
            _iter = 0;
            //run kMeans nInit times
            for(int i = 0; i < _nInit; i++)
            {
                _fitOnce(X, y);
                centroidsLs.Add(new Tuple<double, double[][]>(_inertia, _centroids));
            }
            //choose best attempt
            _inertia = centroidsLs.OrderBy(x => x.Item1).First().Item1;
            _centroids = centroidsLs.OrderBy(x => x.Item1).First().Item2;
        }

        public int[] Predict(double[][] X)
        {
            return X.AsParallel().AsOrdered().Select(x => Predict(x)).ToArray();
        }

        public int Predict(double[] x)
        {
            //choose nearest centroid
            return Math.Stats.ArgKSmallest(_centroids.Select(c => _distance(c, x)).ToArray(), 1)[0];
        }

        public object Clone()
        {
            return this.MemberwiseClone();
        }
    }
}
