#include "SplineEvaluation.hpp"
#include "boost-hdrs/remez.hpp"
#include <boost/math/tools/polynomial.hpp>
#include <iostream>
#include <fstream>
#include <algorithm>
#include <cstring>
#include <cmath>
#include <vector>
#include <iomanip>
#include <cstdlib>
#include <map>
#include <chrono>
#include <iomanip>
#include <cassert>
#include <unordered_map>
#include <cassert>
#include <random>
#include <algorithm>
#include <immintrin.h>
//---------------------------------------------------------------------------
using namespace std;
//---------------------------------------------------------------------------
typedef std::pair<double,double> Coord;
typedef std::pair<double, unsigned> segCoord;
typedef btree::btree_map<double, double> btree_map;
typedef btree::btree_map<double, unsigned> btree_map_segments;
//---------------------------------------------------------------------------
// modes for splines
#define WORK_ON_SPLINE_MODE 1
#define SSE_SPLIT_MODE 0
#define JUST_FIT_MODE 1
#define LOWER_BOUND_MODE 2
#define MAP_MODE 3
#define BTREE_MODE 4
#define BTREE_SPLINE_MODE 5
#define TEST_MODE_ACTIVATED 1
//---------------------------------------------------------------------------
const double precision = std::numeric_limits<double>::epsilon();
const double MIN_ANGLE = 1e-16; // 10 ^ (-precisionOfDouble)
const double MAX_RUNS = 2;
//---------------------------------------------------------------------------
double globalDiff;
double powers[2]; // for sse or avx
std::vector<segCoord> globalSpline; // we need it to define the function for the Remez algorithm
static double interpolate(const vector<Coord>& spline, double pos);
static std::vector<Coord> preciseApproximation(const std::vector<Coord>& function, unsigned desiredSize);
static pair<double, double> computeOverallPrecision(const std::vector<Coord>& func, const std::vector<Coord>& reduced);
//---------------------------------------------------------------------------
SplineEvaluation::SplineEvaluation() {
}
//---------------------------------------------------------------------------
SplineEvaluation::SplineEvaluation(const std::vector<Coord>& function, unsigned desiredSize, const unsigned useSplineMode, const unsigned searchMode, const unsigned testMode) {
    if (useSplineMode == WORK_ON_SPLINE_MODE) {
        // Compute the spline
        this->spline = preciseApproximation(function, min(desiredSize, this->limitSplineSize));
        // Get the errors
        pair<double, double> errors = computeOverallPrecision(function, this->spline);
        cerr << "maxError = " << errors.first << " avgError = " << errors.second << endl;
#if 0
        cerr << "Print spline : " << endl;
        for (auto coord: this->spline)
            cerr << "(" << coord.first << "," << coord.second << ")" << endl;
#endif
        switch (searchMode) {
            case SSE_SPLIT_MODE : {
                cerr << "enter Split Mode" << endl;
                // separate the spline into countOfBuckets smaller splines
                splitSegmentSpline();
                this->polyBuckets = new long double*[countOfBuckets];

                // fit each spline
                for (unsigned index = 0; index < this->countOfBuckets; ++index) {
                    globalSpline = this->splineBuckets[index];
                    changeIntoArray(&this->polyBuckets[index], fitSpline(this->splineBuckets[index]));
                }
                cerr << "quit Split Mode" << endl;
                break;
            }
            case JUST_FIT_MODE : {
                cerr << "enter just fit mode" << endl;
                // get the segment spline
                transformSpline();

                // now fit it
                globalSpline = this->segmentSpline;
                changeIntoArray(&this->poly, fitSpline(this->segmentSpline));
                cerr << "quit just fit mode" << endl;
                break;
            }
            case BTREE_SPLINE_MODE : {
                cerr << "enter btree spline mode" << endl;
                transformSpline();
                this->btreeSegments = new btree_map_segments;
                for (auto coord: this->segmentSpline) {
                    this->btreeSegments->insert(coord);
                }
                cerr << "done with btree spline mode" << endl;
                break;
            }
        }
        // Pad the spline with dummy values
        convertValues();
    } else {
        switch (searchMode) {
            case BTREE_MODE : {
                cerr << "Enter values in btree" << endl;
                this->btreeMap = new btree_map;
                for (auto coord: function) {
                    this->btreeMap->insert(coord);
                }
                cerr << "Done with btree" << endl;
                if (testMode != TEST_MODE_ACTIVATED)
                    break;
            }
            case MAP_MODE : {
                cerr << "Enter the values into map" << endl;
                for (auto coord: function) {
                    this->coordinates[coord.first] = coord.second;
                }
                cerr << "Done with the map" << endl;
                break;
            }
        }
    }
}
//---------------------------------------------------------------------------
SplineEvaluation::~SplineEvaluation()
// deconstructor
{
    // TODO: save the modes, in order to know what to free
    // there are only 4 pointers, which have to be deleted (spl)
    // There may be a problem with "new" on vector.
}
//---------------------------------------------------------------------------
void SplineEvaluation::splitSegmentSpline()
// split the segment spline into 5 smaller splines
{
    // alloc the smaller splines
    this->splineBuckets = new std::vector<segCoord>[countOfBuckets];

    // Get the index of the first split
    unsigned splittedPos = this->spline.size() / countOfBuckets;
    unsigned ptr = splittedPos, segmentCount = 0, splineSize = this->spline.size();
    for (unsigned splineOrder = 0; splineOrder < countOfBuckets; ++splineOrder) {
        // Construct the next smaller spline
        for (unsigned index = ptr - splittedPos; (index < ptr) && (index < splineSize); ++index) {
            this->splineBuckets[splineOrder].push_back(make_pair(this->spline[index].first, segmentCount++));
        }
        // add the next splitted position
        if (splineOrder != this->countOfBuckets - 1)
            this->splittedPositions[splineOrder] = this->spline[segmentCount - 1].first;
        // go to the next spline
        ptr += splittedPos;
    }
    // Load the ymm-register with the limits of the splines
    saveLimits = _mm_loadu_pd(this->splittedPositions);
}
//---------------------------------------------------------------------------
SplineEvaluation::SplineEvaluation(unsigned size, const double* x, const double* y) {
    this->spline.resize(size);
    for (unsigned index = 0; index < size; ++index)
        this->spline[index] = make_pair(x[index], y[index]);
}
//---------------------------------------------------------------------------
void SplineEvaluation::save(std::ofstream& file_out)
// save in binary file the spline and the polynomial
{
    file_out.write(reinterpret_cast<char*>(this->spline.data()),this->spline.size()*sizeof(pair<double,double>));
    file_out.write(reinterpret_cast<char*>(this->poly),(polyGrade + 1) * sizeof(long double));
}
//---------------------------------------------------------------------------
void SplineEvaluation::load(std::ifstream& file_in)
// load from binary file the spline and the polynomial
{
    auto pos = file_in.tellg();
    file_in.seekg( 0, ios::end );
    auto size=file_in.tellg()-pos;
    file_in.seekg(0,ios::beg);

    unsigned splineSize = ((size / sizeof(double)) - (polyGrade + 1)) / 2;

    this->spline.resize(splineSize);
    this->poly = new long double[polyGrade + 1];

    file_in.read(reinterpret_cast<char*>(this->spline.data()),splineSize * sizeof(pair<double,double>));
    file_in.read(reinterpret_cast<char*>(this->poly),(polyGrade + 1)*sizeof(long double));
    convertValues();
}
//---------------------------------------------------------------------------
void SplineEvaluation::transformSpline()
// marks every x with the segment where it is to be found. This is called from now on a "segment Spline"
{
    unsigned segmentCount = 0;
    for (auto elem: this->spline)
        this->segmentSpline.push_back(make_pair(elem.first, segmentCount++));
}
//---------------------------------------------------------------------------
unsigned SplineEvaluation::size() const {
    return this->spline.size();
}
//---------------------------------------------------------------------------
static double square(double x) {
    return x * x;
}
//---------------------------------------------------------------------------
static bool greaterThan(double x, double y) {
    return (x - y) > precision;
}
//---------------------------------------------------------------------------
static bool equals(double x, double y) {
    return fabs(x - y) < precision;
}
//---------------------------------------------------------------------------
static double distance(Coord a, Coord b)
// returns euclidian distance between both points
{
    return sqrt(square(a.first - b.first) + square(a.second - b.second));
}
//---------------------------------------------------------------------------
static double cosinus(Coord u, Coord v, Coord w)
// computes cos(angle) between (u, v, w) with v in middle
{
    double a = distance(u, w), b = distance(u, v), c = distance(v, w);
    return (square(a) - square(b) - square(c)) / (2.0 * b * c);
}
//---------------------------------------------------------------------------
static double computeAngle(Coord u, Coord v, Coord w)
// computes angle between (u, v, w) using cosinus
{
    return 1.0 - cosinus(u, v, w);
}
//---------------------------------------------------------------------------
static double determinant(Coord u, Coord v, Coord w)
// computes the determinant, which represents the spanned area.
{
    return (v.first - u.first) * (w.second - u.second) - (v.second - u.second) * (w.first - u.first);
}
//---------------------------------------------------------------------------
static int orientation(Coord u, Coord v, Coord w)
// computes in which direction w is placed in respect to (u, v)
{
    double det = determinant(u, v, w);
    if (det < -precision)
        return -1;
    else if (det > precision)
        return 1;
    else
        return 0;
}
//---------------------------------------------------------------------------
static pair<double, double> computeOverallPrecision(const std::vector<Coord>& func, const std::vector<Coord>& reduced)
// returns (maxError, avgError)
{
    // compute errors
    double sumError=0, maxError=0, relError = 0;
    for (auto e: func) {
        double estimate=interpolate(reduced, e.first);
        double real=e.second;
        double diff=estimate-real;
        if (diff<0)
            diff=-diff;
        if (diff>maxError)
            maxError=diff;
        sumError+=diff;
        // relError += (!real) ? diff : diff / real;
    }
    // relError /= func.size();
    sumError /= func.size();
    return make_pair(maxError, sumError);
}
#if 0
//---------------------------------------------------------------------------
static double inflexionAngleOfFunction(const std::vector<Coord>& data)
// get an angle which creates small inflexions
{
    if (data.size() == 0) {
        cerr << "invalid data!" << endl;
        assert(0);
    }
    if (data.size() == 1)
        return 0;
   // Taut the string
   std::vector<Coord>::const_iterator iter = data.begin(), limit = data.end();
   Coord prev = *iter, curr = *(++iter);
   double inflexionAngle = 0, currSum = 0;

    for (++iter; iter != limit; ++iter) {
        double deviation = 0.0;
        int orient = orientation(prev, curr, *iter);
        if (orient)
            deviation = computeAngle(prev, curr, *iter);
        currSum += orient * deviation;

        // update the indexes.
        prev = curr;
        curr = *iter;
    }
    return fabs(currSum) / data.size();
}
#endif
//---------------------------------------------------------------------------
static std::vector<Coord> angleApproximation(const std::vector<Coord>& data, double maxAngle, double maxDistance = 0)
// Compress the functions by storing only the inflections with angle lower than maxAngle
{
   // Add the first point
   std::vector<Coord> result;
   if (data.empty())
      return result;
   result.push_back(data[0]);
   if (data.size()==1)
      return result;

   // Taut the string
   std::vector<Coord>::const_iterator iter = data.begin(), limit = data.end();
   Coord prev = *iter, curr = *(++iter);
   double currSum = 0;

    for (++iter; iter != limit; ++iter) {
        double deviation = 0.0;
        int orient = orientation(prev, curr, *iter);
        if (orient) {
            deviation = computeAngle(prev, curr, *iter);
            currSum += orient * deviation;
        }

        // Check for inflexion.
        if (greaterThan(deviation, maxAngle) || greaterThan(fabs(currSum), maxAngle)) {
            result.push_back(curr);
            currSum = 0;
        }

        // update the indexes.
        prev = curr;
        curr = *iter;
    }
    // Add the last point.
    if (result.back() != *(--data.end()))
        result.push_back(*(--data.end()));
    return result;
}
//---------------------------------------------------------------------------
static int cmpDevs(const Coord& a,const Coord& b,const Coord& c)
   // Compare derivations. -1 if (a->b) is more swallow than (a->c), +1 if it is more steeper
{
   double dx1=b.first-a.first,dx2=c.first-a.first;
   double dy1=b.second-a.second,dy2=c.second-a.second;
   if ((dy1*dx2)<(dy2*dx1)) return -1;
   if ((dy1*dx2)>(dy2*dx1)) return 1;
   return 0;
}
//---------------------------------------------------------------------------
static std::vector<Coord> tautString(const std::vector<Coord>& data, double epsilon)
{
   // Add the first point
   std::vector<Coord> result;
   if (data.empty())
      return result;
   result.push_back(data[0]);
   if (data.size()==1)
      return result;

   // Taut the string
   std::vector<Coord>::const_iterator iter=data.begin(),limit=data.end();
   Coord upperLimit,lowerLimit,last=result.back();

   for (++iter;iter!=limit;++iter) {
      // Add the new bounds
      Coord u = *iter, l = *iter, b=result.back();
      u.second+=epsilon; l.second-=epsilon;

      // Check if we cut the error corridor
      if ((last!=b)&&((cmpDevs(b,upperLimit,*iter)<0)||(cmpDevs(b,lowerLimit,*iter)>0))) {
         result.push_back(last);
         b=last;
        }

      // Update the error margins
      if ((last==b)||(cmpDevs(b,upperLimit,u)>0))
         upperLimit=u;
      if ((last==b)||(cmpDevs(b,lowerLimit,l)<0))
        lowerLimit=l;

      // And remember the current point
      last=*iter;
   }

   // Add the last point
   result.push_back(*(--data.end()));
   return result;
}
//---------------------------------------------------------------------------
static std::vector<Coord> compressFunction(const std::vector<Coord>& function,unsigned desiredSize)
// Compress to the desired size
{
   // Relax a bit to speed up compression
   unsigned maxSize=desiredSize+(desiredSize/100),minSize=desiredSize-(desiredSize/100);

   // Fits?
   if (function.size()<=maxSize)
      return function;

   // No, binary search
   unsigned left=0,right=function.size();
   while (left<right) {
      unsigned middle=(left+right)/2;
      std::vector<Coord> candidate=tautString(function, middle);
      if (candidate.size()<minSize) {
         right=middle;
      } else if (candidate.size()>maxSize) {
         left=middle+1;
      } else {
        return candidate;
    }
   }
   // Final call, this is the best we could get
   return tautString(function, left);
}
//---------------------------------------------------------------------------
static std::vector<Coord> preciseApproximation(const std::vector<Coord>& function, unsigned desiredSize)
// first use a sieve for the smallest angles, then apply the previous algorithm
{
    // Apply the last algorithm to compare its results with the new ones
    std::vector<Coord> candidate = compressFunction(function, desiredSize);
    pair<double, double> previousAlg = computeOverallPrecision(function, candidate);
    double previousMaxError = previousAlg.first, previousAvgError = previousAlg.second;

    cerr << "previous has " << candidate.size() << "points and errors = " << previousMaxError << " and " << previousAvgError << endl;

    // let 0 if you want to search for a better spline (which may take 1-2 minutes)
#if 1
    return candidate;
#endif

#if 0
    // Map the sizes of the splines to pairs of (maxError, avgError)
    unordered_map<int, pair<int, pair<double, double>>> answers;
    answers[candidate.size()] = make_pair(candidate.size(), previousAlg);
#endif
    // it could be modified, for precision reasons
    double currAngle = MIN_ANGLE;
    bool foundBetter = true;
    int runs = MAX_RUNS;

    double smallestAvgError = previousAvgError, smallestMaxError = previousMaxError, cand = -1;
    for (int runs = MAX_RUNS; (runs--) || (foundBetter); currAngle *= 10) {
        foundBetter = false;

        cerr << "progress = " << runs << " with " << currAngle << endl;
        for (unsigned mul = 2; mul < 20; ++mul) {
            double eps = currAngle * mul / 2;

            // Sharp the function then apply the previous algorithm
            auto samplePoints = angleApproximation(function, eps);
            samplePoints = compressFunction(samplePoints, desiredSize);

            pair<double, double> errors = computeOverallPrecision(function, samplePoints);
            double maxError = errors.first, avgError = errors.second;

            cerr << "check with " << eps << " errors = " << maxError << " " << avgError << endl;

            if ((previousAvgError > avgError) && (previousMaxError > maxError)) {
#if 1
                cerr << "count Points : " << samplePoints.size() << endl;
                cerr << "it's better at " << eps << " with " << maxError << " and " << avgError << endl;
                cerr << "rate for avg is " << ((-avgError + previousAvgError) / previousAvgError) * 100 << endl;
                cerr << "rate for max is " << ((-maxError + previousMaxError) / previousMaxError) * 100 << endl;
#endif

#if 0
                // Check the existance of this size
                if (answers.find(samplePoints.size()) == answers.end()) {
                    // not yet in map this size
                    candidate = compressFunction(function, samplePoints.size());
                    pair<double, double> curr = computeOverallPrecision(function, candidate);

                    // And map it to its errors
                    answers[samplePoints.size()] = make_pair(candidate.size(), curr);
                }

                auto curr = answers.find(samplePoints.size())->second;

                cerr << "to " << samplePoints.size() << " found : " << curr.first << " with " << "(" << curr.second.first << "," << curr.second.second << ")" << endl;
#endif
                candidate = compressFunction(function, samplePoints.size());
                pair<double, double> curr = computeOverallPrecision(function, candidate);

                if (curr == previousAlg) {
                    cerr << "Even better than the new one!" << endl;
                    foundBetter = true;

                    // get the best avgError
                    if (smallestAvgError > avgError) {
                        smallestAvgError = avgError;
                        smallestMaxError = maxError;
                        cand = eps;
                    } else if (equals(smallestAvgError, avgError)) {
                        if (smallestMaxError > maxError) {
                            smallestMaxError = maxError;
                            cand = eps;
                        }
                    }
                } else {
                    if (candidate.size() <= samplePoints.size())
                        cerr << "Has been beaten at " << samplePoints.size() << endl;
                }
            }
        }
    }
    if (cand == -1) {
        cerr << "Not found better than that!" << endl;
        return candidate;
    }

    // Get the best through the sieve of angles
    auto ret = angleApproximation(function, cand);
    ret = compressFunction(ret, desiredSize);
    return ret;
}
//---------------------------------------------------------------------------
static unsigned realSegment(const vector<Coord>& spline, double pos)
// DEBUG: get the real segment
{
    if (pos<=spline.front().first)
      return 0;
   if (pos>=spline.back().first)
      return spline.size()-1;
   auto iter=lower_bound(spline.begin(),spline.end(),pos,[](const Coord& a,double b) { return a.first<b; });

   // cerr << setprecision(12) << "for " << pos << "come with " << (iter-1)->first << " " << iter->first << endl;

   if (iter->first == pos)
       return iter - spline.begin();
   return (iter-1) - spline.begin();
}
//---------------------------------------------------------------------------
static double interpolateWithSegments(const vector<segCoord>& spline, double pos)
// Helper function for defineFunction. Evaluates the segment Spline at pos.
{
   if (pos<=spline.front().first)
      return spline.front().second;
   if (pos>=spline.back().first)
      return spline.back().second;
   auto iter=std::lower_bound(spline.begin(),spline.end(),pos,[](const Coord& a,double b) { return a.first<b; });
   if (iter->first==pos)
       return iter->second;

   double dx=(iter+0)->first-(iter-1)->first;
   double dy=(iter+0)->second-(iter-1)->second;

   double ofs=pos-(iter-1)->first;
   return (iter-1)->second+ofs*(dy/dx);
}
//---------------------------------------------------------------------------
static double defineFunction(const double& x)
// defines for the Remez algorithm how the spline works through interpolation
{
    return interpolateWithSegments(globalSpline, x);
}
//---------------------------------------------------------------------------
void SplineEvaluation::changeIntoArray(long double** moveHere, boost::math::tools::polynomial<long double> poly)
// change the STL-vectors into simple arrays. Searching on them is faster.
{
    *moveHere = new long double[this->polyGrade + 1];
    // Copy the polynomial
    for (unsigned index = 0; index < poly.size(); ++index)
        (*moveHere)[index] = poly[index];
}
//---------------------------------------------------------------------------
void SplineEvaluation::convertValues()
// pad values with values lower than the first one, respectively greater than the last one
// this enhances the binary searches, not having the overflow condition anymore
{
    this->length = this->spline.size();
#if 0
    // Compute offset as a power of 2
    this->offset = 1;
    while (this->offset < this->length)
        this->offset <<= 1;
#endif
    this->offset = 0;
    // Alloc enough memory for the inital values and 2 offsets
    this->values = new double[offset + this->length + offset];

    // Pad with values lower than the first one
    for (unsigned index = 0; index < offset; ++index)
        this->values[index] = this->spline[0].first - 1;

    // Fill in with the initial values from the spline x-coordinates
    for (unsigned index = 0; index < this->spline.size(); ++index)
        this->values[offset + index] = this->spline[index].first;

    // Pad with values greater than the last one
    for (unsigned index = 0; index < offset; ++index)
        this->values[offset + index + this->length] = this->spline[this->length - 1].first + 1;
}
//---------------------------------------------------------------------------
boost::math::tools::polynomial<long double> SplineEvaluation::fitSpline(const std::vector<segCoord>& spline)
// use Remez algorithm to fit a polynomial to spline
{
    // Prepare parameters
    unsigned sizeOfNominator = this->polyGrade, sizeOfDenominator = 0;
    double a = spline.front().first, b = spline.back().first;
    bool pin = false, relError = false;
    int skew = 0, workingPrecision = boost::math::tools::digits<long double>() * 2;

    // Activate the algorithm with long double to get a better precision. It's recommanded to keep it so. Otherwise it could be the case
    // that an internal function from the algorithm fails because of the poor precision of double
    boost::math::tools::remez_minimax<long double> remez(defineFunction, sizeOfNominator, sizeOfDenominator, a, b, pin, relError, skew, workingPrecision);
    return remez.numerator();
}
//---------------------------------------------------------------------------
double SplineEvaluation::hornerEvaluate(double& x) const
// evaluate the polynomial with Horner's schema
{
#if 1
    long double ret = this->poly[polyGrade];
    for (unsigned i = polyGrade; i != 0; --i)
        ret = ret * x + this->poly[i - 1];
    return static_cast<double>(ret);
#elif 1
    // Use loop unrolling which maybe is done by the compiler. Only when polyGrade = 2. Otherwise, add as many instructions as you want
    long double ret = this->poly[this->polyGrade];
    ret = ret * x + this->poly[this->polyGrade - 1];
    ret = ret * x + this->poly[this->polyGrade - 2];
    return static_cast<double>(ret);
#else
    // for regression-type of spline
    return x * this->poly[this->polyGrade] + this->poly[this->polyGrade - 1];
#endif
}
//---------------------------------------------------------------------------
double SplineEvaluation::hornerEvaluateByBuckets(unsigned bucketIndex, double& x) const
// evaluate x in the polynomial bucketIndex
{
    long double ret = this->polyBuckets[bucketIndex][polyGrade];
    for (unsigned i = polyGrade; i != 0; --i)
        ret = ret * x + this->polyBuckets[bucketIndex][i - 1];
    return static_cast<double>(ret);
}
//---------------------------------------------------------------------------
static const int popcount[4] = {0, 1, 1, 2};
//---------------------------------------------------------------------------
double SplineEvaluation::splittedHornerEvaluate(double& x) const
// evaluate the polynomials, depending on the spline where x lies in
{
    // load 2 x-values
    __m128d constant = _mm_set1_pd(x);
    // compare with greater than
    constant = _mm_cmpgt_pd(constant, this->saveLimits);
    // Get the result of the comparison
    int ret = _mm_movemask_pd(constant);
    // Count how many splitter x is greater than.
    return hornerEvaluateByBuckets(popcount[ret], x);
}
//---------------------------------------------------------------------------
unsigned SplineEvaluation::binarySearch(unsigned lower, unsigned upper, double& x) const
// classic branch-free binary search over the interval [lower, upper)
{
    // dodge the branch-misses
    for (unsigned half, n = upper - lower; (n / 2) != 0; n -= half) {
        half = (n / 2);
        unsigned middle = lower + half;
        lower = ((this->values[middle] < x) ? middle : lower);
    }
    return lower;
}
//---------------------------------------------------------------------------
unsigned SplineEvaluation::searchLeft(int index, double& x) const
// search in the left part of the point "index" to exact place where x lies
{
#if 1
    // We've already padded the left part with values lower than the first element of the array
    // In this way we don't have any overflows
    int pow = 1;
    while ((index >= pow) && (this->values[index - pow] > x))
        pow *= 2;
    return SplineEvaluation::binarySearch(max(index - pow, 0), index - pow / 2, x);
#else
    // We've already padded the left part with values lower than the first element of the array
    // In this way we don't have any overflows
    // Make a fast check for the first 2 values
    if (this->values[index - 1] < x)
        return index - 1;
    if (this->values[index - 2] < x)
        return index - 2;

    int pow = 4;
    // Load x in an ymm register
    __m128d constant = _mm_set1_pd(x);
    leftSseExpSearch : {
        // Load the next steps
        powers[0] = this->values[index - pow];
        powers[1] = this->values[index - (pow << 1)];
        __m128d xmm = _mm_loadu_pd(powers);

        // Compare x with the further steps
        xmm = _mm_cmplt_pd(constant, xmm);
        // Get the result of the comparison
        int ret = _mm_movemask_pd(xmm);

        // If x is still lower than all steps, continue searching
        if (ret == 0x3) {
            pow <<= 2;
            goto leftSseExpSearch;
        }
        // The last used power is shifted with the count of steps which are greater than x (the number of set bits in ret)
        pow <<= ret;
    }
    // Continue with binary search, based on the last power
    // Don't forget that we padded the array, so use "offset" a the lowest index of the array
    return SplineEvaluation::binarySearch(max(index - pow, this->offset), index - pow / 2, x);
#endif
}
//---------------------------------------------------------------------------
unsigned SplineEvaluation::searchRight(int index, double& x) const
// search in the right part of the point "index" to exact place where x lies
{
#if 1
    // We've already padded the right part with values greater than the last element of the array
    // In this way we don't have any overflows
    int pow = 1;
    while ((index + pow < this->length) && (this->values[index + pow] < x))
        pow <<= 1;
    return SplineEvaluation::binarySearch(index + pow / 2, min(index + pow, this->length - 1), x);
#else
    // We've already padded the right part with values greater than the last element of the array
    // In this way we don't have any overflows
    // Make a fast check for the first 2 values
    if (this->values[index + 1] > x)
        return index;
    if (this->values[index + 2] > x)
        return index + 1;

    int pow = 4;
    // Load x in an ymm register
    __m128d constant = _mm_set1_pd(x);
    rightSseExpSearch : {
        // Load the next steps
        powers[0] = this->values[index + pow];
        powers[1] = this->values[index + (pow << 1)];

        // Compare x with the further steps
        __m128d xmm = _mm_loadu_pd(powers);
        xmm = _mm_cmpgt_pd(constant, xmm);
        // Get the result of the comparison
        int ret = _mm_movemask_pd(xmm);

        // If x is still greater than all steps, continue searching
        if (ret == 0x3) {
            pow <<= 2;
            goto rightSseExpSearch;
        }
        // The last used power is shifted with the count of steps which are lower than x (the number of set bits in ret)
        pow <<= ret;
    }
    // Continue with binary search, based on the last power
    // Don't forget that we padded the array.
    return SplineEvaluation::binarySearch(index + pow / 2, min(index + pow, this->length - 1 + this->offset), x);
#endif
}
//---------------------------------------------------------------------------
static unsigned filterSegment(unsigned countOfSegments, double& segment)
// round segment to its int value, which should lie between [0, countOfSegments - 1]
{
    if (segment < -precision) {
       return 0;
    } else if (segment > countOfSegments) {
       return countOfSegments - 1;
    }
    // round and trunc can be used interchangeably
    return static_cast<unsigned>(segment);
}
//---------------------------------------------------------------------------
unsigned SplineEvaluation::findExactSegment(double segment, double& x) const
// find the segment where x lies
{
    // first approximate where x could lie
    unsigned candSegment = filterSegment(this->length - 1, segment);

    // find the exact segment with exponential search
    if (x < this->values[candSegment]) {
        // Continue searching in the left part
#if 1
        return searchLeft(candSegment, x);
#else
        return SplineEvaluation::binarySearch(this->offset, candSegment, x);
#endif
    } else if (x > this->values[candSegment + 1]) {
        // The condition evaluation doesn't make overflow, because noOfSegments == spline.size() - 1
        // Continue searching in the right part
#if 1
        return searchRight(candSegment + 1, x);
#else
        return SplineEvaluation::binarySearch(candSegment + 1, this->spline.size() - 1 + this->offset, x);
#endif
    }
    // x already lies on this segment
    return candSegment;
}
//---------------------------------------------------------------------------
double SplineEvaluation::selfInterpolate(unsigned segment, double& x) const
// get f(x) at segment
{
    Coord up = this->spline[segment + 1], down = this->spline[segment];
    double dx = up.first - down.first;
    double dy = up.second - down.second;
    return down.second + (x - down.first) * (dy / dx);
}
//---------------------------------------------------------------------------
double SplineEvaluation::chebyshevEvaluate(double& x) const
// evaluates f(x), after having fitted a polynomial to the spline
{
    // Check the boundaries
    if (x <= spline.front().first)
        return spline.front().second;
    if (x >= spline.back().first)
        return spline.back().second;
    //cerr << findExactSegment(hornerEvaluate(x), x) << endl;
    return selfInterpolate(findExactSegment(hornerEvaluate(x), x), x);
}
//---------------------------------------------------------------------------
double SplineEvaluation::splittedChebyshevEvaluate(double& x) const
// evaluates f(x), after having fitted a polynomial to the spline
{
    // Check the boundaries
    if (x <= spline.front().first)
        return spline.front().second;
    if (x >= spline.back().first)
        return spline.back().second;
    //cerr << findExactSegment(splittedHornerEvaluate(x), x) << endl;
    return selfInterpolate(findExactSegment(splittedHornerEvaluate(x), x), x);
}
//---------------------------------------------------------------------------
double SplineEvaluation::btreeSegmentsEvaluate(double& x) const
// evaluates f(x) with btree spline
{
    // Check the boundaries
    if (x <= spline.front().first)
        return spline.front().second;
    if (x >= spline.back().first)
        return spline.back().second;
    //cerr << "segment = " << (this->btreeSegments->lower_bound(x)->second) << "soll gleich " << realSegment(this->spline, x) << endl;
    auto iter = this->btreeSegments->lower_bound(x);
    return selfInterpolate(iter->second - 1, x);
}
//---------------------------------------------------------------------------
static double interpolate(const vector<Coord>& spline, double pos)
// evaluates the spline at pos
{
   if (pos<=spline.front().first)
      return spline.front().second;
   if (pos>=spline.back().first)
      return spline.back().second;
   auto iter=lower_bound(spline.begin(),spline.end(),pos,[](const Coord& a,double b) { return a.first<b; });
   if (iter->first==pos)
       return iter->second;

   double dx=(iter+0)->first-(iter-1)->first;
   double dy=(iter+0)->second-(iter-1)->second;

   double ofs=pos-(iter-1)->first;
   return (iter-1)->second+ofs*(dy/dx);
}
//---------------------------------------------------------------------------
double SplineEvaluation::evaluate(double& pos) const
// evaluates f(x)
{
#if 0
   if (pos<=spline.front().first)
      return spline.front().second;
   if (pos>=spline.back().first)
      return spline.back().second;
   bool now = pos < 0;
   auto iter=lower_bound(spline.begin()+now,spline.end()-now,pos,[](const Coord& a,double b) { return a.first<b; });
   if (iter->first==pos)
       return iter->second;

   double dx=(iter+0)->first-(iter-1)->first;
   double dy=(iter+0)->second-(iter-1)->second;

   double ofs=pos-(iter-1)->first;
   return (iter-1)->second+ofs*(dy/dx);
#else
    return interpolate(this->spline, pos);
#endif
}
//---------------------------------------------------------------------------
double SplineEvaluation::mapEvaluate(double& pos) const
// evaluate with the preprocessed map
{
    if (pos <= this->coordinates.begin()->first)
        return this->coordinates.begin()->second;
    if (pos >= (--this->coordinates.end())->first)
        return (--this->coordinates.end())->second;
    auto iter = this->coordinates.lower_bound(pos);
    if (iter->first == pos)
        return iter->second;
    return iter->second;
}
//---------------------------------------------------------------------------
double SplineEvaluation::btreeEvaluate(double& pos) const
// evaluate with the preprocessed map
{
    if (pos <= this->btreeMap->begin()->first)
        return this->btreeMap->begin()->second;
    if (pos >= (--this->btreeMap->end())->first)
        return (--this->btreeMap->end())->second;
    auto iter = this->btreeMap->lower_bound(pos);
    return iter->second;
}
//---------------------------------------------------------------------------
int main(int argc,char* argv[])
{
   if (argc < 7) {
      cerr << "usage: " << argv[0] << " file samplingPoints countToRead(-1 -> all) useOfSpline(0 -> without Spline, 1 -> with Spline) modeOfSearch(0 -> split & fit, 1 -> only fit, 2 -> lower_bound, 3 -> map, 4 -> btree, 5 -> btree & spline) testMode(0, 1) fileRandom" << endl;
      return 1;
   }
#if 1
    // Read from a binary file the function
   unsigned countToRead = atoi(argv[3]);
   vector<Coord> cdf;
   {
        ifstream in(argv[1]);
        auto pos = in.tellg();
        in.seekg( 0, ios::end );
        auto size=in.tellg()-pos;
        in.seekg(0,ios::beg);

        unsigned elements=size/sizeof(pair<double,double>);
        if ((countToRead != -1) && (countToRead < elements))
            elements = countToRead;
        cdf.resize(elements);
        in.read(reinterpret_cast<char*>(cdf.data()),elements*sizeof(Coord));
#if 0
        in.read(reinterpret_cast<char*>(cdf.data()),elements*sizeof(pair<double,double>));
        in.read(reinterpret_cast<char*>(cdf.data()),elements*sizeof(pair<double,double>));
#endif
    }
#else
    // Read from a .txt file
    unsigned countToRead = atoi(argv[3]);
    ifstream in(argv[1]);

    unsigned n;
    in >> n;
    vector<pair<double, double>> cdf(n);

    for (unsigned index = 0; index < n; ++index) {
        double x, y;
        in >> x >> y;
        cdf[index] = make_pair(x, y);
    }
#endif
    cerr << "Reading done!" << endl;

    unsigned size = atoi(argv[2]);
    unsigned useOfSpline = atoi(argv[4]);
    unsigned modeOfSearch = atoi(argv[5]);
    unsigned testMode = atoi(argv[6]);

    if (size > cdf.size()) {
        cerr << "Spline will contain more points than the function itself!" << endl;
        exit(0);
    }

    // Construct the spline and also fit the polynomial if modeSearch allows for that
    SplineEvaluation reduced(cdf, size, useOfSpline, modeOfSearch, testMode);

    // Shuffle the indexes in order to test
    std::vector<int> testIndexes(cdf.size());

    if (argc == 8) {
        ifstream randomFile(argv[7]);
        randomFile.read(reinterpret_cast<char*>(testIndexes.data()),cdf.size()*sizeof(int));
        randomFile.close();
    } else {
        for (unsigned index = 0; index < cdf.size(); ++index)
            testIndexes[index] = index;

        // This also has an influence upon the runtime
        unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
        shuffle(testIndexes.begin(), testIndexes.end(), std::default_random_engine(seed));
    }

    if (testMode == TEST_MODE_ACTIVATED) {
        // compare only those related to spline
        if (useOfSpline == WORK_ON_SPLINE_MODE) {
            cerr << "Start testing with spline!" << endl;
            double curr = 0;
            unsigned alright = 0;
            for (unsigned index = 0; index < cdf.size(); ++index) {
                Coord elem = cdf[testIndexes[index]];
                elem.first += 0.000;
                //cerr << setprecision(12) << "real segment of (" << elem.first << ", " << elem.second << ") is " << realSegment(reduced.spline, elem.first) << endl;
                double y;
                if (modeOfSearch == SSE_SPLIT_MODE)
                    y = reduced.splittedChebyshevEvaluate(elem.first);
                else if (modeOfSearch == JUST_FIT_MODE)
                    y = reduced.chebyshevEvaluate(elem.first);
                else if (modeOfSearch == BTREE_SPLINE_MODE)
                    y = reduced.btreeSegmentsEvaluate(elem.first);
                else {
                    cerr << "Option in test with Spline not available " << modeOfSearch << endl;
                    exit(0);
                }
                curr = max(curr, fabs(y - elem.second));
#if 1
                double cmp = reduced.evaluate(elem.first);
                if (!(fabs(cmp - y) < precision)) {
                    cerr << "Should be equal for " << setprecision(12) << fabs(cmp - y) << " " << elem.first << " --> " << setprecision(12) << y << " with " << setprecision(12)<<cmp << " against " << elem.second << endl;
                } else {
                    alright++;
                }
#endif
            }
            if (alright == cdf.size())
                cerr << "parfum!!!" << endl;
            else
                cerr << "try again!!!Sth went wrong!!" << endl;
            cerr << setprecision(12) << "check: " << curr / cdf.size() << endl;
        } else {
            cerr << "Start testing without spline!" << endl;
            double curr = 0;
            unsigned alright = 0;
            for (unsigned index = 0; index < cdf.size(); ++index) {
                Coord elem = cdf[testIndexes[index]];
                elem.first += 0.005;
                //cerr << setprecision(12) << "real segment of (" << elem.first << ", " << elem.second << ") is " << realSegment(reduced.spline, elem.first) << endl;
                double y;
                if (modeOfSearch == BTREE_MODE)
                    y = reduced.btreeEvaluate(elem.first);
                else {
                    cerr << "Option in test without Spline not available " << modeOfSearch << endl;
                    exit(0);
                }
                curr = max(curr, fabs(y - elem.second));
#if 1
                double cmp = reduced.mapEvaluate(elem.first);
                if (!(fabs(cmp - y) < precision)) {
                    cerr << "Should be equal for " << setprecision(12) << fabs(cmp - y) << " " << elem.first << " --> " << setprecision(12) << y << " with " << setprecision(12)<<cmp << " against " << elem.second << endl;
                } else {
                    alright++;
                }
#endif
            }
            if (alright == cdf.size())
                cerr << "parfume!!!" << endl;
            else
                cerr << "try again!!! Something went wrong!" << endl;
            cerr << setprecision(12) << "check: " << curr / cdf.size() << endl;
        }
    } else {
        auto start = std::chrono::high_resolution_clock::now();
        switch (modeOfSearch) {
            // use a split through SSE comparing, then fit
            case SSE_SPLIT_MODE : {
                double sum = 0, maxDiff = 0;
                for (unsigned index = 0; index < cdf.size(); ++index) {
#if 1
                    reduced.splittedChebyshevEvaluate(cdf[testIndexes[index]].first);
#else
                    // Test the deviations from the polynomial to the exact segment of each point
                    Coord elem = cdf[testIndexes[index]];
                    double segment = reduced.splittedHornerEvaluate(elem.first);
                    unsigned candSegment = filterSegment(reduced.spline.size() - 1, segment);
                    unsigned exactSeg = realSegment(reduced.spline, elem.first);

                    double diff = (candSegment > exactSeg) ? (candSegment - exactSeg) : (exactSeg - candSegment);
                    sum += diff;
                    maxDiff = max(maxDiff, diff);
#endif
                }
                cerr << "globalDiff = " << globalDiff << " mit " << globalDiff / cdf.size() << endl;
                cerr << "Avg deviation : " << (sum / cdf.size()) << " Max deviation : " << maxDiff << endl;
                break;
            }
            // Apply the fit of polynomial O(log(chebysevNorm))
            case JUST_FIT_MODE : {
                double sum = 0, maxDiff = 0;
                for (unsigned index = 0; index < cdf.size(); ++index) {
#if 1
                    reduced.chebyshevEvaluate(cdf[testIndexes[index]].first);
#else
                    // Test the deviations from the polynomial to the exact segment of each point
                    Coord elem = cdf[testIndexes[index]];
                    double segment = reduced.hornerEvaluate(elem.first);
                    unsigned candSegment = filterSegment(reduced.spline.size() - 1, segment);
                    unsigned exactSeg = realSegment(reduced.spline, elem.first);

                    double diff = (candSegment > exactSeg) ? (candSegment - exactSeg) : (exactSeg - candSegment);
                    sum += diff;
                    maxDiff = max(maxDiff, diff);
#endif
                }
                cerr << "globalDiff = " << globalDiff << " mit " << globalDiff / cdf.size() << endl;
                cerr << "Avg deviation : " << (sum / cdf.size()) << " Max deviation : " << maxDiff << endl;
                break;
            }
            // Use only lower_bound O(log(size of spline))
            case LOWER_BOUND_MODE : {
                cerr << cdf.size() << endl;
                for (unsigned index = 0; index < cdf.size(); ++index) {
                    reduced.evaluate(cdf[testIndexes[index]].first);
                }
                break;
            }
            case MAP_MODE : {
                for (unsigned index = 0; index < cdf.size(); ++index) {
                    reduced.mapEvaluate(cdf[testIndexes[index]].first);
                }
                break;
            }
            case BTREE_MODE : {
                for (unsigned index = 0; index < cdf.size(); ++index) {
                    reduced.btreeEvaluate(cdf[testIndexes[index]].first);
                }
                break;
            }
            case BTREE_SPLINE_MODE : {
                for (unsigned index = 0; index < cdf.size(); ++index) {
                    reduced.btreeSegmentsEvaluate(cdf[testIndexes[index]].first);
                }
                break;
            }
            default : {
                cerr << "Choose again!" << endl;
            }
        }
        auto stop = std::chrono::high_resolution_clock::now();
        auto duration = stop - start;
        cerr << "modeOfSearch (0 -> split & fit, 1 -> only fit, 2 -> lower_bound, 3 -> map) : " << modeOfSearch << " took: " << duration.count() * 1e-6 << "ms" << endl;
    } // end of TEST_MODE_ACTIVATED
    return 0;
}
//---------------------------------------------------------------------------


