/**
   *  @brief Robust estimation library
   *
   *  @author  Andrey Kudryavtsev (avkudr.github.io)
   *  @date    01/03/2018
   *  @version 1.0
   */
  
  #ifndef ROBUST_ESTIMATOR_H
  #define ROBUST_ESTIMATOR_H
  
  #include <vector>
  #include <random>
  #include <numeric>
  #include <functional>
  #include <algorithm>
  #include <iostream>
  #include <iterator>
  
  namespace robest {
  
  class EstimationProblem{
  
  public:
      // Functions to overload in your class:
      virtual double estimErrorForSample(int i) = 0;
      virtual   void estimModelFromSamples(std::vector<int> samplesIdx) = 0;
      virtual    int getTotalNbSamples() const = 0;
  
      int getNbParams()     const{return nbParams;}
      int getNbMinSamples() const{return nbMinSamples;}
  
  protected:
      void setNbParams(int i)    {nbParams     = i;}
      void setNbMinSamples(int i){nbMinSamples = i;}
  
  private:
      int nbParams = -1;
      int nbMinSamples = -1;
  };
  
  static std::random_device rd;  // random device engine, usually based on /dev/random on UNIX-like systems
  static std::mt19937 rng(rd()); // initialize Mersennes' twister using rd to generate the seed
  //static std::mt19937 rng((unsigned int) - 1);
  
  // Base class for Robust Estimators
  class AbstractEstimator
  {
  public:
      // Generate X !different! random numbers from 0 to N-1
      // X - minimal number of samples
      // N - total number of samples
      // generated numbers are indices of data points
  
      std::vector<int> randomSampleIdx(){
  
          int minNbSamples = problem->getNbMinSamples();
          int totalNbSamples = problem->getTotalNbSamples();
  
          std::vector<int> allIdx(totalNbSamples);
          std::iota (std::begin(allIdx), std::end(allIdx), 0); // Fill with 0, 1, ..., totalNbSamples-1.
  
          // shuffle the elements of allIdx : Fisher–Yates shuffle
          for (int i = 0; i < minNbSamples; i++){
              std::uniform_int_distribution<int> dist(0,totalNbSamples-i-1);
              int randInt = dist(rng);
              std::swap(allIdx[totalNbSamples-i-1],allIdx[randInt]);
          }
  
          //take last <minNbSamples> elements
          std::vector<int> idx( allIdx.end() - minNbSamples, allIdx.end());
  //        std::cout << "[";
  //        for (auto i : idx) std::cout << i << ", ";
  //        std::cout << "]
  ";
          return idx;
      }
  
      double getInliersFraction() const {return inliersFraction;}
      std::vector<int> getInliersIndices() const {return inliersIdx;}
  
  protected:
      void getInliers(double thres){
          int totalNbSamples = problem->getTotalNbSamples();
          inliersIdx.clear();
          for(int j = 0; j < totalNbSamples; j++){
              double error = problem->estimErrorForSample(j);
              error = error*error;
              if (error < thres){
                  inliersIdx.push_back(j);
              }
          }
          this->inliersFraction = (double)(inliersIdx.size()) / (double)(totalNbSamples);
      }
  
      EstimationProblem * problem;
      std::vector<int> bestIdxSet;
      std::vector<int> inliersIdx;
      double inliersFraction = -1.0;
  };
  
  class RANSAC : public AbstractEstimator
  {
  public:
      RANSAC(){
  
      }
  
      void solve(EstimationProblem * pb, double thres = 0.1, int nbIter = 10000){
          problem = pb;
  
          int totalNbSamples = problem->getTotalNbSamples();
          for (int i = 0; i < nbIter; i++){
              std::vector<int> indices = randomSampleIdx();
  
              problem->estimModelFromSamples(indices);
  
              //getInliersNb
              int nbInliers = 0;
              for(int j = 0; j < totalNbSamples; j++){
                  double error = problem->estimErrorForSample(j);
                  error = error*error;
                  if (error < thres){
                      nbInliers++;
                  }
              }
  
              double inliersFraction = (double)(nbInliers) / (double)(problem->getTotalNbSamples());
              if (inliersFraction > this->inliersFraction){
                  this->inliersFraction = inliersFraction;
                  this->bestIdxSet = indices;
              }
          }
  
          problem->estimModelFromSamples(bestIdxSet);
          getInliers(thres);
          problem->estimModelFromSamples(inliersIdx);
      }
  };
  
  class LMedS : public AbstractEstimator
  {
  public:
      LMedS(){
          med = 1000000.0;
      }
  
      template <typename T>
      double median(std::vector<T> & v){
          std::sort(v.begin(), v.end());
          if (v.size() % 2 == 0){
              return (double)(v[v.size()/2-1] + v[v.size()/2]) / 2;
          }else{
              return v[v.size()/2];
          }
      }
  
      void solve(EstimationProblem * pb, double thres = 0.1, int nbIter = 1000){
          problem = pb;
          int totalNbSamples = problem->getTotalNbSamples();
  
          for (int i = 0; i < nbIter; i++){
              std::vector<int> indices = randomSampleIdx();
  
              problem->estimModelFromSamples(indices);
  
              std::vector<double> errorsVec(totalNbSamples);
              for(int j = 0; j < problem->getTotalNbSamples(); j++){
                  double error = problem->estimErrorForSample(j);
                  //errorsVec[j] = error; // error must be squared!
                  errorsVec[j] = error*error; // error must be squared!
              }
              double med = median(errorsVec);
              if (med < this->med){
                  this->med = med;
                  this->bestIdxSet = indices;
              }
          }
  
          problem->estimModelFromSamples(bestIdxSet);
          getInliers(thres);
          problem->estimModelFromSamples(inliersIdx);
      }
  
  private:
      double med;
  };
  
  }
  #endif // ROBUST_ESTIMATOR_H