//
//  E_step.cpp
//  synthetic_data_helper
//
//  Created by Cristián Garay on 11/16/16.
//  Revised and edited by Anirudhan Badrinath on 27/02/20.
//

#define NPY_NO_DEPRECATED_API NPY_1_11_API_VERSION

#include <iostream>
#include <stdint.h>
#include <alloca.h>
#include <Eigen/Core>
#include <omp.h>
#include <Python.h>
#include <numpy/ndarrayobject.h>

using namespace Eigen;
using namespace std;

//original comment:
//"TODO if we aren't outputting gamma, don't need to write it to memory (just
//need t and t+1), so we can save the stack array for each HMM at the cost of
//a branch"
//

static double extract_double(PyArrayObject *arr, int i) {
    return ((double*) PyArray_DATA(arr))[i];
}

void capsule_cleanup(PyObject *capsule) {
    void *memory = PyCapsule_GetPointer(capsule, NULL);
    delete memory;
}

static PyObject* run(PyObject * module, PyObject * args) {
    //TODO: check if parameters are null.
    //TODO: check that dicts have the required members.
    //TODO: check that all parameters have the right sizes.
    //TODO: i'm not sending any error messages.
    IOFormat CleanFmt(4, 0, ", ", "\n", "[", "]");
    Eigen::initParallel();
    import_array();

    PyObject *data_ptr = NULL, *model_ptr = NULL, *t_softcounts = NULL, *e_softcounts = NULL, *i_softcounts = NULL, *fixed = NULL;
    PyArrayObject *t_softcounts_np = NULL, *e_softcounts_np = NULL, *i_softcounts_np = NULL,
                  *alldata = NULL, *allresources = NULL, *starts = NULL, *lengths = NULL, *learns = NULL, *forgets = NULL, *guesses = NULL, *slips = NULL;   // Extended Numpy/C API
    int num_outputs, parallel;
    double prior;
    // dict& data, dict& model, numpy::ndarray& trans_softcounts, numpy::ndarray& emission_softcounts, numpy::ndarray& init_softcounts, int num_outputs

    // "O" format -> read argument as a PyObject type into argy (Python/C API)
    if (!PyArg_ParseTuple(args, "OOiiO", &data_ptr, &model_ptr, &num_outputs, &parallel, &fixed)) {
        PyErr_SetString(PyExc_ValueError, "Error parsing arguments.");
        return NULL;
    }

    if (!parallel)
        omp_set_num_threads(1);

    // Load all the numpy arrays in data & model
    char* DM_NAMES[] = {"data", "resources", "starts", "lengths", "learns", "forgets", "guesses", "slips"};
    PyArrayObject** DM_PTRS[] = {&alldata, &allresources, &starts, &lengths, &learns, &forgets, &guesses, &slips};
    for (int i = 0; i < 8; i++) {
        PyObject *dp = PyDict_GetItemString(i < 4 ? data_ptr : model_ptr, DM_NAMES[i]);
        int DTYPE = PyArray_ObjectType(dp, (i < 4 ? (i < 1 ? NPY_INT : NPY_INT64) : NPY_FLOAT)); // hack to force correct type
        *DM_PTRS[i] = (PyArrayObject *)PyArray_FROM_OTF(dp, DTYPE, NPY_ARRAY_IN_ARRAY);
    }
    prior = PyFloat_AsDouble(PyDict_GetItemString(model_ptr, "prior"));

    int bigT = (int) PyArray_DIM(alldata, 1), num_subparts = (int) PyArray_DIM(alldata, 0);
    int len_allresources = (int) PyArray_DIM(allresources, 0);
    int num_sequences = (int) PyArray_DIM(starts, 0);
    int len_lengths = (int) PyArray_DIM(lengths, 0);
    int num_resources = (int) PyArray_DIM(learns, 0);

    Map<Array<int32_t,Dynamic,Dynamic,RowMajor>,Aligned> alldata_arr(reinterpret_cast<int*>(PyArray_DATA(alldata)),num_subparts,bigT);
    Map<Array<int64_t, Eigen::Dynamic, 1>,Aligned> allresources_arr(reinterpret_cast<int64_t*>(PyArray_DATA(allresources)),len_allresources,1);
    Map<Array<int64_t, Eigen::Dynamic, 1>,Aligned> starts_arr(reinterpret_cast<int64_t*>(PyArray_DATA(starts)),num_sequences,1);
    Map<Array<int64_t, Eigen::Dynamic, 1>,Aligned> lengths_arr(reinterpret_cast<int64_t*>(PyArray_DATA(lengths)),len_lengths,1);

    bool normalizeLengths = false;
    //then the original code goes to find the optional parameters.

    PyObject* fixed_prior = PyDict_GetItemString(fixed, "prior");
    PyObject* fixed_learn = (PyObject*)PyDict_GetItemString(fixed, "learn");
    PyObject* fixed_forget = (PyObject*)PyDict_GetItemString(fixed, "forget");
    PyObject* fixed_guess = (PyObject*)PyDict_GetItemString(fixed, "guess");
    PyObject* fixed_slip = (PyObject*)PyDict_GetItemString(fixed, "slip");

    if (fixed_prior) {
        prior = PyFloat_AsDouble(fixed_prior);
    }
    Array2d initial_distn;
    initial_distn << 1-prior, prior;

    MatrixXd As(2,2*num_resources);
    double learn = -1;
    double forget = -1;

    for (int n=0; n<num_resources; n++) {
        if (fixed_learn) {
            learn = extract_double((PyArrayObject*)fixed_learn, n);
        }
        if (learn < 0) {
            learn = extract_double(learns, n);
        }
        if (fixed_forget) {
            forget = extract_double((PyArrayObject*)fixed_forget, n);
        }
        if (forget < 0) {
            forget = extract_double(forgets, n);
        }
        As.col(2*n) << 1-learn, learn;
        As.col(2*n+1) << forget, 1-forget;
    }

    Array2Xd Bn(2,2*num_subparts);
    double guess = -1;
    double slip = -1;
    for (int n=0; n<num_subparts; n++) {
        if (fixed_guess) {
            guess = extract_double((PyArrayObject*)fixed_guess, n);
        }
        if (guess < 0) {
            guess = extract_double(guesses, n);
        }
        if (fixed_slip) {
            slip = extract_double((PyArrayObject*)fixed_slip, n);
        }
        if (slip < 0) {
            slip = extract_double(slips, n);
        }
        Bn.col(2*n) << 1-guess, slip; // incorrect
        Bn.col(2*n+1) << guess, 1-slip; // correct
    }

    //// outputs

    //TODO: NEED TO FIX THIS I'M CREATING NEW ARRAYS AND I NEED TO USE THE ARGUMENTS!!!
    //TODO: FIX THIS!!!
    /*Map<ArrayXXd,Aligned> all_trans_softcounts(trans_softcounts,2,2*num_resources);
    all_trans_softcounts.setZero();
    Map<Array2Xd,Aligned> all_emission_softcounts(emission_softcounts,2,2*num_subparts);
    all_emission_softcounts.setZero();
    Map<Array2d,Aligned> all_initial_softcounts(init_softcounts);
    all_initial_softcounts.setZero();*/
    //TODO: I replaced the pointers to the arguments for new Eigen arrays.
    /*ArrayXXd all_trans_softcounts(2,2*num_resources);
    all_trans_softcounts.setZero(); //why is he setting all these to zero???
    Array2Xd all_emission_softcounts(2,2*num_subparts);
    all_emission_softcounts.setZero();
    Array2d all_initial_softcounts(2, 1); //should i use these dimensions? the same as the original vector??
    all_initial_softcounts.setZero();*/
    double* r_trans_softcounts = new double[2*2*num_resources];
    double* r_emission_softcounts = new double[2*2*num_subparts];
    double* r_init_softcounts = new double[2*1];
    Map<ArrayXXd,Aligned> all_trans_softcounts(r_trans_softcounts,2,2*num_resources);
    all_trans_softcounts.setZero();
    Map<Array2Xd,Aligned> all_emission_softcounts(r_emission_softcounts,2,2*num_subparts);
    all_emission_softcounts.setZero();
    Map<Array2d,Aligned> all_initial_softcounts(r_init_softcounts);
    all_initial_softcounts.setZero();

    //TODO: FIX THIS!!! I'll replace all these weird arrays for zeroes ones.
    //Array2Xd likelihoods_out(2,bigT);
    //likelihoods_out.setZero();
    //Array2Xd gamma_out(2,bigT);
    //gamma_out.setZero();
    //Array2Xd alpha_out(2,bigT);
    //alpha_out.setZero();
    Map<Array2Xd,Aligned> alpha_out(NULL,2,bigT);
    double s_total_loglike = 0;
    double *total_loglike = &s_total_loglike;

    //TODO: FIX THIS!!! why is he doing this??
    /* switch (num_outputs)
    {
        case 4:
            plhs[3] = mxCreateDoubleMatrix(2,bigT,mxREAL);
            new (&likelihoods_out) Map<Array2Xd,Aligned>(mxGetPr(plhs[3]),2,bigT);
        case 3:
            plhs[2] = mxCreateDoubleMatrix(2,bigT,mxREAL);
            new (&gamma_out) Map<Array2Xd,Aligned>(mxGetPr(plhs[2]),2,bigT);
        case 2:
            plhs[1] = mxCreateDoubleMatrix(2,bigT,mxREAL);
            new (&alpha_out) Map<Array2Xd,Aligned>(mxGetPr(plhs[1]),2,bigT);
        case 1:
            plhs[0] = mxCreateDoubleScalar(0.);
            total_loglike = mxGetPr(plhs[0]);
    }*/
    double* r_alpha_out = new double[2 * bigT];
    new (&alpha_out) Map<Array2Xd,Aligned>(r_alpha_out,2,bigT);

    /* COMPUTATION */
    #pragma omp parallel
    {
        double s_trans_softcounts[2*2*num_resources] __attribute__((aligned(16)));
        double s_emission_softcounts[2*2*num_subparts] __attribute__((aligned(16)));
        Map<ArrayXXd,Aligned> trans_softcounts_temp(s_trans_softcounts,2,2*num_resources);
        Map<ArrayXXd,Aligned> emission_softcounts_temp(s_emission_softcounts,2,2*num_subparts);
        Array2d init_softcounts_temp;
        double loglike;

        trans_softcounts_temp.setZero();
        emission_softcounts_temp.setZero();
        init_softcounts_temp.setZero();
        loglike = 0;
        int num_threads = omp_get_num_threads();
        int blocklen = 1 + ((num_sequences - 1) / num_threads);
        int sequence_idx_start = blocklen * omp_get_thread_num();
        int sequence_idx_end = min(sequence_idx_start+blocklen,num_sequences);
        //mexPrintf("start:%d   end:%d\n", sequence_idx_start, sequence_idx_end);

        for (int sequence_index=sequence_idx_start; sequence_index < sequence_idx_end; sequence_index++) {

            // NOTE: -1 because Matlab indexing starts at 1
            int64_t sequence_start = starts_arr(sequence_index, 0) - 1;

            int64_t T = lengths_arr(sequence_index, 0);

            //// likelihoods
            double s_likelihoods[2*T];
            Map<Array2Xd,Aligned> likelihoods(s_likelihoods,2,T);

            likelihoods.setOnes();
             for (int t=0; t<T; t++) {
                 for (int n=0; n<num_subparts; n++) {
                    int32_t data_temp = alldata_arr(n, sequence_start+t);
                    if (data_temp != 0) {
                        for (int i = 0; i < likelihoods.rows(); i++)
                            if (Bn(i, 2*n + (data_temp == 2)) != 0)
                                likelihoods(i, t) *= Bn(i, 2*n + (data_temp == 2));
                    }
                 }
             }

            //// forward messages
            double norm;
            double s_alpha[2*T] __attribute__((aligned(16)));
            double contribution;
            Map<MatrixXd,Aligned> alpha(s_alpha,2,T);
            alpha.col(0) = initial_distn * likelihoods.col(0);
            norm = alpha.col(0).sum();
            alpha.col(0) /= norm;
            loglike += log(norm) / (normalizeLengths? T : 1);

            for (int t=0; t<T-1; t++) {
                int64_t resources_temp = allresources_arr(sequence_start+t, 0);
                alpha.col(t+1) = (As.block(0,2*(resources_temp-1),2,2) * alpha.col(t)).array()
                    * likelihoods.col(t+1);
                norm = alpha.col(t+1).sum();
                alpha.col(t+1) /= norm;
                loglike += log(norm) / (normalizeLengths? T : 1);
            }

            //// backward messages and statistic counting

            double s_gamma[2*T] __attribute__((aligned(16)));
            Map<Array2Xd,Aligned> gamma(s_gamma,2,T);
            gamma.col(T-1) = alpha.col(T-1);
            for (int n=0; n<num_subparts; n++) {
                int32_t data_temp = alldata_arr(n, sequence_start+(T-1));
                if (data_temp != 0) {
                    emission_softcounts_temp.col(2*n + (data_temp == 2)) += gamma.col(T-1);
                }
            }

            for (int t=T-2; t>=0; t--) {

				int64_t resources_temp = allresources_arr(sequence_start+t, 0);
                Matrix2d A = As.block(0,2*(resources_temp-1),2,2);
                Array22d pair = A.array();
                pair.rowwise() *= alpha.col(t).transpose().array();
                pair.colwise() *= gamma.col(t+1);
                pair.colwise() /= (A*alpha.col(t)).array();
                pair = (pair != pair).select(0.,pair); // NOTE: replace NaNs
                trans_softcounts_temp.block(0,2*(resources_temp-1),2,2) += pair;

                gamma.col(t) = pair.colwise().sum().transpose();
                // NOTE: we have to touch the data again here
                for (int n=0; n<num_subparts; n++) {
                    int32_t data_temp = alldata_arr(n, sequence_start+t);
                    if (data_temp != 0) {
                        emission_softcounts_temp.col(2*n + (data_temp == 2)) += gamma.col(t);
                    }
                }
            }
            init_softcounts_temp += gamma.col(0);

            //TODO: FIX THIS!!!
            /* switch (nlhs)
            {
                case 4:
                    likelihoods_out.block(0,sequence_start,2,T) = likelihoods;
                case 3:
                    gamma_out.block(0,sequence_start,2,T) = gamma;
                case 2:
                    alpha_out.block(0,sequence_start,2,T) = alpha;
            } */
            alpha_out.block(0,sequence_start,2,T) = alpha;
        }
        #pragma omp critical
        {
            all_trans_softcounts += trans_softcounts_temp;
            all_emission_softcounts += emission_softcounts_temp;
            all_initial_softcounts += init_softcounts_temp;
            *total_loglike += loglike;
        }
    }

    PyObject *result = PyDict_New();

    npy_intp dims1[] = {num_resources, 2, 2};
    PyObject *all_trans_softcounts_arr = (PyObject *) PyArray_SimpleNewFromData(3, dims1, NPY_DOUBLE, r_trans_softcounts);
    PyObject *capsule1 = PyCapsule_New(r_trans_softcounts, NULL, capsule_cleanup);
    PyArray_SetBaseObject((PyArrayObject *) all_trans_softcounts_arr, capsule1);

    npy_intp dims2[] = {num_subparts, 2, 2};
    PyObject *all_emission_softcounts_arr = (PyObject *) PyArray_SimpleNewFromData(3, dims2, NPY_DOUBLE, r_emission_softcounts);
    PyObject *capsule2 = PyCapsule_New(r_emission_softcounts, NULL, capsule_cleanup);
    PyArray_SetBaseObject((PyArrayObject *) all_emission_softcounts_arr, capsule2);

    npy_intp dims3[] = {2, 1};
    PyObject *all_initial_softcounts_arr = (PyObject *) PyArray_SimpleNewFromData(2, dims3, NPY_DOUBLE, r_init_softcounts);
    PyObject *capsule3 = PyCapsule_New(r_init_softcounts, NULL, capsule_cleanup);
    PyArray_SetBaseObject((PyArrayObject *) all_initial_softcounts_arr, capsule3);

    npy_intp dims4[] = {2, bigT};
    PyObject *alpha_out_arr = (PyObject *) PyArray_SimpleNewFromData(2, dims4, NPY_DOUBLE, r_alpha_out);
    PyObject *capsule4 = PyCapsule_New(r_alpha_out, NULL, capsule_cleanup);
    PyArray_SetBaseObject((PyArrayObject *) alpha_out_arr, capsule4);

    PyDict_SetItemString(result, "all_trans_softcounts", all_trans_softcounts_arr);
    PyDict_SetItemString(result, "all_emission_softcounts", all_emission_softcounts_arr);
    PyDict_SetItemString(result, "all_initial_softcounts", all_initial_softcounts_arr);
    PyDict_SetItemString(result, "alpha", alpha_out_arr);
    PyDict_SetItemString(result, "total_loglike", PyLong_FromLong(*total_loglike));

    for (int i = 0; i < 8; i++)
        Py_XDECREF(*DM_PTRS[i]);
    Py_XDECREF(all_trans_softcounts_arr);
    Py_XDECREF(all_emission_softcounts_arr);
    Py_XDECREF(all_initial_softcounts_arr);
    Py_XDECREF(alpha_out_arr);

    return(result);
}

static PyMethodDef E_step_Methods[] = {
    {"run",  run, METH_VARARGS,
     "Runs E-step of Expectation Maximization in C++ module"},
    {NULL, NULL, 0, NULL}        /* Sentinel */
};

static struct PyModuleDef E_step_module = {
   PyModuleDef_HEAD_INIT,
   "E_step",   /* name of module */
   NULL, /* module documentation, may be NULL */
   -1,       /* size of per-interpreter state of the module,
                or -1 if the module keeps state in global variables. */
   E_step_Methods
};

PyMODINIT_FUNC PyInit_E_step() {
    return PyModule_Create(&E_step_module);
}


//---------------------------------------------------------------------------//
//!
//! \file   tstPolynomialDistribution.cpp
//! \author Alex Robinson
//! \brief  Polynomial distribution unit tests
//!
//---------------------------------------------------------------------------//

// Std Lib Includes
#include <iostream>

// Boost Includes
#include <boost/units/systems/si.hpp>
#include <boost/units/systems/cgs.hpp>
#include <boost/units/io.hpp>

// Trilinos Includes
#include <Teuchos_UnitTestHarness.hpp>
#include <Teuchos_RCP.hpp>
#include <Teuchos_Array.hpp>
#include <Teuchos_ParameterList.hpp>
#include <Teuchos_XMLParameterListCoreHelpers.hpp>
#include <Teuchos_VerboseObject.hpp>

// FRENSIE Includes
#include "Utility_UnitTestHarnessExtensions.hpp"
#include "Utility_OneDDistribution.hpp"
#include "Utility_PolynomialDistribution.hpp"
#include "Utility_RandomNumberGenerator.hpp"
#include "Utility_PhysicalConstants.hpp"
#include "Utility_UnitTraits.hpp"
#include "Utility_QuantityTraits.hpp"
#include "Utility_ElectronVoltUnit.hpp"

using boost::units::quantity;
using namespace Utility::Units;
namespace si = boost::units::si;
namespace cgs = boost::units::cgs;

//---------------------------------------------------------------------------//
// Testing Variables
//---------------------------------------------------------------------------//

Teuchos::RCP<Teuchos::ParameterList> test_dists_list;

Teuchos::RCP<Utility::OneDDistribution> distribution;

Teuchos::RCP<Utility::UnitAwareOneDDistribution<MegaElectronVolt,si::amount> >
  unit_aware_distribution;

//---------------------------------------------------------------------------//
// Tests.
//---------------------------------------------------------------------------//
// Check that the distribution can be evaluated
TEUCHOS_UNIT_TEST( PolynomialDistribution, evaluate )
{
  TEST_EQUALITY_CONST( distribution->evaluate( -1.0 ), 0.0 );
  TEST_EQUALITY_CONST( distribution->evaluate( 0.0 ), 1.0 );
  TEST_EQUALITY_CONST( distribution->evaluate( 0.5 ), 2.75 );
  TEST_EQUALITY_CONST( distribution->evaluate( 1.0 ), 6.0 );
  TEST_EQUALITY_CONST( distribution->evaluate( 2.0 ), 0.0 );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be evaluated
TEUCHOS_UNIT_TEST( UnitAwarePolynomialDistribution, evaluate )
{
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluate( -1.0*MeV ),
		       0.0*si::mole );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluate( 0.0*MeV ),
		       1.0*si::mole );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluate( 0.5*MeV ),
		       2.75*si::mole );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluate( 1.0*MeV ),
		       6.0*si::mole );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluate( 2.0*MeV ),
		       0.0*si::mole );
}

//---------------------------------------------------------------------------//
// Check that the PDF can be evaluated
TEUCHOS_UNIT_TEST( PolynomialDistribution, evaluatePDF )
{
  TEST_EQUALITY_CONST( distribution->evaluatePDF( -1.0 ), 0.0 );
  TEST_EQUALITY_CONST( distribution->evaluatePDF( 0.0 ), 1.0/3.0 );
  TEST_EQUALITY_CONST( distribution->evaluatePDF( 0.5 ), 2.75/3.0 );
  TEST_EQUALITY_CONST( distribution->evaluatePDF( 1.0 ), 2.0 );
  TEST_EQUALITY_CONST( distribution->evaluatePDF( 2.0 ), 0.0 );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware PDF can be evaluated
TEUCHOS_UNIT_TEST( UnitAwarePolynomialDistribution, evaluatePDF )
{
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluatePDF( -1.0*MeV ),
		       0.0/MeV );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluatePDF( 0.0*MeV ),
		       (1.0/3.0)/MeV );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluatePDF( 0.5*MeV ),
		       (2.75/3.0)/MeV );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluatePDF( 1.0*MeV ),
		       2.0/MeV );
  TEST_EQUALITY_CONST( unit_aware_distribution->evaluatePDF( 2.0*MeV ),
		       0.0/MeV );
}

//---------------------------------------------------------------------------//
// Check that the distribution can be sampled
TEUCHOS_UNIT_TEST( PolynomialDistribution, sample )
{
  std::vector<double> fake_stream( 6 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.0;
  fake_stream[2] = 0.0;
  fake_stream[3] = 0.5;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );
  double sample = distribution->sample();
  TEST_EQUALITY_CONST( sample, 0.0 );

  sample = distribution->sample();
  TEST_EQUALITY_CONST( sample, 0.5 );

  fake_stream[0] = 0.5;
  fake_stream[1] = 0.0;
  fake_stream[2] = 0.5;
  fake_stream[3] = 0.5;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );
  sample = distribution->sample();
  TEST_EQUALITY_CONST( sample, 0.0 );

  sample = distribution->sample();
  TEST_EQUALITY_CONST( sample, sqrt(0.5) );

  fake_stream[0] = 0.7;
  fake_stream[1] = 0.0;
  fake_stream[2] = 0.7;
  fake_stream[3] = 0.5;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );
  sample = distribution->sample();
  TEST_EQUALITY_CONST( sample, 0.0 );

  sample = distribution->sample();
  TEST_EQUALITY_CONST( sample, pow( 0.5, 1.0/3.0 ) );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be sampled
TEUCHOS_UNIT_TEST( UnitAwarePolynomialDistribution, sample )
{
  std::vector<double> fake_stream( 6 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.0;
  fake_stream[2] = 0.0;
  fake_stream[3] = 0.5;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );
  quantity<MegaElectronVolt> sample = unit_aware_distribution->sample();
  TEST_EQUALITY_CONST( sample, 0.0*MeV );

  sample = unit_aware_distribution->sample();
  TEST_EQUALITY_CONST( sample, 0.5*MeV );

  fake_stream[0] = 0.5;
  fake_stream[1] = 0.0;
  fake_stream[2] = 0.5;
  fake_stream[3] = 0.5;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );
  sample = unit_aware_distribution->sample();
  TEST_EQUALITY_CONST( sample, 0.0*MeV );

  sample = unit_aware_distribution->sample();
  TEST_EQUALITY_CONST( sample, sqrt(0.5)*MeV );

  fake_stream[0] = 0.7;
  fake_stream[1] = 0.0;
  fake_stream[2] = 0.7;
  fake_stream[3] = 0.5;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );
  sample = unit_aware_distribution->sample();
  TEST_EQUALITY_CONST( sample, 0.0*MeV );

  sample = unit_aware_distribution->sample();
  TEST_EQUALITY_CONST( sample, pow( 0.5, (1.0/3.0) )*MeV );
}

//---------------------------------------------------------------------------//
// Check that the distribution can be sampled
TEUCHOS_UNIT_TEST( PolynomialDistribution, sampleAndRecordTrials )
{
  std::vector<double> fake_stream( 6 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.0;
  fake_stream[2] = 0.0;
  fake_stream[3] = 0.5;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  unsigned trials = 0;

  double sample = distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, 0.0 );
  TEST_EQUALITY_CONST( 1.0/trials, 1.0 );

  sample = distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, 0.5 );
  TEST_EQUALITY_CONST( 2.0/trials, 1.0 );

  fake_stream[0] = 0.5;
  fake_stream[1] = 0.0;
  fake_stream[2] = 0.5;
  fake_stream[3] = 0.5;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  trials = 0;

  sample = distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, 0.0 );
  TEST_EQUALITY_CONST( 1.0/trials, 1.0 );

  sample = distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, sqrt(0.5) );
  TEST_EQUALITY_CONST( 2.0/trials, 1.0 );

  fake_stream[0] = 0.7;
  fake_stream[1] = 0.0;
  fake_stream[2] = 0.7;
  fake_stream[3] = 0.5;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  trials = 0;

  sample = distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, 0.0 );
  TEST_EQUALITY_CONST( 1.0/trials, 1.0 );

  sample = distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, pow( 0.5, 1.0/3.0 ) );
  TEST_EQUALITY_CONST( 2.0/trials, 1.0 );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be sampled
TEUCHOS_UNIT_TEST( UnitAwarePolynomialDistribution, sampleAndRecordTrials )
{
  std::vector<double> fake_stream( 6 );
  fake_stream[0] = 0.0;
  fake_stream[1] = 0.0;
  fake_stream[2] = 0.0;
  fake_stream[3] = 0.5;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  unsigned trials = 0;

  quantity<MegaElectronVolt> sample =
    unit_aware_distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, 0.0*MeV );
  TEST_EQUALITY_CONST( 1.0/trials, 1.0 );

  sample = unit_aware_distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, 0.5*MeV );
  TEST_EQUALITY_CONST( 2.0/trials, 1.0 );

  fake_stream[0] = 0.5;
  fake_stream[1] = 0.0;
  fake_stream[2] = 0.5;
  fake_stream[3] = 0.5;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  trials = 0;

  sample = unit_aware_distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, 0.0*MeV );
  TEST_EQUALITY_CONST( 1.0/trials, 1.0 );

  sample = unit_aware_distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, sqrt(0.5)*MeV );
  TEST_EQUALITY_CONST( 2.0/trials, 1.0 );

  fake_stream[0] = 0.7;
  fake_stream[1] = 0.0;
  fake_stream[2] = 0.7;
  fake_stream[3] = 0.5;

  Utility::RandomNumberGenerator::setFakeStream( fake_stream );

  trials = 0;

  sample = unit_aware_distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, 0.0*MeV );
  TEST_EQUALITY_CONST( 1.0/trials, 1.0 );

  sample = unit_aware_distribution->sampleAndRecordTrials( trials );
  TEST_EQUALITY_CONST( sample, pow( 0.5, 1.0/3.0 )*MeV );
  TEST_EQUALITY_CONST( 2.0/trials, 1.0 );
}

//---------------------------------------------------------------------------//
// Check that the upper bound of the distribution independent variable can be
// returned
TEUCHOS_UNIT_TEST( PolynomialDistribution, getUpperBoundOfIndepVar )
{
  TEST_EQUALITY_CONST( distribution->getUpperBoundOfIndepVar(), 1.0 );
}

//---------------------------------------------------------------------------//
// Check that the upper bound of the unit-aware distribution independent
// variable can be returned
TEUCHOS_UNIT_TEST( UnitAwarePolynomialDistribution, getUpperBoundOfIndepVar )
{
  TEST_EQUALITY_CONST( unit_aware_distribution->getUpperBoundOfIndepVar(),
		       1.0*MeV );
}

//---------------------------------------------------------------------------//
// Check that the lower bound of the unit-aware distribution independent
// variable can be returned
TEUCHOS_UNIT_TEST( UnitAwarePolynomialDistribution, getLowerBoundOfIndepVar )
{
  TEST_EQUALITY_CONST( unit_aware_distribution->getLowerBoundOfIndepVar(),
		       0.0*MeV );
}

//---------------------------------------------------------------------------//
// Check that the distribution type can be returned
TEUCHOS_UNIT_TEST( PolynomialDistribution, getDistributionType )
{
  TEST_EQUALITY_CONST( distribution->getDistributionType(),
		       Utility::POLYNOMIAL_DISTRIBUTION );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution type can be returned
TEUCHOS_UNIT_TEST( UnitAwarePolynomialDistribution, getDistributionType )
{
  TEST_EQUALITY_CONST( unit_aware_distribution->getDistributionType(),
		       Utility::POLYNOMIAL_DISTRIBUTION );
}

//---------------------------------------------------------------------------//
// Check if the distribution is tabular
TEUCHOS_UNIT_TEST( PolynomialDistribution, isTabular )
{
  TEST_ASSERT( !distribution->isTabular() );
}

//---------------------------------------------------------------------------//
// Check if the unit-aware distribution is tabular
TEUCHOS_UNIT_TEST( UnitAwarePolynomialDistribution, isTabular )
{
  TEST_ASSERT( !unit_aware_distribution->isTabular() );
}

//---------------------------------------------------------------------------//
// Check if the distribution is continuous
TEUCHOS_UNIT_TEST( PolynomialDistribution, isContinuous )
{
  TEST_ASSERT( distribution->isContinuous() );
}

//---------------------------------------------------------------------------//
// Check if the unit-aware distribution is continuous
TEUCHOS_UNIT_TEST( UnitAwarePolynomialDistribution, isContinuous )
{
  TEST_ASSERT( unit_aware_distribution->isContinuous() );
}

//---------------------------------------------------------------------------//
// Check that the distribution can be written to and read from an xml file
TEUCHOS_UNIT_TEST( PolynomialDistribution, toParameterList )
{
  Teuchos::RCP<Utility::PolynomialDistribution> true_distribution =
    Teuchos::rcp_dynamic_cast<Utility::PolynomialDistribution>( distribution );

  Teuchos::ParameterList parameter_list;

  parameter_list.set<Utility::PolynomialDistribution>( "test_distribution",
						       *true_distribution );

  Teuchos::writeParameterListToXmlFile( parameter_list,
					"polynomial_dist_test_list.xml" );

  Teuchos::RCP<Teuchos::ParameterList> read_parameter_list =
    Teuchos::getParametersFromXmlFile( "polynomial_dist_test_list.xml" );

  TEST_EQUALITY( parameter_list, *read_parameter_list );

  Teuchos::RCP<Utility::PolynomialDistribution>
    copy_distribution( new Utility::PolynomialDistribution );

  *copy_distribution =
    read_parameter_list->get<Utility::PolynomialDistribution>(
							 "test_distribution" );
  TEST_EQUALITY( *copy_distribution, *true_distribution );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be written to and read from an
// xml file
TEUCHOS_UNIT_TEST( UnitAwarePolynomialDistribution, toParameterList )
{
  typedef Utility::UnitAwarePolynomialDistribution<MegaElectronVolt,si::amount> UnitAwarePolynomialDistribution;

  Teuchos::RCP<UnitAwarePolynomialDistribution> true_distribution =
    Teuchos::rcp_dynamic_cast<UnitAwarePolynomialDistribution>( unit_aware_distribution );

  Teuchos::ParameterList parameter_list;

  parameter_list.set<UnitAwarePolynomialDistribution>( "test_distribution",
						       *true_distribution );

  Teuchos::writeParameterListToXmlFile( parameter_list,
					"unit_aware_polynomial_dist_test_list.xml" );

  Teuchos::RCP<Teuchos::ParameterList> read_parameter_list =
    Teuchos::getParametersFromXmlFile( "unit_aware_polynomial_dist_test_list.xml" );

  TEST_EQUALITY( parameter_list, *read_parameter_list );

  Teuchos::RCP<UnitAwarePolynomialDistribution>
    copy_distribution( new UnitAwarePolynomialDistribution );

  *copy_distribution =
    read_parameter_list->get<UnitAwarePolynomialDistribution>(
							 "test_distribution" );
  TEST_EQUALITY( *copy_distribution, *true_distribution );
}

//---------------------------------------------------------------------------//
// Check that the distribution can be read from an xml file
TEUCHOS_UNIT_TEST( PolynomialDistribution, fromParameterList )
{
  Utility::PolynomialDistribution read_distribution =
    test_dists_list->get<Utility::PolynomialDistribution>( "Polynomial Distribution A" );

  TEST_EQUALITY_CONST( read_distribution.getLowerBoundOfIndepVar(), 0.0 );
  TEST_EQUALITY_CONST( read_distribution.getUpperBoundOfIndepVar(), 1.0 );

  read_distribution =
    test_dists_list->get<Utility::PolynomialDistribution>( "Polynomial Distribution B" );

  TEST_EQUALITY_CONST( read_distribution.getLowerBoundOfIndepVar(), 0.0 );
  TEST_EQUALITY_CONST( read_distribution.getUpperBoundOfIndepVar(),
		       Utility::PhysicalConstants::pi/2 );
}

//---------------------------------------------------------------------------//
// Check that the unit-aware distribution can be read from an xml file
TEUCHOS_UNIT_TEST( UnitAwarePolynomialDistribution, fromParameterList )
{
  typedef Utility::UnitAwarePolynomialDistribution<MegaElectronVolt,si::amount> UnitAwarePolynomialDistribution;

  UnitAwarePolynomialDistribution read_distribution =
    test_dists_list->get<UnitAwarePolynomialDistribution>( "Unit-Aware Polynomial Distribution A" );

  TEST_EQUALITY_CONST( read_distribution.getLowerBoundOfIndepVar(), 0.0*MeV );
  TEST_EQUALITY_CONST( read_distribution.getUpperBoundOfIndepVar(), 1.0*MeV );
  TEST_EQUALITY_CONST( read_distribution.evaluate( 0.0*MeV ),
		       1.0*si::mole );
  TEST_EQUALITY_CONST( read_distribution.evaluate( 0.5*MeV ),
		       2.75*si::mole );
  TEST_EQUALITY_CONST( read_distribution.evaluate( 1.0*MeV ),
		       6.0*si::mole );
}

//---------------------------------------------------------------------------//
// Check that a unit-aware distribution can be scaled
TEUCHOS_UNIT_TEST_TEMPLATE_4_DECL( UnitAwarePolynomialDistribution,
				   explicit_conversion,
				   IndepUnitA,
				   DepUnitA,
				   IndepUnitB,
				   DepUnitB )
{
  typedef typename Utility::UnitTraits<IndepUnitA>::template GetQuantityType<double>::type IndepQuantityA;
  typedef typename Utility::UnitTraits<typename Utility::UnitTraits<IndepUnitA>::InverseUnit>::template GetQuantityType<double>::type InverseIndepQuantityA;

  typedef typename Utility::UnitTraits<IndepUnitB>::template GetQuantityType<double>::type IndepQuantityB;
  typedef typename Utility::UnitTraits<typename Utility::UnitTraits<IndepUnitB>::InverseUnit>::template GetQuantityType<double>::type InverseIndepQuantityB;

  typedef typename Utility::UnitTraits<DepUnitA>::template GetQuantityType<double>::type DepQuantityA;
  typedef typename Utility::UnitTraits<DepUnitB>::template GetQuantityType<double>::type DepQuantityB;

  // Copy from a unitless distribution to distribution type A (static method)
  Utility::UnitAwarePolynomialDistribution<IndepUnitA,DepUnitA>
    unit_aware_dist_a_copy = Utility::UnitAwarePolynomialDistribution<IndepUnitA,DepUnitA>::fromUnitlessDistribution( *Teuchos::rcp_dynamic_cast<Utility::PolynomialDistribution>( distribution ) );

  // Copy from distribution type A to distribution type B (explicit cast)
  Utility::UnitAwarePolynomialDistribution<IndepUnitB,DepUnitB>
    unit_aware_dist_b_copy( unit_aware_dist_a_copy );

  IndepQuantityA indep_quantity_a =
    Utility::QuantityTraits<IndepQuantityA>::initializeQuantity( 0.0 );
  InverseIndepQuantityA inv_indep_quantity_a =
    Utility::QuantityTraits<InverseIndepQuantityA>::initializeQuantity( 1/3. );
  DepQuantityA dep_quantity_a =
    Utility::QuantityTraits<DepQuantityA>::initializeQuantity( 1.0 );

  IndepQuantityB indep_quantity_b( indep_quantity_a );
  InverseIndepQuantityB inv_indep_quantity_b( inv_indep_quantity_a );
  DepQuantityB dep_quantity_b( dep_quantity_a );

  UTILITY_TEST_FLOATING_EQUALITY(
			   unit_aware_dist_a_copy.evaluate( indep_quantity_a ),
			   dep_quantity_a,
			   1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			unit_aware_dist_a_copy.evaluatePDF( indep_quantity_a ),
			inv_indep_quantity_a,
			1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			   unit_aware_dist_b_copy.evaluate( indep_quantity_b ),
			   dep_quantity_b,
			   1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			unit_aware_dist_b_copy.evaluatePDF( indep_quantity_b ),
			inv_indep_quantity_b,
			1e-15 );

  Utility::setQuantity( indep_quantity_a, 1.0 );
  Utility::setQuantity( inv_indep_quantity_a, 2.0 );
  Utility::setQuantity( dep_quantity_a, 6.0 );

  indep_quantity_b = IndepQuantityB( indep_quantity_a );
  inv_indep_quantity_b = InverseIndepQuantityB( inv_indep_quantity_a );
  dep_quantity_b = DepQuantityB( dep_quantity_a );

  UTILITY_TEST_FLOATING_EQUALITY(
			   unit_aware_dist_a_copy.evaluate( indep_quantity_a ),
			   dep_quantity_a,
			   1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			unit_aware_dist_a_copy.evaluatePDF( indep_quantity_a ),
			inv_indep_quantity_a,
			1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			   unit_aware_dist_b_copy.evaluate( indep_quantity_b ),
			   dep_quantity_b,
			   1e-15 );
  UTILITY_TEST_FLOATING_EQUALITY(
			unit_aware_dist_b_copy.evaluatePDF( indep_quantity_b ),
			inv_indep_quantity_b,
			1e-15 );
}

typedef si::energy si_energy;
typedef cgs::energy cgs_energy;
typedef si::length si_length;
typedef cgs::length cgs_length;
typedef si::mass si_mass;
typedef cgs::mass cgs_mass;
typedef si::amount si_amount;
typedef si::dimensionless si_dimensionless;
typedef cgs::dimensionless cgs_dimensionless;

TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      si_energy,
				      si_amount,
				      cgs_energy,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      cgs_energy,
				      si_amount,
				      si_energy,
				      si_amount );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      si_length,
				      si_energy,
				      cgs_length,
				      cgs_energy );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      cgs_length,
				      cgs_energy,
				      si_length,
				      si_energy );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      si_length,
				      si_mass,
				      cgs_length,
				      cgs_mass );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      cgs_length,
				      cgs_mass,
				      si_length,
				      si_mass );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      si_length,
				      void,
				      cgs_length,
				      void );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      cgs_length,
				      void,
				      si_length,
				      void );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      void,
				      si_mass,
				      void,
				      cgs_mass );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      void,
				      cgs_mass,
				      void,
				      si_mass );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      si_length,
				      si_dimensionless,
				      cgs_length,
				      cgs_dimensionless );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      cgs_length,
				      cgs_dimensionless,
				      si_length,
				      si_dimensionless );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      si_dimensionless,
				      si_length,
				      cgs_dimensionless,
				      cgs_length );
TEUCHOS_UNIT_TEST_TEMPLATE_4_INSTANT( UnitAwarePolynomialDistribution,
				      explicit_conversion,
				      cgs_dimensionless,
				      cgs_length,
				      si_dimensionless,
				      si_length );

//---------------------------------------------------------------------------//
// Custom main function
//---------------------------------------------------------------------------//
int main( int argc, char** argv )
{
  std::string test_dists_xml_file;

  Teuchos::CommandLineProcessor& clp = Teuchos::UnitTestRepository::getCLP();

  clp.setOption( "test_dists_xml_file",
		 &test_dists_xml_file,
		 "Test distributions xml file name" );

  const Teuchos::RCP<Teuchos::FancyOStream> out =
    Teuchos::VerboseObjectBase::getDefaultOStream();

  Teuchos::CommandLineProcessor::EParseCommandLineReturn parse_return =
    clp.parse(argc,argv);

  if ( parse_return != Teuchos::CommandLineProcessor::PARSE_SUCCESSFUL ) {
    *out << "\nEnd Result: TEST FAILED" << std::endl;
    return parse_return;
  }

  TEUCHOS_ADD_TYPE_CONVERTER( Utility::PolynomialDistribution );
  typedef Utility::UnitAwarePolynomialDistribution<MegaElectronVolt,si::amount> UnitAwarePolynomialDistribution;
  TEUCHOS_ADD_TYPE_CONVERTER( UnitAwarePolynomialDistribution );

  test_dists_list = Teuchos::getParametersFromXmlFile( test_dists_xml_file );

  // Initialize the random number generator
  Utility::RandomNumberGenerator::createStreams();

  // Initialize the polynomial distribution
  Teuchos::Array<double> coeffs( 4 );
  coeffs[0] = 1.0;
  coeffs[1] = 2.0;
  coeffs[2] = 3.0;
  coeffs[3] = 0.0;

  distribution.reset( new Utility::PolynomialDistribution( coeffs, 0.0, 1.0 ));

  // Initialize the unit-aware polynomial distribution
  unit_aware_distribution.reset(
     new Utility::UnitAwarePolynomialDistribution<MegaElectronVolt,si::amount>(
						  coeffs, 0.0*eV, 1e6*eV ) );

  // Run the unit tests
  Teuchos::GlobalMPISession mpiSession( &argc, &argv );

  const bool success = Teuchos::UnitTestRepository::runUnitTests(*out);

  if (success)
    *out << "\nEnd Result: TEST PASSED" << std::endl;
  else
    *out << "\nEnd Result: TEST FAILED" << std::endl;

  clp.printFinalTimerSummary(out.ptr());

  return (success ? 0 : 1);
}

//---------------------------------------------------------------------------//
// end tstPolynomialDistribution.cpp
//---------------------------------------------------------------------------//


#include <boost/test/unit_test.hpp>
#include <memory>
#include "algorithms/data_structures/sll/add_two_numbers_as_linked_lists.hpp"

BOOST_AUTO_TEST_SUITE(AddTwoNumbersAsLinkedList)

BOOST_AUTO_TEST_CASE(empty_lists)
{
    NodeSLL<int> *left = nullptr, *right = nullptr;
    BOOST_CHECK(Algo::DS::SLL::AddNumbersSLL<int>::Add(left, right) == nullptr);

    NodeSLL<int> shortList(0);
    {
        NodeSLL<int>* result =
                Algo::DS::SLL::AddNumbersSLL<int>::Add(left, &shortList);
        BOOST_CHECK(result->value == 0);
        DeleteSLL(&result);
    }

    {
        shortList.value = 5;
        NodeSLL<int>* result =
                Algo::DS::SLL::AddNumbersSLL<int>::Add(&shortList, right);
        BOOST_CHECK(result->value == 5);
        DeleteSLL(&result);
    }
}

BOOST_AUTO_TEST_CASE(equal_size_lists)
{
    NodeSLL<int>* first = CreateSLL<int>({0, 1, 2, 3, 4});
    NodeSLL<int>* second = CreateSLL<int>({8, 7, 9, 9, 4});
    std::vector<int> expected({8, 8, 1, 3, 9});

    NodeSLL<int>* result = Algo::DS::SLL::AddNumbersSLL<int>::Add(first, second);
    BOOST_CHECK(expected == ValuesInSLL(&result));

    DeleteSLL(&first);
    DeleteSLL(&second);
    DeleteSLL(&result);
}

BOOST_AUTO_TEST_CASE(equal_size_lists_with_carry)
{
    NodeSLL<int>* first = CreateSLL<int>({0, 1, 2, 3, 4});
    NodeSLL<int>* second = CreateSLL<int>({8, 7, 9, 9, 9});
    std::vector<int> expected({8, 8, 1, 3, 4, 1});

    NodeSLL<int>* result = Algo::DS::SLL::AddNumbersSLL<int>::Add(first, second);
    BOOST_CHECK(expected == ValuesInSLL(&result));

    DeleteSLL(&first);
    DeleteSLL(&second);
    DeleteSLL(&result);
}

BOOST_AUTO_TEST_CASE(diff_size_lists)
{
    NodeSLL<int>* first = CreateSLL<int>({0, 1, 2});
    NodeSLL<int>* second = CreateSLL<int>({8, 7, 9, 9, 4});
    std::vector<int> expected({8, 8, 1, 0, 5});

    NodeSLL<int>* result = Algo::DS::SLL::AddNumbersSLL<int>::Add(first, second);
    BOOST_CHECK(expected == ValuesInSLL(&result));

    DeleteSLL(&first);
    DeleteSLL(&second);
    DeleteSLL(&result);
}

BOOST_AUTO_TEST_SUITE_END()


#include <boost/hana/list.hpp>
#include <boost/hana/type.hpp>

<%= render('_main.erb') do |f, state, xs|
    <<-EOS
    decltype(
        boost::hana::foldl(
            boost::hana::lift<#{f}::apply>,
            boost::hana::type<#{state}>,
            boost::hana::list_t<#{xs.join(', ')}>
        )
    )::type
    EOS
end %>


/**
 * @file BayesianNetwork.hpp
 * @brief Declaration of the BayesianNetwork class.
 * @author Ankit Srivastava <asrivast@gatech.edu>
 *
 * Copyright 2020 Georgia Institute of Technology
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef BAYESIANNETWORK_HPP_
#define BAYESIANNETWORK_HPP_

#include "graph/Graph.hpp"

#include <boost/graph/filtered_graph.hpp>

/**
 * @brief Class which provides functionality for a Bayesian network.
 *
 * @tparam Var Type of variable indices (expected to be an integer type).
 */
template <typename Var>
class BayesianNetwork : public Graph<BidirectionalAdjacencyList, VertexLabel, Var> {
private:
  class AntiParallelEdgeFilter;
  class EdgeCycleCounter;

public:
  using Vertex = typename Graph<BidirectionalAdjacencyList, VertexLabel, Var>::Vertex;
  using Edge = typename Graph<BidirectionalAdjacencyList, VertexLabel, Var>::Edge;

private:
  using GraphImpl = typename Graph<BidirectionalAdjacencyList, VertexLabel, Var>::GraphImpl;
  using FilteredGraph = Graph<GenericBoostGraph, boost::filtered_graph<GraphImpl, AntiParallelEdgeFilter>, Var>;

public:
  BayesianNetwork(const std::vector<std::string>&);

  using Graph<BidirectionalAdjacencyList, VertexLabel, Var>::addEdge;

  void
  addEdge(const Var, const Var, const bool);

  void
  applyVStructures(std::vector<std::tuple<double, Var, Var, Var>>&&);

  bool
  hasDirectedCycles() const;

  void
  breakDirectedCycles();

  bool
  applyMeekRules();

  void
  writeGraphviz(const std::string&) const;

  ~BayesianNetwork() { }

private:
  Graph<GenericBoostGraph, boost::filtered_graph<GraphImpl, AntiParallelEdgeFilter>, Var>
  filterAntiParallelEdges() const;

  std::unordered_map<Edge, size_t, typename Edge::Hash>
  countEdgeCycles() const;

  bool
  removeEdgeAcyclic(Edge&&);

  bool
  unshieldedColliderRule(const Vertex&, const Vertex&) const;

  bool
  acyclicityRule(const Vertex&, const Vertex&) const;

  bool
  hybridRule(const Vertex&, const Vertex&) const;

private:
  // View of the network with only directed edges
  Graph<GenericBoostGraph, boost::filtered_graph<GraphImpl, AntiParallelEdgeFilter>, Var> m_directed;
}; // class BayesianNetwork

#include "detail/BayesianNetwork.hpp"

#endif // BAYESIANNETWORK_HPP_