#include "StereoSlic.h"
#include <cmath>
#include <float.h>
#include <exception>
#include <Eigen/Dense>

// Default parameter
const int STEREOSLIC_DEFAULT_ITERATION_TOTAL = 10;
const double STEREOSLIC_DEFAULT_COLOR_WEIGHT = 3000.0;
const double STEREOSLIC_DEFAULT_DISPARITY_WEIGHT = 30.0;
const double STEREOSLIC_DEFAULT_NO_DISPARITY_PENALTY = 3.0;

// Pixel offsets of 4- and 8-neighbors
const int fourNeighborTotal = 4;
const int fourNeighborOffsetX[4] = {-1, 0, 1, 0};
const int fourNeighborOffsetY[4] = { 0,-1, 0, 1};
const int eightNeighborTotal = 8;
const int eightNeighborOffsetX[8] = {-1,-1, 0, 1, 1, 1, 0,-1};
const int eightNeighborOffsetY[8] = { 0,-1,-1,-1, 0, 1, 1, 1};

// Prototype declaration
int computeRequiredSamplingTotal(const int drawTotal, const int inlierTotal, const int pointTotal,
                                 const int currentSamplingTotal, const double confidenceLevel);

StereoSlic::StereoSlic() : iterationTotal_(STEREOSLIC_DEFAULT_ITERATION_TOTAL),
                           colorWeight_(STEREOSLIC_DEFAULT_COLOR_WEIGHT),
                           disparityWeight_(STEREOSLIC_DEFAULT_DISPARITY_WEIGHT),
                           noDisparityPenalty_(STEREOSLIC_DEFAULT_NO_DISPARITY_PENALTY) {}

void StereoSlic::setIterationTotal(const int iterationTotal) {
    if (iterationTotal < 1) {
        throw std::runtime_error("StereoSlic: the number of iterations is less than 1");
    }

    iterationTotal_ = iterationTotal;
}

void StereoSlic::setEnergyParameter(const double colorWeight, const double disparityWeight, const double noDisparityPenalty) {
    if (colorWeight <= 0 || disparityWeight < 0 || noDisparityPenalty < 0) {
        throw std::runtime_error("StereoSlic: energy parameter is less than zero");
    }

    colorWeight_ = colorWeight;
    disparityWeight_ = disparityWeight;
    noDisparityPenalty_ = noDisparityPenalty;
}

void StereoSlic::segment(const int superpixelTotal,
                         const png::image<png::rgb_pixel>& leftImage,
                         png::image<png::gray_pixel_16>& segmentImage)
{
    energyType_ = 0;   // Color(left)
    setColorImage(leftImage);
    performSegmentation(superpixelTotal);
    makeSegmentImage(segmentImage);
}

void StereoSlic::segment(const int superpixelTotal,
                         const png::image<png::rgb_pixel>& leftImage,
                         const png::image<png::gray_pixel_16>& leftDisparityImage,
                         png::image<png::gray_pixel_16>& segmentImage,
                         png::image<png::gray_pixel_16>& disparityImage)
{
    energyType_ = 1;   // Color(left) + Disparity(left)
    setColorImage(leftImage);
    setDisparityImage(leftDisparityImage);
    performSegmentation(superpixelTotal);
    makeSegmentImage(segmentImage);
    makeDisparityImage(disparityImage);
}

void StereoSlic::segment(const int superpixelTotal,
                         const png::image<png::rgb_pixel>& leftImage,
                         const png::image<png::gray_pixel_16>& leftDisparityImage,
                         const png::image<png::gray_pixel_16>& rightDisparityImage,
                         png::image<png::gray_pixel_16>& segmentImage,
                         png::image<png::gray_pixel_16>& disparityImage)
{
    energyType_ = 2;   // Color(left) + Disparity(left/right)
    setColorImage(leftImage);
    setDisparityImage(leftDisparityImage, rightDisparityImage);
    performSegmentation(superpixelTotal);
    makeSegmentImage(segmentImage);
    makeDisparityImage(disparityImage);
}

void StereoSlic::segment(const int superpixelTotal,
                         const png::image<png::rgb_pixel>& leftImage,
                         const png::image<png::gray_pixel_16>& leftDisparityImage,
                         const png::image<png::gray_pixel_16>& rightDisparityImage,
                         const png::image<png::rgb_pixel>& rightImage,
                         png::image<png::gray_pixel_16>& segmentImage,
                         png::image<png::gray_pixel_16>& disparityImage)
{
    energyType_ = 3;   // Color(left/right) + Disparity(left/right)
    setColorImage(leftImage, rightImage);
    setDisparityImage(leftDisparityImage, rightDisparityImage);
    performSegmentation(superpixelTotal);
    makeSegmentImage(segmentImage);
    makeDisparityImage(disparityImage);
}

void StereoSlic::setColorImage(const png::image<png::rgb_pixel>& leftImage) {
    width_ = static_cast<int>(leftImage.get_width());
    height_ = static_cast<int>(leftImage.get_height());

    convertRGBToLab(leftImage, leftLabImage_);
    calcLeftLabEdges();
}

void StereoSlic::setColorImage(const png::image<png::rgb_pixel>& leftImage,
                               const png::image<png::rgb_pixel>& rightImage)
{
    width_ = static_cast<int>(leftImage.get_width());
    height_ = static_cast<int>(leftImage.get_height());

    convertRGBToLab(leftImage, leftLabImage_);
    calcLeftLabEdges();
    convertRGBToLab(rightImage, rightLabImage_);
}

void StereoSlic::setDisparityImage(const png::image<png::gray_pixel_16>& leftDisparityImage) {
    leftDisparityImage_.resize(width_*height_);
    for (int y = 0; y < height_; ++y) {
        for (int x = 0; x < width_; ++x) {
            unsigned short pixelValue = leftDisparityImage[y][x];
            if (pixelValue == 0) leftDisparityImage_[width_*y + x] = -1.0;
            else leftDisparityImage_[width_*y + x] = static_cast<float>(pixelValue)/256.0f;
        }
    }
}

void StereoSlic::setDisparityImage(const png::image<png::gray_pixel_16>& leftDisparityImage,
                                   const png::image<png::gray_pixel_16>& rightDisparityImage)
{
    leftDisparityImage_.resize(width_*height_);
    rightDisparityImage_.resize(width_*height_);
    for (int y = 0; y < height_; ++y) {
        for (int x = 0; x < width_; ++x) {
            unsigned short leftPixelValue = leftDisparityImage[y][x];
            if (leftPixelValue == 0) leftDisparityImage_[width_*y + x] = -1.0;
            else leftDisparityImage_[width_*y + x] = static_cast<float>(leftPixelValue)/256.0f;

            unsigned short rightPixelValue = rightDisparityImage[y][x];
            if (rightPixelValue == 0) rightDisparityImage_[width_*y + x] = -1.0;
            else rightDisparityImage_[width_*y + x] = static_cast<float>(rightPixelValue)/256.0f;
        }
    }
}

void StereoSlic::performSegmentation(const int superpixelTotal) {
    initializeSeeds(superpixelTotal);

    for (int iterationCount = 0; iterationCount < iterationTotal_; ++iterationCount) {
        assignLabel();
        updateSeeds();
    }
    enforceLabelConnectivity();
}

void StereoSlic::makeSegmentImage(png::image<png::gray_pixel_16>& segmentImage) const {
    segmentImage.resize(width_, height_);
    for (int y = 0; y < height_; ++y) {
        for (int x = 0; x < width_; ++x) {
            segmentImage[y][x] = labels_[width_*y + x];
        }
    }
}

void StereoSlic::makeDisparityImage(png::image<png::gray_pixel_16>& disparityImage) {
    estimateDisparityPlaneParameter();

    disparityImage.resize(width_, height_);
    for (int y = 0; y < height_; ++y) {
        for (int x = 0; x < width_; ++x) {
            int seedIndex = labels_[width_*y + x];
            double estimatedDisparity = seeds_[seedIndex].disparityPlane[0]*x + seeds_[seedIndex].disparityPlane[1]*y + seeds_[seedIndex].disparityPlane[2];
            if (estimatedDisparity > 0 && estimatedDisparity < 256.0) disparityImage[y][x] = static_cast<unsigned short>(estimatedDisparity*256.0 + 0.5);
            else if (estimatedDisparity >= 256.0) disparityImage[y][x] = 65535;
            else disparityImage[y][x] = 0;
        }
    }
}

void StereoSlic::convertRGBToLab(const png::image<png::rgb_pixel>& rgbImage,
                                 std::vector<float>& labImage)
{
    const int RGB2LABCONVERTER_XYZ_TABLE_SIZE = 1024;
    // CIE standard parameters
    const double epsilon = 0.008856;
    const double kappa = 903.3;
    // Reference white
    const double referenceWhite[3] = {0.950456, 1.0, 1.088754};
    /// Maximum values
    const double maxXYZValues[3] = {0.95047, 1.0, 1.08883};

    std::vector<float> sRGBGammaCorrections(256);
    for (int pixelValue = 0; pixelValue < 256; ++pixelValue) {
        double normalizedValue = pixelValue/255.0;
        double transformedValue = (normalizedValue <= 0.04045) ? normalizedValue/12.92 : pow((normalizedValue+0.055)/1.055, 2.4);

        sRGBGammaCorrections[pixelValue] = transformedValue;
    }

    int tableSize = RGB2LABCONVERTER_XYZ_TABLE_SIZE;
    std::vector<double> xyzTableIndexCoefficients(3);
    xyzTableIndexCoefficients[0] = (tableSize-1)/maxXYZValues[0];
    xyzTableIndexCoefficients[1] = (tableSize-1)/maxXYZValues[1];
    xyzTableIndexCoefficients[2] = (tableSize-1)/maxXYZValues[2];

    std::vector< std::vector<float> > fXYZConversions(3);
    for (int xyzIndex = 0; xyzIndex < 3; ++xyzIndex) {
        fXYZConversions[xyzIndex].resize(tableSize);
        double stepValue = maxXYZValues[xyzIndex]/tableSize;
        for (int tableIndex = 0; tableIndex < tableSize; ++tableIndex) {
            double originalValue = stepValue*tableIndex;
            double normalizedValue = originalValue/referenceWhite[xyzIndex];
            double transformedValue = (normalizedValue > epsilon) ? pow(normalizedValue, 1.0/3.0) : (kappa*normalizedValue + 16.0)/116.0;

            fXYZConversions[xyzIndex][tableIndex] = transformedValue;
        }
    }

    labImage.resize(width_*height_*3);
    for (int y = 0; y < height_; ++y) {
        for (int x = 0; x < width_; ++x) {
            png::rgb_pixel rgbColor = rgbImage[y][x];

            float correctedR = sRGBGammaCorrections[rgbColor.red];
            float correctedG = sRGBGammaCorrections[rgbColor.green];
            float correctedB = sRGBGammaCorrections[rgbColor.blue];

            float xyzColor[3];
            xyzColor[0] = correctedR*0.4124564f + correctedG*0.3575761f + correctedB*0.1804375f;
            xyzColor[1] = correctedR*0.2126729f + correctedG*0.7151522f + correctedB*0.0721750f;
            xyzColor[2] = correctedR*0.0193339f + correctedG*0.1191920f + correctedB*0.9503041f;

            int tableIndexX = static_cast<int>(xyzColor[0]*xyzTableIndexCoefficients[0] + 0.5);
            int tableIndexY = static_cast<int>(xyzColor[1]*xyzTableIndexCoefficients[1] + 0.5);
            int tableIndexZ = static_cast<int>(xyzColor[2]*xyzTableIndexCoefficients[2] + 0.5);

            float fX = fXYZConversions[0][tableIndexX];
            float fY = fXYZConversions[1][tableIndexY];
            float fZ = fXYZConversions[2][tableIndexZ];

            labImage[width_*3*y + 3*x] = 116.0*fY - 16.0;
            labImage[width_*3*y + 3*x + 1] = 500.0*(fX - fY);
            labImage[width_*3*y + 3*x + 2] = 200.0*(fY - fZ);
        }
    }
}

void StereoSlic::calcLeftLabEdges() {
    leftLabEdges_.resize(width_*height_, 0);
    for (int y = 1; y < height_-1; ++y) {
        for (int x = 1; x < width_-1; ++x) {
            double dxL = leftLabImage_[width_*3*y + 3*(x-1)] - leftLabImage_[width_*3*y + 3*(x+1)];
            double dxA = leftLabImage_[width_*3*y + 3*(x-1) + 1] - leftLabImage_[width_*3*y + 3*(x+1) + 1];
            double dxB = leftLabImage_[width_*3*y + 3*(x-1) + 2] - leftLabImage_[width_*3*y + 3*(x+1) + 2];
            double dx = dxL*dxL + dxA*dxA + dxB*dxB;

            double dyL = leftLabImage_[width_*3*(y-1) + 3*x] - leftLabImage_[width_*3*(y+1) + 3*x];
            double dyA = leftLabImage_[width_*3*(y-1) + 3*x + 1] - leftLabImage_[width_*3*(y+1) + 3*x + 1];
            double dyB = leftLabImage_[width_*3*(y-1) + 3*x + 2] - leftLabImage_[width_*3*(y+1) + 3*x + 2];
            double dy = dyL*dyL + dyA*dyA + dyB*dyB;

            leftLabEdges_[width_*y + x] = dx + dy;
        }
    }
}

void StereoSlic::initializeSeeds(const int superpixelTotal) {
    int imageSize = width_*height_;
    gridSize_ = sqrt(static_cast<double>(imageSize)/superpixelTotal);
    stepSize_ = static_cast<int>(gridSize_ + 2.0);
    int offsetX = static_cast<int>(gridSize_/2.0);
    int offsetY = static_cast<int>(gridSize_/2.0);

    seeds_.clear();
    for (int indexY = 0; indexY < height_; ++indexY) {
        int y = static_cast<int>(gridSize_*indexY + offsetY + 0.5);
        if (y >= height_) break;
        for (int indexX = 0; indexX < width_; ++indexX) {
            int x = static_cast<int>(gridSize_*indexX + offsetX + 0.5);
            if (x >= width_) break;

            LabXYD newSeed;
            newSeed.color[0] = leftLabImage_[width_*3*y + 3*x];
            newSeed.color[1] = leftLabImage_[width_*3*y + 3*x + 1];
            newSeed.color[2] = leftLabImage_[width_*3*y + 3*x + 2];
            newSeed.position[0] = x;
            newSeed.position[1] = y;
            seeds_.push_back(newSeed);
        }
    }

    int seedTotal = static_cast<int>(seeds_.size());
    for (int seedIndex = 0; seedIndex < seedTotal; ++seedIndex) {
        int originalX = seeds_[seedIndex].position[0];
        int originalY = seeds_[seedIndex].position[1];
        int originalPixelIndex = width_*originalY + originalX;

        int perturbedPixelIndex = originalPixelIndex;
        for (int neighborIndex = 0; neighborIndex < eightNeighborTotal; ++neighborIndex) {
            int neighborX = originalX + eightNeighborOffsetX[neighborIndex];
            int neighborY = originalY + eightNeighborOffsetY[neighborIndex];
            if (neighborX < 0 || neighborX >= width_ || neighborY < 0 || neighborY >= height_) continue;
            int neighborPixelIndex = width_*neighborY + neighborX;

            if (leftLabEdges_[neighborPixelIndex] < leftLabEdges_[perturbedPixelIndex]) {
                perturbedPixelIndex = neighborPixelIndex;
            }
        }

        if (perturbedPixelIndex != originalPixelIndex) {
            int perturbedX = perturbedPixelIndex%width_;
            int perturbedY = perturbedPixelIndex/width_;

            seeds_[seedIndex].color[0] = leftLabImage_[width_*3*perturbedY + 3*perturbedX];
            seeds_[seedIndex].color[1] = leftLabImage_[width_*3*perturbedY + 3*perturbedX + 1];
            seeds_[seedIndex].color[2] = leftLabImage_[width_*3*perturbedY + 3*perturbedX + 2];
            seeds_[seedIndex].position[0] = perturbedX;
            seeds_[seedIndex].position[1] = perturbedY;
        }
    }
}

void StereoSlic::assignLabel() {
    int seedTotal = static_cast<int>(seeds_.size());

    double positionWeight = colorWeight_/(stepSize_*stepSize_);

    labels_.resize(width_*height_, -1);
    std::vector<double> distancesToSeeds(width_*height_, DBL_MAX);
    for (int seedIndex = 0; seedIndex < seedTotal; ++seedIndex) {
        int minX = (seeds_[seedIndex].position[0] > stepSize_) ? seeds_[seedIndex].position[0] - stepSize_ : 0;
        int maxX = (seeds_[seedIndex].position[0] < width_ - stepSize_) ? seeds_[seedIndex].position[0] + stepSize_ : width_;
        int minY = (seeds_[seedIndex].position[1] > stepSize_) ? seeds_[seedIndex].position[1] - stepSize_ : 0;
        int maxY = (seeds_[seedIndex].position[1] < height_ - stepSize_) ? seeds_[seedIndex].position[1] + stepSize_ : height_;

        double seedL = seeds_[seedIndex].color[0];
        double seedA = seeds_[seedIndex].color[1];
        double seedB = seeds_[seedIndex].color[2];
        double seedX = seeds_[seedIndex].position[0];
        double seedY = seeds_[seedIndex].position[1];
        double seedAlpha = seeds_[seedIndex].disparityPlane[0];
        double seedBeta = seeds_[seedIndex].disparityPlane[1];
        double seedGamma = seeds_[seedIndex].disparityPlane[2];

        for (int y = minY; y < maxY; ++y) {
            for (int x = minX; x < maxX; ++x) {
                float leftPixelL = leftLabImage_[width_*3*y + 3*x];
                float leftPixelA = leftLabImage_[width_*3*y + 3*x + 1];
                float leftPixelB = leftLabImage_[width_*3*y + 3*x + 2];
                double distanceLeftLab = (leftPixelL - seedL)*(leftPixelL - seedL)
                    + (leftPixelA - seedA)*(leftPixelA - seedA)
                    + (leftPixelB - seedB)*(leftPixelB - seedB);
                double distanceXY = (x - seedX)*(x - seedX) + (y - seedY)*(y - seedY);

                // Default distances
                double distanceRightLab = distanceLeftLab;
                double distanceLeftD = noDisparityPenalty_;
                double distanceRightD = noDisparityPenalty_;

                if (energyType_ > 0) {
                    double estimatedDisparity = seedAlpha*x + seedBeta*y + seedGamma;
                    if (estimatedDisparity > 0) {
                        if (leftDisparityImage_[width_*y + x] > 0) {
                            // distance of left disparities
                            double leftDisparityDifference = leftDisparityImage_[width_*y + x] - estimatedDisparity;
                            distanceLeftD = leftDisparityDifference*leftDisparityDifference;
                        }

                        if (energyType_ > 1) {
                            int rightX = static_cast<int>(x - estimatedDisparity + 0.5);
                            if (rightX >= 0 && rightDisparityImage_[width_*y + rightX] > 0) {
                                double rightDisparityDifference = estimatedDisparity - rightDisparityImage_[width_*y + rightX];
                                if (rightDisparityDifference < -1) {
                                    // Occluded
                                    rightDisparityDifference = 0;
                                    distanceRightLab = distanceLeftLab;
                                } else {
                                    // Non-occluded
                                    distanceRightD = rightDisparityDifference*rightDisparityDifference;
                                    if (energyType_ > 2) {
                                        float rightPixelL = rightLabImage_[width_*3*y + 3*rightX];
                                        float rightPixelA = rightLabImage_[width_*3*y + 3*rightX + 1];
                                        float rightPixelB = rightLabImage_[width_*3*y + 3*rightX + 2];
                                        distanceRightLab = (rightPixelL - seedL)*(rightPixelL - seedL)
                                            + (rightPixelA - seedA)*(rightPixelA - seedA)
                                            + (rightPixelB - seedB)*(rightPixelB - seedB);
                                    }
                                }
                            }
                        }
                    }
                }

                double distanceLab = distanceLeftLab;
                if (energyType_ == 3) distanceLab += distanceRightLab;
                double distanceD = 0;
                if (energyType_ > 0) distanceD = distanceLeftD;
                if (energyType_ > 1) distanceD += distanceRightD;

                double distanceLabXYD = distanceLab + positionWeight*distanceXY + disparityWeight_*distanceD;
                if (distanceLabXYD < distancesToSeeds[width_*y + x]) {
                    distancesToSeeds[width_*y + x] = distanceLabXYD;
                    labels_[width_*y + x] = seedIndex;
                }
            }
        }
    }
}

void StereoSlic::updateSeeds() {
    updateSeedsColorAndPosition();
    if (energyType_ > 0) {
        estimateDisparityPlaneParameter();
    }
}

void StereoSlic::updateSeedsColorAndPosition() {
    int seedTotal = static_cast<int>(seeds_.size());

    std::vector<int> segmentSizes(seedTotal, 0);
    std::vector<LabXYD> segmentSigmas(seedTotal);
    for (int y = 0; y < height_; ++y) {
        for (int x = 0; x < width_; ++x) {
            int pixelSeedIndex = labels_[width_*y + x];
            if (pixelSeedIndex >= 0) {
                segmentSigmas[pixelSeedIndex].color[0] += leftLabImage_[width_*3*y + 3*x];
                segmentSigmas[pixelSeedIndex].color[1] += leftLabImage_[width_*3*y + 3*x + 1];
                segmentSigmas[pixelSeedIndex].color[2] += leftLabImage_[width_*3*y + 3*x + 2];
                segmentSigmas[pixelSeedIndex].position[0] += x;
                segmentSigmas[pixelSeedIndex].position[1] += y;
                ++segmentSizes[pixelSeedIndex];
            }
        }
    }

    for (int seedIndex = 0; seedIndex < seedTotal; ++seedIndex) {
        double segmentSizeInverse = 1.0;
        if (segmentSizes[seedIndex] > 0) segmentSizeInverse = 1.0/segmentSizes[seedIndex];

        seeds_[seedIndex].color[0] = segmentSigmas[seedIndex].color[0]*segmentSizeInverse;
        seeds_[seedIndex].color[1] = segmentSigmas[seedIndex].color[1]*segmentSizeInverse;
        seeds_[seedIndex].color[2] = segmentSigmas[seedIndex].color[2]*segmentSizeInverse;
        seeds_[seedIndex].position[0] = segmentSigmas[seedIndex].position[0]*segmentSizeInverse;
        seeds_[seedIndex].position[1] = segmentSigmas[seedIndex].position[1]*segmentSizeInverse;
    }
}

void StereoSlic::estimateDisparityPlaneParameter() {
    std::vector< std::vector<DisparityPixel> > segmentDisparityPixels = makeDisparityPixelList();
    int seedTotal = static_cast<int>(segmentDisparityPixels.size());

    for (int seedIndex = 0; seedIndex < seedTotal; ++seedIndex) {
        std::vector<double> planeParameter = estimateDisparityPlaneParameter(segmentDisparityPixels[seedIndex]);

        seeds_[seedIndex].disparityPlane[0] = planeParameter[0];
        seeds_[seedIndex].disparityPlane[1] = planeParameter[1];
        seeds_[seedIndex].disparityPlane[2] = planeParameter[2];
    }
}

std::vector< std::vector<StereoSlic::DisparityPixel> > StereoSlic::makeDisparityPixelList() const {
    int seedTotal = static_cast<int>(seeds_.size());

    std::vector< std::vector<DisparityPixel> > segmentDisparityPixels(seedTotal);
    for (int y = 0; y < height_; ++y) {
        for (int x = 0; x < width_; ++x) {
            if (leftDisparityImage_[width_*y + x] > 0) {
                int pixelSeedIndex = labels_[width_*y + x];
                if (pixelSeedIndex >= 0) {
                    DisparityPixel newDisparityPixel;
                    newDisparityPixel.x = x;
                    newDisparityPixel.y = y;
                    newDisparityPixel.d = leftDisparityImage_[width_*y + x];
                    segmentDisparityPixels[pixelSeedIndex].push_back(newDisparityPixel);
                }
            }
        }
    }

    return segmentDisparityPixels;
}

std::vector<double> StereoSlic::estimateDisparityPlaneParameter(const std::vector<DisparityPixel>& disparityPixels) const {
    std::vector<double> planeParameter(3);

    int pixelTotal = static_cast<int>(disparityPixels.size());
    if (pixelTotal < 3) {
        planeParameter[0] = 0;
        planeParameter[1] = 0;
        planeParameter[2] = -1;
    } else {
        bool isSameX = true;
        bool isSameY = true;
        for (int pixelIndex = 1; pixelIndex < pixelTotal; ++pixelIndex) {
            if (disparityPixels[pixelIndex].x != disparityPixels[0].x) isSameX = false;
            if (disparityPixels[pixelIndex].y != disparityPixels[0].y) isSameY = false;
            if (!isSameX && !isSameY) break;
        }
        if (isSameX || isSameY) {
            double disparitySum = 0.0;
            for (int pixelIndex = 0; pixelIndex < pixelTotal; ++pixelIndex) disparitySum += disparityPixels[pixelIndex].d;

            planeParameter[0] = 0;
            planeParameter[1] = 0;
            planeParameter[2] = -1;
        } else {
            planeParameter = estimateDisparityPlaneParameterRANSAC(disparityPixels);
        }
    }

    return planeParameter;
}

std::vector<double> StereoSlic::estimateDisparityPlaneParameterRANSAC(const std::vector<DisparityPixel>& disparityPixels) const {
    const double inlierThreshold = 2.0;
    const double confidenceLevel = 0.99;

    int pixelTotal = static_cast<int>(disparityPixels.size());

    int samplingTotal = pixelTotal*2;

    int bestInlierTotal = 0;
    std::vector<bool> bestInlierFlags(pixelTotal);
    int samplingCount = 0;
    while (samplingCount < samplingTotal) {
        // Randomly select 3 pixels
        int drawIndices[3];
        drawIndices[0] = rand()%pixelTotal;
        drawIndices[1] = rand()%pixelTotal;
        while(drawIndices[1] == drawIndices[0]) drawIndices[1] = rand()%pixelTotal;
        drawIndices[2] = rand()%pixelTotal;
        while(drawIndices[2] == drawIndices[1] || drawIndices[2] == drawIndices[0]
              || (disparityPixels[drawIndices[0]].x == disparityPixels[drawIndices[1]].x && disparityPixels[drawIndices[0]].x == disparityPixels[drawIndices[2]].x)
              || (disparityPixels[drawIndices[0]].y == disparityPixels[drawIndices[1]].y && disparityPixels[drawIndices[0]].y == disparityPixels[drawIndices[2]].y))
        {
            drawIndices[2] = rand()%pixelTotal;
        }

        // Compute plane parameters
        Eigen::Matrix3d matPosition;
        Eigen::Vector3d vecDisparity;
        for (int i = 0; i < 3; ++i) {
            matPosition(i, 0) = disparityPixels[drawIndices[i]].x;
            matPosition(i, 1) = disparityPixels[drawIndices[i]].y;
            matPosition(i, 2) = 1.0;
            vecDisparity(i) = disparityPixels[drawIndices[i]].d;
        }
        Eigen::Vector3d planeParameter = matPosition.colPivHouseholderQr().solve(vecDisparity);

        // Count the number of inliers
        int inlierTotal = 0;
        std::vector<bool> inlierFlags(pixelTotal);
        for (int pixelIndex = 0; pixelIndex < pixelTotal; ++pixelIndex) {
            double estimatedDisparity = planeParameter(0)*disparityPixels[pixelIndex].x + planeParameter(1)*disparityPixels[pixelIndex].y + planeParameter(2);
            if (fabs(estimatedDisparity - disparityPixels[pixelIndex].d) <= inlierThreshold) {
                ++inlierTotal;
                inlierFlags[pixelIndex] = true;
            } else {
                inlierFlags[pixelIndex] = false;
            }
        }

        // Update best inliers
        if (inlierTotal > bestInlierTotal) {
            bestInlierTotal = inlierTotal;
            bestInlierFlags = inlierFlags;

            samplingTotal = computeRequiredSamplingTotal(3, bestInlierTotal, pixelTotal, samplingTotal, confidenceLevel);
        }

        // Increment
        ++samplingCount;
    }

    // Least-square estimation
    Eigen::MatrixXd matPosition(bestInlierTotal, 3);
    Eigen::VectorXd vecDisparity(bestInlierTotal);
    int inlierIndex = 0;
    for (int pixelIndex = 0; pixelIndex < pixelTotal; ++pixelIndex) {
        if (bestInlierFlags[pixelIndex]) {
            matPosition(inlierIndex, 0) = disparityPixels[pixelIndex].x;
            matPosition(inlierIndex, 1) = disparityPixels[pixelIndex].y;
            matPosition(inlierIndex, 2) = 1.0;
            vecDisparity(inlierIndex) = disparityPixels[pixelIndex].d;
            ++inlierIndex;
        }
    }
    Eigen::Vector3d vecPlaneParameter = matPosition.colPivHouseholderQr().solve(vecDisparity);
    std::vector<double> planeParameter(3);
    planeParameter[0] = vecPlaneParameter(0);
    planeParameter[1] = vecPlaneParameter(1);
    planeParameter[2] = vecPlaneParameter(2);

    return planeParameter;
}

void StereoSlic::enforceLabelConnectivity() {
    int imageSize = width_*height_;
    int seedTotal = static_cast<int>(seeds_.size());

    const int minimumSegmentSizeThreshold = imageSize/seedTotal/4;

    std::vector<int> newLabels(width_*height_, -1);
    int newLabelIndex = 0;
    for (int y = 0; y < height_; ++y) {
        for (int x = 0; x < width_; ++x) {
            if (newLabels[width_*y + x] >= 0) continue;

            newLabels[width_*y + x] = newLabelIndex;

            int adjacentLabel = 0;
            for (int neighborIndex = 0; neighborIndex < fourNeighborTotal; ++neighborIndex) {
                int neighborX = x + fourNeighborOffsetX[neighborIndex];
                int neighborY = y + fourNeighborOffsetY[neighborIndex];
                if (neighborX < 0 || neighborX >= width_ || neighborY < 0 || neighborY >= height_) continue;

                if (newLabels[width_*neighborY + neighborX] >= 0) adjacentLabel = newLabels[width_*neighborY + neighborX];
            }

            std::vector<int> connectedXs(1);
            std::vector<int> connectedYs(1);
            connectedXs[0] = x;
            connectedYs[0] = y;
            labelConnectedPixels(x, y, newLabelIndex, newLabels, connectedXs, connectedYs);

            int segmentPixelTotal = static_cast<int>(connectedXs.size());
            if (segmentPixelTotal <= minimumSegmentSizeThreshold) {
                for (int i = 0; i < segmentPixelTotal; ++i) {
                    newLabels[width_*connectedYs[i] + connectedXs[i]] = adjacentLabel;
                }
                --newLabelIndex;
            }

            ++newLabelIndex;
        }
    }

    labels_ = newLabels;
    seeds_.resize(newLabelIndex);
}

void StereoSlic::labelConnectedPixels(const int x, const int y, const int newLabelIndex,
                                      std::vector<int>& newLabels, std::vector<int>& connectedXs, std::vector<int>& connectedYs) const
{
    int originalLabelIndex = labels_[width_*y + x];
    for (int neighborIndex = 0; neighborIndex < fourNeighborTotal; ++neighborIndex) {
        int neighborX = x + fourNeighborOffsetX[neighborIndex];
        int neighborY = y + fourNeighborOffsetY[neighborIndex];
        if (neighborX < 0 || neighborX >= width_ || neighborY < 0 || neighborY >= height_) continue;

        if (newLabels[width_*neighborY + neighborX] < 0 && labels_[width_*neighborY + neighborX] == originalLabelIndex) {
            connectedXs.push_back(neighborX);
            connectedYs.push_back(neighborY);
            newLabels[width_*neighborY + neighborX] = newLabelIndex;
            labelConnectedPixels(neighborX, neighborY, newLabelIndex, newLabels, connectedXs, connectedYs);
        }
    }
}

int computeRequiredSamplingTotal(const int drawTotal, const int inlierTotal, const int pointTotal, const int currentSamplingTotal, const double confidenceLevel) {
    double ep = 1 - static_cast<double>(inlierTotal)/static_cast<double>(pointTotal);
    if (ep == 1.0) {
        ep = 0.5;
    }

    int newSamplingTotal = static_cast<int>(log(1 - confidenceLevel)/log(1 - pow(1 - ep, drawTotal)) + 0.5);
    if (newSamplingTotal < currentSamplingTotal) {
        return newSamplingTotal;
    } else {
        return currentSamplingTotal;
    }
}

png::image<png::rgb_pixel> drawSegmentBoundary(const png::image<png::rgb_pixel>& originalImage,
                                               const png::image<png::gray_pixel_16>& segmentImage,
                                               const png::rgb_pixel boundaryColor)
{
    // Pixel offsets of 8-neighbors
    const int eightNeighborTotal = 8;
    const int eightNeighborOffsetX[8] = {-1,-1, 0, 1, 1, 1, 0,-1};
    const int eightNeighborOffsetY[8] = { 0,-1,-1,-1, 0, 1, 1, 1};

    // Check image size
    int width = static_cast<int>(originalImage.get_width());
    int height = static_cast<int>(originalImage.get_height());
    if (segmentImage.get_width() != width || segmentImage.get_height() != height) {
        throw std::runtime_error("drawSegmentBoundary: image sizes are different");
    }

    // Copy input image
    png::image<png::rgb_pixel> boundaryImage(originalImage);

    // Draw boundary
    std::vector<bool> boundaryFlags(width*height, false);
    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            int pixelLabelIndex = segmentImage[y][x];

            bool drawBoundary = false;
            for (int neighborIndex = 0; neighborIndex < eightNeighborTotal; ++neighborIndex) {
                int neighborX = x + eightNeighborOffsetX[neighborIndex];
                int neighborY = y + eightNeighborOffsetY[neighborIndex];
                if (neighborX < 0 || neighborX >= width || neighborY < 0 || neighborY >= height) continue;

                int neighborPixelIndex = width*neighborY + neighborX;
                if (boundaryFlags[neighborPixelIndex]) continue;
                if (segmentImage[neighborY][neighborX] != pixelLabelIndex) {
                    drawBoundary = true;
                    break;
                }
            }

            if (drawBoundary) {
                boundaryImage[y][x] = boundaryColor;
            }
        }
    }

    return boundaryImage;
}


#ifndef STAN_MATH_TORSTEN_EVENT_HISTORY_HPP
#define STAN_MATH_TORSTEN_EVENT_HISTORY_HPP

#include <iomanip>
#include <stan/math/prim/fun/value_of.hpp>
#include <stan/math/rev/fun/value_of.hpp>
#include <stan/math/prim/meta/return_type.hpp>
#include <stan/math/prim/err/check_greater_or_equal.hpp>
#include <stan/math/torsten/PKModel/functions.hpp>
#include <stan/math/torsten/pk_nsys.hpp>
#include <Eigen/Dense>
#include <iostream>
#include <algorithm>
#include <numeric>
#include <vector>
#include <utility>
#include <functional>

namespace torsten {
  /**
   * Implementation of Non event parameters structure for
   * bioavailability, lag time, real ODEdata, integer ODE data. Since
   * we allow elision of any of these we need template discern the
   * combination. The default is that the param pack has both
   * bioavailability and lag time
   *
   * @tparam Ts types of non-event parameters. Each in the parameter
   * pack is a 2d array.
   */
  template<typename... Ts>
  struct NonEventParameters_Impl {
    /// total # of params plus one for theta
    static constexpr int npar = sizeof...(Ts) + 1;
    using biovar_t = std::tuple_element_t<0, std::tuple<Ts...> >;
    using lag_t = std::tuple_element_t<1, std::tuple<Ts...> >;

    static const auto& bioavailability(int i, int j,
                                       const std::tuple<const std::vector<std::vector<Ts> >&...> array_2d_params) {
      return std::get<0>(array_2d_params).at(i).at(j);
    }
    static const auto& lag_time(int i, int j,
                                const std::tuple<const std::vector<std::vector<Ts> >&...> array_2d_params) {
      return std::get<1>(array_2d_params).at(i).at(j);
    }
  };

  /**
   * Specialization: there's bioavailabity but no lag time.
   */
  template<typename T>
  struct NonEventParameters_Impl<T> {
    static constexpr int npar = 2;
    using biovar_t = T;
    using lag_t = double;

    static const auto& bioavailability(int i, int j,
                                       const std::tuple<const std::vector<std::vector<T> >&> array_2d_params) {
      return std::get<0>(array_2d_params).at(i).at(j);
    }
    static double lag_time(int i, int j,
                                const std::tuple<const std::vector<std::vector<T> >&> array_2d_params) {
      return 0.0;
    }
  };

  /**
   * Specialization: there's neither bioavailabity but nor lag time.
   */
  template<>
  struct NonEventParameters_Impl<> {
    static constexpr int npar = 1;
    using biovar_t = double;
    using lag_t = double;

    static double bioavailability(int i, int j,
                                  const std::tuple<> array_2d_params) {
      return 1.0;
    }
    static double lag_time(int i, int j,
                                const std::tuple<> array_2d_params) {
      return 0.0;
    }
  };

  /**
   * Parameters & controls that are not part of event sequence,
   * including
   * theta,
   * bioavailability,
   * tlag(lag time),
   * x_r(real data for ODE functor),
   * x_i(integer data for ODE functor)
   *
   * except theta, the above parameters are optional and
   * are handled by parameter packs. Here two packas are used for
   * different purposes:
   * 1. bioavailability and lag time are related to PMX events and
   *    grouped into a tuple. One can specify none,
   *    bioavailability only, or bioavailability + tlag.
   * 2. ODE's real and int data are grouped into a pack. One can specify none,
   *    real data only, or real + int data.
   *
   * @tparam T0 type for time
   * @tparam T4 type for system parameter, usually refered as pMatrix or theta.
   * @tparam theta_container type of container,  std::vector or
   *                         Eigen::Matrix(for linear * system)
   * @tparam params_tuple_type tuple with a parameter pack for PMX events: bioavailability and lag time.
   * @tparam Ts pack of types for ODE data: real & integer.
   *
   */
  template<typename T0,
           typename T4,
           template<typename...> class theta_container,
           typename params_tuple_type,
           typename... Ts>
  struct NonEventParameters;

  /**
   * Parameters & controls that are not part of event sequence,
   * specialization when types for both bioavailability and lag time
   * are provided in tuple_pars_t.
   **/
  template<typename T0,
           typename T4,
           template<typename...> class theta_container,
           typename... tuple_pars_t,
           typename... Ts>
  struct NonEventParameters<T0, T4, theta_container,
                            std::tuple<tuple_pars_t...>, Ts...> {
    static constexpr int npar = NonEventParameters_Impl<tuple_pars_t..., Ts...>::npar;
    /// time & index for an event entry's parameters
    using par_t = std::pair<double, std::array<int, npar> >;
    using biovar_t = typename NonEventParameters_Impl<tuple_pars_t...>::biovar_t;
    using lag_t = typename NonEventParameters_Impl<tuple_pars_t...>::lag_t;
    using T5 = biovar_t;
    using T6 = lag_t;

    /// mapping between time and corresponding parameter index
    std::vector<par_t> pars;
    const std::vector<T0>& time_;
    const std::vector<theta_container<T4>>& theta_;
    /// put parameter pack for bioavailability and lag time into tuple for later retrieval
    const std::tuple<const std::vector<std::vector<tuple_pars_t> >&...> event_ctrl;
    /// put parameter pack for ODE's real & integer data into tuple for later retrieval
    const std::tuple<const std::vector<std::vector<Ts> >&...> ode_data;

    template <typename rec_t>
    NonEventParameters(int id, const rec_t& rec,
                       const std::vector<theta_container<T4>>& theta,
                       const std::vector<std::vector<tuple_pars_t> >&... event_ctrl0,
                       const std::vector<std::vector<Ts> >&... ode_data0) :
      pars(rec.len_[id]),
      time_(rec.time_),
      theta_(theta),
      event_ctrl{std::forward_as_tuple(event_ctrl0...)},
      ode_data{std::forward_as_tuple(ode_data0...)} {
      int ibegin = rec.begin_[id];
      for (int i = 0; i < rec.len_[id]; ++i) {
        int theta_i = rec.len_param(id, theta) > 1 ? rec.begin_param(id, theta) + i : rec.begin_param(id, theta);
        pars[i] = std::make_pair<double, std::array<int, npar> >(double(stan::math::value_of(time_[ibegin + i])),
          {theta_i,index_param(id,i,rec,event_ctrl0)..., index_param(id,i,rec,ode_data0)...});
      }
      sort();
    }

    /**
     * For unit test.
     *
     * @param id  subject id
     * @param rec nonmen event record
     * @param ibegin_theta beginning index in popultion array for subject id
     * @param isize_theta length in popultion array for subject id
     * @param ibegin_biovar beginning index in popultion array for subject id
     * @param isize_biovar length in popultion array for subject id
     * @param ibegin_tlag beginning index in popultion array for subject id
     * @param isize_tlag length in popultion array for subject id
     * @param theta model param
     * @param array_2d_params0 additional params
     *
     */
    // template <typename rec_t, std::enable_if_t<sizeof...(Ts) == 2>* = nullptr>
    template <typename rec_t>
    NonEventParameters(int id, const rec_t& rec,
                       int ibegin_theta, int isize_theta,
                       int ibegin_biovar, int isize_biovar,
                       int ibegin_tlag, int isize_tlag,
                       const std::vector<theta_container<T4>>& theta,
                       const std::vector<std::vector<tuple_pars_t> >&... event_ctrl0,
                       const std::vector<std::vector<Ts> >&... ode_data0) :
      pars(rec.len_[id]),
      time_(rec.time_),
      theta_(theta),
      event_ctrl{std::forward_as_tuple(event_ctrl0...)},
      ode_data{std::forward_as_tuple(ode_data0...)} {
      int ibegin = rec.begin_[id];
      for (int i = 0; i < rec.len_[id]; ++i) {
        int j = isize_theta   > 1 ? ibegin_theta  + i : ibegin_theta;
        int k = isize_biovar  > 1 ? ibegin_biovar + i : ibegin_biovar;
        int l = isize_tlag    > 1 ? ibegin_tlag   + i : ibegin_tlag;
        pars[i] = std::make_pair<double, std::array<int, npar> >(double(stan::math::value_of(time_[ibegin + i])), {j, k, l });
      }
      sort();
    }

    inline void set_par_time(int i, double t) {
      std::get<0>(pars[i]) = t;
    }

    inline void set_par_array(int i, const std::array<int, npar>& a) {
      std::get<1>(pars[i]) = a;
    }

    inline double get_par_time(int i) const {
      return std::get<0>(pars[i]);
    }

    inline const std::array<int, npar>& get_par_array(int i) const {
      return std::get<1>(pars[i]);
    }

    inline const theta_container<T4>& theta(int i) const {
      return theta_[std::get<0>(pars[i].second)];
    }

    template<size_t Is>
    inline auto& get_model_array_1d_param(int i) const {
      using Tuple = std::tuple<const std::vector<std::vector<Ts> >&...>;
      constexpr size_t Is_param = Is + NonEventParameters_Impl<tuple_pars_t...>::npar;
      return std::get<Is>(ode_data)[std::get<Is_param>(pars[i].second)];
    }

    inline const T5 bioavailability(int iEvent, int iParameter) const {
      return NonEventParameters_Impl<tuple_pars_t...>::bioavailability(get_par_array(iEvent)[1], iParameter,
                                                                       event_ctrl);
    }

    inline const T6 lag_time(int iEvent, int iParameter) const {
      return NonEventParameters_Impl<tuple_pars_t...>::lag_time(get_par_array(iEvent)[2], iParameter,
                                                                event_ctrl);
    }

    inline int size() { return pars.size(); }

    void sort() {
      std::sort(pars.begin(), pars.end(),
                [](const par_t& a, const par_t& b)
                { return a.first  0 && ordered) {
        ordered = (pars[i].first >= pars[i-1].first);
        i--;
      }
      return ordered;
    }
  };

  /**
   * find i'th param index for subject id
   *
   * @param id subject index
   * @param i param index for this subject
   * @param rec populatoin event record
   * @param param param for entire population
   *
   * @return i'th param index for subject id
   */
  template<typename rec_t, typename T>
  inline int index_param(int id, int i, const rec_t& rec,
                         const std::vector<std::vector<T> >& param) {
    return rec.len_param(id, param) > 1 ? rec.begin_param(id, param) + i : rec.begin_param(id, param);
  }

  /**
   * The EventHistory class defines objects that contain a vector of Events,
   * along with a series of functions that operate on them.
   */
  template<typename T0, typename T1, typename T2, typename T3, typename T_lag>
  struct EventHistory {
    // using T_scalar = typename stan::return_type_t<T0, T1, T2, T3, T4, T5, T6>;
    using T_time = typename stan::return_type_t<T0, T1, T2, T3, T_lag>;
    // using T_rate = typename stan::return_type_t<T2, T5>;
    // using T_amt = typename stan::return_type_t<T1, T5>;
    /// <time, <theta index, F index, lag index> >
    using Param = std::pair<double, std::array<int, 3> >;
    using rate_t = std::pair<double, std::vector<T2> >;

    const std::vector<T0>& time_;
    const std::vector<T1>& amt_;
    const std::vector<T2>& rate_;
    const std::vector<T3>& ii_;
    const std::vector<int>& evid_;
    const std::vector<int>& cmt_;
    const std::vector<int>& addl_;
    const std::vector<int>& ss_;

    // internally generated events
    // these events are not from user input but required for event
    // specification. Such examples include those from tlag or
    // infusion dosing.
    const size_t num_event_times;
    std::vector<T_time> gen_time;
    std::vector<T1> gen_amt;
    std::vector<T2> gen_rate;
    std::vector<T3> gen_ii;
    std::vector<int> gen_cmt;
    std::vector<int> gen_addl;
    std::vector<int> gen_ss;

    // 0: original(0)/generated(1)
    // 1: index in original/generated arrays
    // 2: evid
    // 3: is new?(0/1)
    using IDVec = std::array<int, 4>;
    std::vector<IDVec> idx;

    // rate at distinct time
    std::vector<rate_t> rates;
    std::vector<int> rate_index;

    inline bool keep(const IDVec& id)  const { return id[0] == 0; }
    inline bool isnew(const IDVec& id) const { return id[3] == 1; }
    inline int evid (const IDVec& id) const { return id[2] ; }

    inline bool keep(int i)  const { return keep(idx[i]); }
    inline bool isnew(int i) const { return isnew(idx[i]); }
    inline int evid (int i) const { return evid(idx[i]); }

    /*
     * for a population with data in ragged array form, we
     * form the events history using the population data and
     * the location of the individual in the ragged arrays.
     * In this constructor we assume @c p_ii.size() > 1 and
     * @c p_ss.size() > 1.
     *
     * @param ncmt nb. of compartments
     * @param ibegin beginning index of the subject in data array
     * @param isize # of indices for current subject in data array
     * @param p_time event time
     * @param p_amt dosing amount
     * @param p_rate dosing rate
     * @param p_ii dosing interval
     * @param p_evid event id
     * @param p_cmt event compartment
     * @param p_addl additional events
     * @param p_ss steady states flag
     */
    EventHistory(int ncmt, int ibegin, int isize,
                 const std::vector<T0>& p_time, const std::vector<T1>& p_amt,
                 const std::vector<T2>& p_rate, const std::vector<T3>& p_ii,
                 const std::vector<int>& p_evid, const std::vector<int>& p_cmt,
                 const std::vector<int>& p_addl, const std::vector<int>& p_ss) :
      time_(p_time),
      amt_(p_amt),
      rate_(p_rate),
      ii_(p_ii),
      evid_(p_evid),
      cmt_(p_cmt),
      addl_(p_addl),
      ss_(p_ss),
      num_event_times(isize),
      idx(isize, {0, 0, 0, 0})
    {
      const int iend = ibegin + isize;
      using stan::math::check_greater_or_equal;
      static const char* caller = "EventHistory::EventHistory";
      check_greater_or_equal(caller, "isize", isize , 1);
      check_greater_or_equal(caller, "time size", p_time.size() , size_t(iend));
      check_greater_or_equal(caller, "amt size", p_amt.size()   , size_t(iend));
      check_greater_or_equal(caller, "rate size", p_rate.size() , size_t(iend));
      check_greater_or_equal(caller, "ii size", p_ii.size()     , size_t(iend));
      check_greater_or_equal(caller, "evid size", p_evid.size() , size_t(iend));
      check_greater_or_equal(caller, "cmt size", p_cmt.size()   , size_t(iend));
      check_greater_or_equal(caller, "addl size", p_addl.size() , size_t(iend));
      check_greater_or_equal(caller, "ss size", p_ss.size()     , size_t(iend));
      for (size_t i = 0; i < isize; ++i) {
        idx[i][1] = ibegin + i;
        idx[i][2] = evid_[ibegin + i];
      }
      insert_addl_dose();
      sort_state_time();
    }

    /**
     * EventHistory constructor for an invidual, based on population
     * version.
     *
     * @param ncmt nb. of compartments
     * @param p_time event time
     * @param p_amt dosing amount
     * @param p_rate dosing rate
     * @param p_ii dosing interval
     * @param p_evid event id
     * @param p_cmt event compartment
     * @param p_addl additional events
     * @param p_ss steady states flag
     * @param theta parameter vector
     * @param biovar bioavailability
     * @param tlag lag time
     */
    EventHistory(int ncmt,
                 const std::vector<T0>& p_time, const std::vector<T1>& p_amt,
                 const std::vector<T2>& p_rate, const std::vector<T3>& p_ii,
                 const std::vector<int>& p_evid, const std::vector<int>& p_cmt,
                 const std::vector<int>& p_addl, const std::vector<int>& p_ss)
    : EventHistory(ncmt,
                   0, p_time.size(), p_time, p_amt, p_rate, p_ii, p_evid, p_cmt, p_addl, p_ss)
    {}

    bool is_reset(int i) const {
      return evid(i) == 3 || evid(i) == 4;
    }

    bool is_reset_event(const IDVec& id) {
      return evid(id) == 3 || evid(id) == 4;
    }

    bool is_dosing(int i) const {
      return evid(i) == 1 || evid(i) == 4;
    }

    /*
     * if an event is steady-state dosing event.
     */
    bool is_ss_dosing(int i) const {
      return (is_dosing(i) && (ss(i) == 1 || ss(i) == 2)) || ss(i) == 3;
    }

    static bool is_dosing(const std::vector<int>& event_id, int i) {
      return event_id[i] == 1 || event_id[i] == 4;
    }

    bool is_bolus_dosing(int i) const {
      const double eps = 1.0E-12;
      return is_dosing(i) && rate(i) < eps;
    }

    /*
     * use current event #i as template to @c push_back
     * another event.
     */
    void insert_event(int i) {
      idx.push_back({1, int(gen_time.size()), idx[i][2], 1});
      gen_time. push_back(time (i));
      gen_amt.  push_back(amt  (i));
      gen_rate. push_back(rate (i));
      gen_ii.   push_back(ii   (i));
      gen_cmt.  push_back(cmt  (i));
      gen_addl. push_back(addl (i));
      gen_ss.   push_back(ss   (i));
    }

    /**
     * Add events to EventHistory object, corresponding to additional dosing,
     * administered at specified inter-dose interval. This information is stored
     * in the addl and ii members of the EventHistory object.
     *
     * Events is sorted at the end of the procedure.
     */
    void insert_addl_dose() {
      for (int i = 0; i < size(); i++) {
        if (is_dosing(i) && ((addl(i) > 0) && (ii(i) > 0))) {
          for (int j = 1; j <= addl(i); j++) {
            insert_event(i);
            gen_time.back() += j * ii(i);
            gen_ii.back() = 0;
            gen_addl.back() = 0;
            gen_ss.back() = 0;
          }
        }
      }
    }

    /*
     * sort PMX events and nonevents times
     */
    void sort_state_time() { std::stable_sort(idx.begin(), idx.end(),
                                              [this](const IDVec &a, const IDVec &b) {
                                                using stan::math::value_of;
                                                double ta = keep(a) ? value_of(time_[a[1]]) : value_of(gen_time[a[1]]);
                                                double tb = keep(b) ? value_of(time_[b[1]]) : value_of(gen_time[b[1]]);
                                                return ta < tb;
                                              });
    }

    /**
     * Generate end-of-infusion event to indicate the range of the infusion.
     * The new event has a special EVID = 9 and it introduces no action.
     *
     * @param nCmt nb of compartments for the model
     */
    void generate_rates(int nCmt) {
      using std::vector;
      using stan::math::value_of;

      const int n = size();
      for (size_t i = 0; i < n; ++i) {
        if ((is_dosing(i)) && (rate(i) > 0 && amt(i) > 0)) {
          insert_event(i);
          idx.back()[2] = 9;    // set evid to a special type "9"
          gen_time. back() += amt(i)/rate(i);
          gen_amt.  back() = 0;
          gen_rate. back() = 0;
          gen_ii.   back() = 0;
          gen_addl. back() = 0;
          gen_ss.   back() = 0;
        }
      }
      sort_state_time();

      rate_t newRate{0.0, std::vector<T2>(nCmt, 0.0)};
      // unique_times is sorted
      std::vector<int> ut(unique_times());
      for (auto i : ut) {
        newRate.first = value_of(time(i));
        rates.push_back(newRate);
      }
      // check sorted?
      std::sort(rates.begin(), rates.end(), [](rate_t const &a, rate_t const &b) {
          return a.first < b.first;
        });

      for (size_t i = 0; i < size(); ++i) {
        if ((is_dosing(i)) && (rate(i) > 0 && amt(i) > 0)) {
          double t0 = value_of(time(i));
          double t1 = t0 + value_of(amt(i)/rate(i));
          for (auto&& r : rates) {
            if (r.first > t0 && r.first <= t1) {
              r.second[cmt(i) - 1] += rate(i);
            }
          }
        }
      }

      /*
       * rate index points to the rates for each state time,
       * since there is one rates vector per time, not per event.
       */
      rate_index.resize(idx.size());
      int iRate = 0;
      for (size_t i = 0; i < idx.size(); ++i) {
        if (rates[iRate].first != time(i)) iRate++;
        rate_index[i] = iRate;
      }
    }

    // Access functions
    inline T_time time (const IDVec& id) const { return keep(id) ? time_[id[1]] : gen_time[id[1]] ; }
    inline const T1& amt (const IDVec& id) const { return keep(id) ? amt_[id[1]] : gen_amt[id[1]] ; }
    inline const T2& rate (const IDVec& id) const { return keep(id) ? rate_[id[1]] : gen_rate[id[1]] ; }
    inline const T3& ii (const IDVec& id) const { return keep(id) ? ii_[id[1]] : gen_ii[id[1]] ; }
    inline int cmt (const IDVec& id) const { return keep(id) ? cmt_[id[1]] : gen_cmt[id[1]] ; }
    inline int addl (const IDVec& id) const { return keep(id) ? addl_[id[1]] : gen_addl[id[1]] ; }
    inline int ss (const IDVec& id) const { return keep(id) ? ss_[id[1]] : gen_ss[id[1]] ; }

    inline T_time time (int i) const { return time(idx[i]); }
    inline const T1& amt (int i) const { return amt (idx[i]); }
    inline const T2& rate (int i) const { return rate(idx[i]); }
    inline const T3& ii (int i) const { return ii  (idx[i]); }
    inline int cmt (int i) const { return cmt (idx[i]); }
    inline int addl (int i) const { return addl(idx[i]); }
    inline int ss (int i) const { return ss  (idx[i]); }

    inline size_t size() const { return idx.size(); }

    std::vector<int> unique_times() {
      std::vector<int> t(idx.size());
      std::iota(t.begin(), t.end(), 0);
      auto last = std::unique(t.begin(), t.end(),
                              [this](const int& i, const int& j) {return time(i) == time(j);});
      t.erase(last, t.end());
      return t;
    }

    /*
     * Overloading the << Operator
     */
    friend std::ostream& operator<<(std::ostream& os, const EventHistory& ev) {
      const int w = 6;
      os << "\n";
      os << std::setw(w) << "time" <<
        std::setw(w) << "amt" <<
        std::setw(w) << "rate" <<
        std::setw(w) << "ii" <<
        std::setw(w) << "evid" <<
        std::setw(w) << "cmt" <<
        std::setw(w) << "addl" <<
        std::setw(w) << "ss" <<
        std::setw(w) << "keep" <<
        std::setw(w) << "isnew" << "\n";
      for (size_t i = 0; i < ev.size(); ++i) {
        os <<
          std::setw(w)   << ev.time(i) << " " <<
          std::setw(w-1) << ev.amt(i) << " " <<
          std::setw(w-1) << ev.rate(i) << " " <<
          std::setw(w-1) << ev.ii(i) << " " <<
          std::setw(w-1) << ev.evid(i) << " " <<
          std::setw(w-1) << ev.cmt(i) << " " <<
          std::setw(w-1) << ev.addl(i) << " " <<
          std::setw(w-1) << ev.ss(i) << " " <<
          std::setw(w-1) << ev.keep(i) << " " <<
          std::setw(w-1) << ev.isnew(i) << "\n";
      }
      return os;
    }
  };

}    // torsten namespace
#endif


/* fflame_generator.hpp -- part of the fractal flames implementation
 *
 * Copyright (C) 2015 Alrik Firl
 *
 * This software may be modified and distributed under the terms
 * of the MIT license.  See the LICENSE file for details.
 */

#ifndef FF_FFLAME_GENERATOR_HPP
#define FF_FFLAME_GENERATOR_HPP

#include "fractal_flame.hpp"
#include "ifs.hpp"
#include "ifs_types.hpp"
#include "render_flame.hpp"
#include "util/EventQueue.hpp"
#include "util/ff_utils.hpp"

#include <opencv2/opencv.hpp>
//#include <boost/lockfree/queue.hpp>
#include <boost/lockfree/spsc_queue.hpp>

#include <atomic>
#include <condition_variable>
#include <limits>
#include <memory>
#include <mutex>
#include <random>
#include <string>
#include <thread>

namespace FFlames {

template <template <class> class frame_t, typename data_t, typename pixel_t>
class fflame_generator {
public:
  // NOTE: can (try to) get a reproducable sequence by not initializing
  // flame_gen(flame_rd())?
  fflame_generator(const int imheight, const int imwidth,
                   const int num_variants,
                   const int num_workers = std::thread::hardware_concurrency())
      : num_workers(num_workers), render_imheight(imheight),
        render_imwidth(imwidth), num_working_variants(num_variants),
        fflame_histoqueue(100, 30), flame_prebarrier(num_workers),
        flame_postbarrier(num_workers), fflame_th(nullptr),
        fflame_histdata(nullptr), flame_gen(flame_rd()),
        total_variant_rng(
            0, affine_fcns::variant_list<data_t>::variant_names.size() - 1) {
    // NOTE: this has to be shared between all the worker threads
    fflame_histdata = std::unique_ptr<fflame_data<data_t, pixel_t>>(
        new fflame_data<data_t, pixel_t>());

    initialize_variants();
    fflame_state.store(false);
  }

  fflame_generator(const int imheight, const int imwidth,
                   std::vector<std::string> &&manual_variants,
                   const int num_workers = std::thread::hardware_concurrency())
      : num_workers(num_workers), render_imheight(imheight),
        render_imwidth(imwidth), num_working_variants(manual_variants.size()),
        fflame_histoqueue(100, 30), flame_prebarrier(num_workers),
        flame_postbarrier(num_workers), fflame_th(nullptr),
        fflame_histdata(nullptr), flame_gen(flame_rd()),
        total_variant_rng(
            0, affine_fcns::variant_list<data_t>::variant_names.size() - 1) {
    // NOTE: this has to be shared between all the worker threads
    fflame_histdata = std::unique_ptr<fflame_data<data_t, pixel_t>>(
        new fflame_data<data_t, pixel_t>());

    initialize_variants(std::move(manual_variants));
    fflame_state.store(false);
  }

  ~fflame_generator() {
    if (fflame_state.load()) {
      stop_generation();
    }
  }

#if 1
  template <typename T> using ts_queue_t = EventQueue<T>;
#else
  template <typename T>
  using ts_queue_t =
      boost::lockfree::spsc_queue<T, boost::lockfree::capacity<1024>>;
#endif

  void start_generation() {
    // if already started, don't try to start again
    if (fflame_state.load()) {
      return;
    }

    fflame_state.store(true);
    fflame_th = std::unique_ptr<std::thread>(
        new std::thread(&fflame_generator::start_fflame_generation, this));
  }

  void stop_generation() {
    if (fflame_state.load()) {
      fflame_state.store(false);
      for (size_t th_idx = 0; th_idx < fflame_workers.size(); ++th_idx) {
        fflame_workers.at(th_idx)->finish_flame();
      }
      fflame_th->join();
    }
  }

  void register_framequeue(ts_queue_t<frame_t<pixel_t>> *queue) {
    fflame_imagequeue = queue;
  }

private:
  void initialize_variants(std::vector<std::string> &&manual_variants);
  void initialize_variants();
  void start_fflame_generation();
  void generate_fflame(fflame_util::fast_rand &rand_gen);
  void render_fflame();

  // number of threads used for the generation (not counting rendering)
  int num_workers;
  int render_imheight;
  int render_imwidth;

  // the number of variants to have active
  uint8_t num_working_variants;
  std::thread::id worker_overlord_id;

  // pause generation if above the max, resume if paused and below the min
  static constexpr size_t max_image_thresh = 25;
  static constexpr size_t min_image_thresh = 5;

  // controls the starting and stopping of the worker threads
  std::mutex gen_mtx;
  std::condition_variable gen_cv;

  // for passing results asynchronously between the generate & render steps
  // EventQueue<cv::Mat_<histogram_info<pixel_t>>> fflame_histoqueue;

  ts_queue_t<std::vector<histogram_info<pixel_t>>> fflame_histoqueue;
  // for holding the final images. provided by the caller
  ts_queue_t<frame_t<pixel_t>> *fflame_imagequeue;

  fflame_util::barrier flame_prebarrier;
  fflame_util::barrier flame_postbarrier;

  // the rendering/controller flame thread
  std::unique_ptr<std::thread> fflame_th;
  // the fractal flame generation threads
  std::vector<std::unique_ptr<flame_thread>> fflame_workers;

  // flag for starting/stopping the whole fflame pipeline
  std::atomic<bool> fflame_state;

  // holds the list of the current variation functions to use
  std::unique_ptr<affine_fcns::invoker<data_t>> flamer;
  std::unique_ptr<fflame_data<data_t, pixel_t>> fflame_histdata;

  // the various RNGs needed for the generation
  std::random_device flame_rd;
  std::mt19937 flame_gen;
  // std::uniform_int_distribution<uint8_t> working_variant_rng;
  std::uniform_int_distribution<> total_variant_rng;
};

template <template <class> class frame_t, typename data_t, typename pixel_t>
void fflame_generator<frame_t, data_t, pixel_t>::initialize_variants(
    std::vector<std::string> &&manual_variants) {
  flamer = std::unique_ptr<affine_fcns::invoker<data_t>>(
      new affine_fcns::invoker<data_t>(std::move(manual_variants)));
  flamer->randomize_parameters(-2, 2);
}

template <template <class> class frame_t, typename data_t, typename pixel_t>
void fflame_generator<frame_t, data_t, pixel_t>::initialize_variants() {
  flamer = std::unique_ptr<affine_fcns::invoker<data_t>>(
      new affine_fcns::invoker<data_t>(num_working_variants));
  // NOTE: we always want to have the linear variant (variant #0)
  int variant_idx = 0;
  auto linear_variant_id =
      affine_fcns::variant_list<data_t>::variant_names[variant_idx];
  flamer->set_variant(variant_idx, linear_variant_id);
  for (int i = variant_idx + 1; i < num_working_variants; ++i) {
    std::string selected_variant;
    // keep re-rolling the variant until we get a valid one
    do {
      variant_idx = total_variant_rng(flame_gen);
      selected_variant =
          affine_fcns::variant_list<data_t>::variant_names[variant_idx];
    } while (!flamer->check_variant_valid(selected_variant));

    // can assume by here that the selected variant is a valid one
    flamer->set_variant(i, selected_variant);
  }
  flamer->randomize_parameters(-2, 2);
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

template <template <class> class frame_t, typename data_t, typename pixel_t>
void fflame_generator<frame_t, data_t, pixel_t>::start_fflame_generation() {
  fflame_state.store(true);

  std::uniform_int_distribution<uint64_t> flame_dist(
      0, std::numeric_limits<uint64_t>::max());

  // spawn the worker threads
  for (int th_idx = 0; th_idx < num_workers; ++th_idx) {
    // seeds the flame thread's fast random number generator
    fflame_workers.emplace_back(
        new flame_thread(flame_dist(flame_gen), flame_dist(flame_gen)));
    // this step is the one that actually spawns the thread
    bool flame_run;
    std::thread::id tid;
    std::tie(flame_run, tid) = fflame_workers.at(th_idx)->do_flame(
        &fflame_generator::generate_fflame, this);

    // this would be sort of a big deal. should probably throw something here
    if (!flame_run) {
      std::cout << "NOTE: flame task failed" << std::endl;
    }

    if (th_idx == 0) {
      worker_overlord_id = tid;
    }
  }

  // prepare the image and push the image to the shared queue
  render_fflame();
}

template <typename pixel_t>
histogram_info<pixel_t>
print_mat(const cv::Mat_<histogram_info<pixel_t>> &hdata, int row, int col) {
  return hdata(row, col);
}

// makes the fractal flame histogram using the chaos game. Is invoked by N
// threads
template <template <class> class frame_t, typename data_t, typename pixel_t>
void fflame_generator<frame_t, data_t, pixel_t>::generate_fflame(
    fflame_util::fast_rand &rand_gen) {
  while (fflame_state.load()) {
    auto start_iter_time = std::chrono::high_resolution_clock::now();

    // 0. check if the generation should pause (i.e. too many frames in queue)
    // --> leave this for after you get it working
    // 1. start the next flame generation
    run_fflame<data_t, pixel_t>(flamer.get(),
                                fflame_constants::num_samples / num_workers,
                                fflame_histdata.get(), rand_gen);

    auto end_iter_time = std::chrono::high_resolution_clock::now();
    double time_elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
                              end_iter_time - start_iter_time)
                              .count();
    std::cout << "thread " << std::this_thread::get_id()
              << " generation took -- " << time_elapsed << " ms" << std::endl;

    // 2. wait for all the threads to finish the previous round
    flame_prebarrier.wait();

    // merge the N thread results using the overlord thread
    if (std::this_thread::get_id() == worker_overlord_id) {
      static std::vector<std::string> mutated_variant_ids(num_working_variants,
                                                          "default");

      // somewhat unfortunate, but need to dynamically allocate to avoid scoping
      // problems
      auto hist_info = std::unique_ptr<std::vector<histogram_info<pixel_t>>>(
          new std::vector<histogram_info<pixel_t>>(fflame_constants::imheight *
                                                   fflame_constants::imwidth));
      fflame_histdata->get_and_reset(*hist_info);

      // 3.5 pass the finished histogram to the shared-buffer for rendering
      fflame_histoqueue.push(std::move(hist_info));

      // 4. mutate the variants
      std::string selected_variant;
      // keep re-rolling the variant until we get a valid one
      do {
        selected_variant =
            affine_fcns::variant_list<data_t>::variant_names[total_variant_rng(
                flame_gen)];
      } while (!flamer->check_variant_valid(selected_variant));

      // replace a random variant (that's not the linear variant)
      int mod_idx =
          (num_working_variants > 1)
              ? total_variant_rng(flame_gen) % (num_working_variants - 1) + 1
              : 0;
      // can assume by here that the selected variant is a valid one
      flamer->set_variant(mod_idx, selected_variant);
      flamer->randomize_parameters(-2, 2);

      flamer->print_variant_list();

      // there should be no threads waiting at the pre-barrier at this point;
      // the other (non-overlord) threads should be waiting at the post-barrier
      flame_prebarrier.reset(num_workers);
    }
  }
}

// generates an image based on the fflame histogram
template <template <class> class frame_t, typename data_t, typename pixel_t>
void fflame_generator<frame_t, data_t, pixel_t>::render_fflame() {
  bool got_histdata = false;
  // int raw_counter = 0;

  fflame_renderer<data_t> flame_image_renderer(render_imheight, render_imwidth);

  double total_render_time = 0;
  while (fflame_state.load()) {
    // 1. get the histogram
    auto hist_info = fflame_histoqueue.pop(got_histdata);
    if (got_histdata && hist_info) {
      auto start_render_time = std::chrono::high_resolution_clock::now();

      std::unique_ptr<frame_t<pixel_t>> image =
          std::unique_ptr<frame_t<pixel_t>>(
              new frame_t<pixel_t>(render_imheight, render_imwidth));
      std::fill(image->data, image->data + image->rows * image->cols, 0);

      // 2. call the rendering routine, get resultant image
      flame_image_renderer.template render<frame_t, pixel_t>(
          image.get(), std::move(hist_info));

      /*
                  //filter out the flames that are too sparse
                  int num_nonzero = 0;
                  const double image_threshold = 0.1 * imwidth * imheight;
                  const double px_threshold = 1.0;
                  for (int r = 0; r < imheight; ++r)
                  {
                      for (int c = 0; c < imwidth; ++c)
                      {
                          size_t px_sum = cv::sum(outfile_image.at<pixel_t>(r,
         c))[0]; if(px_sum > px_threshold) { num_nonzero++;
                          }
                      }
                  }
                  if(num_nonzero < image_threshold)
                  {
                      std::cout << "NOTE: flame too dark" << std::endl;
                      continue;
                  }
                  //mostly for debugging -- save the images to disk
                  const std::string raw_impath = "FFLAME_baseimage_" +
         std::to_string(raw_counter++) + ".png"; cv::imwrite(raw_impath, image);
      */

      fflame_imagequeue->push(std::move(image));

      auto end_render_time = std::chrono::high_resolution_clock::now();
      total_render_time +=
          std::chrono::duration_cast<std::chrono::milliseconds>(
              end_render_time - start_render_time)
              .count();
    }
  }

  std::cout << "rendering took -- " << total_render_time << " ms" << std::endl;
}

} // namespace FFlames
#endif


/*
 *  libtraj_gen.cpp
 *  Trajectory Generator Library
 *
 *  Created by Matthew O'Kelly on 7/17/15.
 *  Copyright (c) 2015 Matthew O'Kelly. All rights reserved.
 *  mokelly@seas.upenn.edu
 *
*/

/*
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions are met:
 *
 *  * Redistributions of source code must retain the above copyright notice, this
 *    list of conditions and the following disclaimer.
 *
 *  * Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 *  * Neither the name of Autoware nor the names of its
 *    contributors may be used to endorse or promote products derived from
 *    this software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 *  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 *  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 *  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 *  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 *  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 *  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 *  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 *  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 *  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
*/

/*
 * This file contains the vehicle model for use in predicting the ego vehicle's motion
 * It also contains a heuristic for initial trajetories and
 * a gradient descent method for refining them...
 *
 * First some constants are defined, these will be specific to your vehicle
 * Nine functions are included: speedControlLogic(),
 * responsetoControlInputs(), and most
 * importantly, the vehicle's motionModel() etc...
 *
 * motionModel() calls responseToControlInputs()
 * responseToControlInputs calls speedControlLogic()
 *
 * For a standalone build use the following command: g++ trajectorygenerator.cpp -o trajectorygenerator -O2 -larmadillo
 * Dependancies: Armadillo, http://arma.sourceforge.net/
 *
*/

#include <iostream>
#include <vector>
#include <armadillo>
#include <cmath>
#include <algorithm>
#include <iostream>
#include <fstream>
#include <stdlib.h>
#include <stdio.h>
#include "libtraj_gen.h"
// #include "trajectorygenerator.h"

using namespace std;

//-----------------------------FUNCTIONS------------------------------------//

// ------------INITIALIZE PARAMETERS----------//
// Heuristic to compute initial paramters given
// INPUT: Initial state, and a goal state
// OUTPUT: Estimated control parameters
// Use #define USE_HEURISTIC_INIT if using Nagy 2001 method
// This is necessary for cubic splines
// Use #define FIRST_ORDER if working with first order splines

union Spline initParams(union State veh, union State goal)
{

    // Local variables for init and goal:
    double sx_f = goal.sx;
    double sy_f = goal.sy;
    double theta_f = goal.theta;
    double kappa_0 = veh.kappa;
    double kappa_f = goal.kappa;

    // Initialize output
    union Spline curvature;

    // Convenience
    double d_theta = abs(theta_f);

    // Heuristic for initial guess from Nagy and Kelly, 2001
    // 'Trajectory Generation for Car-like Robots using
    //  Cubic Curvature Polynomials'
    #ifdef USE_HEURISTIC_INIT
        // Order of computation matters, need to compute d before s and so on...

        // Commented out, because unused for cubic splines
        //double c = 0.0;

        // Note we must cast to double to avoid ambiguity
        double d = sqrt((double)pow(sx_f,2) +(double) pow(sy_f,2));

        double s = d * ((pow(d_theta,2))/5.0 + 1.0) + (2.0/5.0)*d_theta;

        // Commented out because unused for cubic splines
        //double a = (6* (theta_f/(pow(s,2)))) - (2*kappa_0/s) + (4*kappa_f/s);

        double b = (3/(pow(s,2))) * (kappa_0 + kappa_f) + (6*theta_f/(pow(s,3)));

        double si=0.00;
        curvature.kappa_0 = veh.kappa;
        curvature.kappa_1=(1.00/49.00)*(8.00*b*si - 8.00*b*curvature.s - 26.00*curvature.kappa_0 - curvature.kappa_3);
        curvature.kappa_2=0.25*(curvature.kappa_3 -2.00*curvature.kappa_0 +5.00*curvature.kappa_1);
        curvature.kappa_3 = goal.kappa;
        curvature.s = s;

    #endif

    // Alternative heuristic initial guess for first order spline
    // From Thomas Howard's dissertation page 38
    #ifdef FIRST_ORDER
        double d = sqrt((double)pow(sx_f,2) +(double) pow(sy_f,2));
        double s = d * ((pow(d_theta,2))/5 + 1) + (2/5)  *d_theta;
        union Spline curvature;
        double si=0.00;
        curvature.kappa_0 = veh.kappa;
        curvature.kappa_1 = 0.00;
        curvature.kappa_2 = 0.00;
        curvature.kappa_3 = goal.kappa;
        curvature.s = s;
    #endif

    curvature.success=TRUE;

    // Return the initialized curvature union
    return curvature;

}

// ------------SPEED CONTROL----------//
// Function for computing safe (feasible) speed and curvature
// INPUT: Next vehicle state
// OUTPUT: Updated next vehicle state

union State speedControlLogic(union State veh_next)
{

    // Calculate speed (look up from next state vector)
    double vcmd = abs(veh_next.v);
    double kappa_next = veh_next.kappa;

    // Compute safe speed
    double compare_v=(kappa_next-ascl)/bscl;
    double vcmd_max = max((double)vscl, compare_v);

    // Compute safe curvature
    double compare_kappa = ascl + (bscl * vcmd);
    double kmax_scl = min((double) kmax, compare_kappa);

    // Check if max curvature for speed is exceeded
    if(kappa_next >= kmax_scl)
    {
        // Check for safe speed
        vcmd = sf * vcmd_max;
    }

    // Update velocity command, this is not quite equivalent to Ferguson, Howard, & Liukhachev
    veh_next.v = vcmd;

    // Return vehicle state
    return veh_next;
}

// ------------RESPONSE TO CONTROL----------//
// Function for computing the vehicles response to control inputs
// Mainly deals with delays associated with controls
// Also considers upper and lower safety bounds
// INPUT: Current vehicle state, next vehicle state, sampling time
// OUTPUT: Next vehicle state

union State responseToControlInputs(union State veh, union State veh_next, double dt)
{

    // Local variables:
    double kappa = veh.kappa;
    double kappa_next = veh_next.kappa;
    double v = veh.v;
    double v_next = veh_next.v;

    // Compute curvature rate command
    double kdot = (kappa_next - kappa)/dt;

    // Check against upper bound on curvature rate
    kdot = min(kdot, (double) dkmax);

    // Check against lower bound on curvature rate
    kdot = max(kdot, (double) dkmin);

    // Call the speedControlLogic function
    veh_next = speedControlLogic(veh_next);

    // Compute curvature at the next vehicle state
    kappa_next = kappa + kdot*dt;

    // Check upper bound on curvature
    kappa_next = min(kappa_next, (double) kmax);

    // Check lower bound on curvature
    kappa_next = max(kappa_next, (double) kmin);

    // Compute acceleration command
    double vdot = (v_next -v)/dt;

    // Check for upper bound on acceleration
    vdot = min(vdot, (double) dvmax);

    // Check for lower bound on acceleration
    vdot = max(vdot, (double) dvmin);

    // Compute velocity at next state
    veh_next.v = v + vdot*dt;

    // Return next vehicle state
    return veh_next;
}

// ------------COMPUTE CURVATURE----------//
// Computes the curvature at the next state
// Still messy but will clean shortly...

double getCurvatureCommand(union Spline curvature, double dt, double v, double t)
{
    // Local variables for curvature constants
    double kappa_0 = curvature.kappa_0;
    double kappa_1 = curvature.kappa_1;
    double kappa_2 = curvature.kappa_2;
    double kappa_3 = curvature.kappa_3;
    double s = curvature.s;

    // Init next command
    double k_next_cmd =0.0;

    // Assume position at t=0, si=0
    double si = 0.0;

    // Estimate distance traveled on spiral
    double st = v*t;

    // These equations can be found in McNaughton, page 76
    // They make the knot points equally spaced.
    // They apply to stable cubic paths
    #ifdef CUBIC_STABLE
        double a =0;
        double b =0;
        double c= 0;
        double d = 0;
        a = kappa_0;
        b = (-0.50)*(-2*kappa_3 + 11*kappa_0 - 18*kappa_1 + 9*kappa_2)/(s-si);
        c = (4.50)*(-kappa_3 + 2*kappa_0 - 5*kappa_1 +4*kappa_2)/(pow((s-si), 2));
        d = (-4.50)*(-kappa_3 + kappa_0 - 3*kappa_1 + 3*kappa_2)/(pow((s-si), 3));
        k_next_cmd = a + b*st + c*pow(st,2) + d*pow(st,3);
    #endif

    // These equations can be found in McNaughton, page 77
    // They make the knot points equally spaced.
    // They apply to stable quintic paths
    #ifdef QUINTIC_STABLE
        double a = kappa_0;
        double b = kappa_1;
        double c = kappa_2/2.0;
        double d = -(575*kappa_0 -648*kappa_3 +81*kappa_4 -8*kappa_5 +170*kappa_1*st +22*kappa_2*pow(st,2)/ (8*pow(st,3));
        double e = 9*(37*kappa_0 -45*kappa_3 +9*kappa_4 -kappa_5 +10*kappa_1*st +kappa_2*pow(st,2))/(2*pow(st,4));
        double f = -9*(85*kappa_0 -108*kappa_3 +27*kappa_4 -4*kappa_5 +22*kappa_1*st +2*kappa_2*pow(st,2))/(8*pow(st,5));
        k_next_cmd = a + b*st + c*pow(st,2) + d*pow(st,3) + e*pow(st,4) +f*pow(st,5);
    #endif

    // Compute next command from Howard page 39
    #ifdef FIRST_ORDER
        k_next_cmd = kappa_0 + ((kappa_1 - kappa_0)/s) * (st);
    #endif

    // Return value
    return k_next_cmd;

}

// ------------COMPUTE VELOCITY----------//
// Updates the vehicles velocity based on velocity profile
// Currently assumes constant
/// Will update to include ramp...

double getVelocityCommand(double v_goal, double v)
{
    #ifdef DEBUG
    cout << "Function: getVelocityCommand()"<< endl;
    #endif
    // Some function based on parameterization of velocity profile
    // Set to constant for now
    double v_next_cmd=v_goal;
    return v_next_cmd;
}

// ------------MOTION MODEL----------//
// Computes update to vehicle state
// INPUT: Current vehicle state, parameterized control inputs, sampling time
// OUTPUT: Next vehicle state

union State motionModel(union State veh, union State goal, union Spline curvature, double dt, double horizon, int flag)
{
    // Initialized the elapsed time to 0.0 s
    double t =0.0;
    // Setup local data structures for holding
    // the previous and next vehicle states
    union State veh_next;
    union State veh_temp;
    // Set veh_temp to the current state
    veh_temp=veh;
    // Compute the stop time for the simulation
    horizon = curvature.s/goal.v;

    while(t < horizon)
    {
        // Read state vector
        double sx = veh_temp.sx;
        double sy = veh_temp.sy;
        double v = veh_temp.v;
        double theta = veh_temp.theta;
        double kappa = veh_temp.kappa;
        double v_goal= goal.v;

        // Compute change in 2D x-position (this is x_dot)
        double sx_next = sx + (v * cos(theta) * dt);

        veh_next.sx = sx_next;

        // Compute change in 2D y-position (this is y_dot)
        double sy_next = sy + (v * sin(theta) * dt);

        veh_next.sy = sy_next;

        // Compute change in 2D orientation (this is theta_dot)
        double theta_next = theta + (v * kappa * dt);

        veh_next.theta = theta_next;

        // Get curvature command
        double kappa_next = getCurvatureCommand(curvature, dt , veh.v,t);

        veh_next.kappa = kappa_next;

        // Get velocity command
        double v_next = getVelocityCommand(v_goal, v);

        veh_next.v = v_next;

        // Get acceleration command, not used yet...
        //double a_next_cmd = 0;

        // Update the next vehicle state
        veh_next = responseToControlInputs(veh, veh_next, dt);

        // Increment the timestep
        t=t+dt;

        // Update veh_temp
        veh_temp=veh_next;

        // Write to log file if we are testing
        if(flag==1)
        {
            fmm_sx <<veh_temp.sx<<", ";
            fmm_sy <<veh_temp.sy<<", ";
            fmm_v <<veh_temp.v<<", ";
            fmm_theta << veh_temp.theta<<", ";
            fmm_kappa<<veh_temp.kappa<<", ";
        }

    }
        // Print information to console
        #ifdef DEBUG
        cout << "Next state sx: " << veh_next.sx << endl;
        cout << "Next state sy: " << veh_next.sy << endl;
        cout << "Next state theta: " << veh_next.theta << endl;
        cout << "Next state v: " << veh_next.v << endl;
        cout << "Next state kappa: " << veh_next.kappa<< endl;
        #endif

    // Return the state at the end of the trajectory
    return veh_next;
}

// ------------CHECK CONVERGENCE----------//
// Checks to see if we have reached target solution
// INPUT: Next vehicle state predicted by the forward motion model and goal state
// OUTPUT: Boolean
// TO DO: if theta is constrained, please check that this is an appropriate metric.
// IE consider 359 degrees vs. 1 degree
// Can use fmod (modulus) over pi or 2*pi

bool checkConvergence(union State veh_next, union State goal)
{
    #ifdef DEBUG
    cout << "Function: checkConvergence()"<< endl;
    #endif

    double sx_error = abs(veh_next.sx - goal.sx);
    double sy_error = abs(veh_next.sy - goal.sy);
    double theta_error = abs(veh_next.theta - goal.theta);

    // Commented out to suppress compiler warning about unused var
    // double v_error = abs(veh_next.v - goal.v);
    // double kappa_error = abs(veh_next.kappa - goal.kappa);

    #ifdef DEBUG_OUTPUT
    cout << "Error sx: " << sx_error << endl;
    cout << "Error sy: " << sy_error << endl;
    cout << "Error theta: " << theta_error << endl;
    // Commented out to suppress compiler warning about unused var
    // cout << "Error v: " << v_error << endl;
    // cout << "Error kappa: " << kappa_error << endl;
    #endif

    // Depending on the order of the spline the full range of checks is:
    // if(sx_error<general_e && sy_error<general_e && theta_error<general_e
    // && theta_error<general_e && v_error<general_e && kappa_error<general_e)
    if(sx_error<general_e && sy_error<general_e && theta_error<general_e)
    {
        #ifdef DEBUG
        cout << "Converged"<< endl;
        #endif
        return TRUE;
    }

    else
    {
        #ifdef DEBUG
        cout << "Not Converged"<< endl;
        #endif
        return FALSE;
    }

}

// ------------ESTIMATE JACOBIAN----------//
// Forward difference method via predictive motion model
// INPUT: Initial State, Target State, Parameterized Action,
//        Predictive Motion Model, Action Parameters
// OUTPUT: Jacobian Estimate

union State pDerivEstimate(union State veh, union State veh_next, union State goal, union Spline curvature, int p_id, double h, double dt, double horizon, int stateIndex)
{
    #ifdef DEBUG
    cout << "Function: pDerivEstimate()"<< endl;
    #endif

    // Index for iterating through parameters
    int i = 0;

    // Compute difference between desired sx and actual
    // May need to update with intial state if not 0,0,0,0,0
    union State delta_state;

    // Compute the difference between the end state of the current spline and the goal
    for (i=0; i<stateIndex; i++)
    {
        delta_state.state_value[i] = (goal.state_value[i] - veh_next.state_value[i]);
    }

    // Create parameter perturbation vector
    union Spline perturb_curve = curvature;

    // Perturb parameter
    perturb_curve.spline_value[p_id] = curvature.spline_value[p_id] +h;

    // Check forward motion model with perturbed parameter
    union State perturb_next = motionModel(veh, goal, perturb_curve, dt, horizon, 0);

    // Now create a new structure to hold the difference between goal and perturbed forward motion model
    union State delta_perturb_state;

    // Compute difference between goal sx and actual
    for (i=0; i<stateIndex; i++)
    {
         delta_perturb_state.state_value[i] = goal.state_value[i] - perturb_next.state_value[i];
    }

    // Structure to hold the difference between the perturbed error and unperturbed erro
    union State partial_state_vector;

    // Estimate partial derivative
    for (i=0; i<stateIndex; i++)
    {
        double num = (delta_perturb_state.state_value[i] - delta_state.state_value[i]);
        partial_state_vector.state_value[i] = num/h;
    }

    return partial_state_vector;

}

// ------------UPDATE PARAMETERS----------//
// Inverts the Jacobian and computes the parameter update

union Spline generateCorrection(union State veh, union State veh_next, union State goal, union Spline curvature, double dt, double horizon)
{
    // Compute jacobian with forward gradient
    int i;
    int stateIndex=3;
    int j;

    // Matrix J will contain the Jacobian
    arma::mat J(stateIndex,stateIndex);

    // Initialize it to zero
    J.fill(0.0);

    // Parameter perturbation vector
    arma::vec h(stateIndex);

    // How much to perturb each parameter
    // Note h(0) is larger than others because it perturbs length ~20m
    // Rather than kappa < ~.19 rad/m
    h(0)= (double) h_global*10;
    h(1)= (double) h_global;
    h(2)= (double) h_global;

    // For each parameter compute the vector associated with a small perturbation of its value
    for (i=0; i<stateIndex; i++)
    {
        union State temp;
        temp= pDerivEstimate(veh, veh_next, goal, curvature, i, h(i), dt, horizon, stateIndex);

        // For each element of the state vector place the value in the Jacobian matrix (column-wise)
        for (j=0; j<stateIndex; j++)
        {
            J(j,i) = temp.state_value[j];
        }
    }

    // Put parameters in a vector
    arma::vec dX(stateIndex);
    dX.fill(0.0);

    for(i=0; i<stateIndex; i++)
    {
        dX(i) = goal.state_value[i]-veh_next.state_value[i];
    }

    // Solve for delta P
    // Note that we use a try catch routine as it is possible for J to be singular
    arma::vec dP(stateIndex);
    try
    {
        J=J.i();
    }
    catch(const std::exception& e)
    {
        curvature.success=FALSE;
        return curvature;
        cout<< "goal.sx: "<< goal.sx<<endl;
        cout<< "goal.sy: "<< goal.sx<<endl;
        cout<< "goal.theta: "<< goal.sx<<endl;
    }

    dP=J*dX;

    // Update curvature parameters
    for(i=0; i<stateIndex; i++)
    {
        curvature.spline_value[i]=curvature.spline_value[i]-dP(i);
    }

    // Return new spline
    return curvature;
}

// ------------NEXT STATE----------//
// Computes update to vehicle state
// for the purpose of control
// Not called in local planner
// INPUT: Current vehicle state, parameterized control inputs, sampling time
// OUTPUT: Next vehicle state

union State nextState(union State veh, union Spline curvature, double vdes, double dt, double elapsedTime)
{
    union State veh_next;
    union State veh_temp;
    double total_time = elapsedTime + dt;
    //cout<<"total_time: "<< total_time<<endl;
    double t = 0.00;

    veh_temp=veh;

    // Read state vector
    double sx = veh_temp.sx;
    double sy = veh_temp.sy;
    double v = veh_temp.v;
    double theta = veh_temp.theta;
    double kappa = veh_temp.kappa;
    double v_goal= vdes;

    //cout<<"Total sim time: " << total_time << endl;

    while(t <= total_time)
    {

        // Compute change in 2D x-position (this is x_dot)
        double sx_next = sx + (v * cos(theta) * dt);

        veh_next.sx = sx_next;

        // Compute change in 2D y-position (this is y_dot)
        double sy_next = sy + (v * sin(theta) * dt);

        veh_next.sy = sy_next;

        // Compute change in 2D orientation (this is theta_dot)
        double theta_next = theta + (v * kappa * dt);

        veh_next.theta = theta_next;

        // Get curvature command
        double kappa_next = getCurvatureCommand(curvature, dt , veh.v,t);

        veh_next.kappa = kappa_next;

        // Get velocity command
        double v_next = getVelocityCommand(v_goal, v);

        veh_next.v = v_next;

        // Get acceleration command
        //double a_next_cmd = 0;

        // Update the next vehicle state
        veh_next = responseToControlInputs(veh, veh_next, dt);

        // Increment the timestep
        t=t+step_size;

        // Update veh_temp
        veh_temp=veh_next;
    }
    cout<<"Exited the while loop t = "<<t<<endl;

    return veh_next;
}

// plotTraj is used by rViz to compute points for line strip,
// it is a lighter weight version of nextState
union State genLineStrip(union State veh, union Spline curvature, double vdes, double t)
{
    union State veh_next;
    union State veh_temp;
    double dt = plot_step_size;

    veh_temp=veh;

    // Read state vector
    double sx = veh_temp.sx;
    double sy = veh_temp.sy;
    double v = veh_temp.v;
    double theta = veh_temp.theta;
    double kappa = veh_temp.kappa;
    double v_goal= vdes;

    // Compute change in 2D x-position (this is x_dot)
    double sx_next = sx + (v * cos(theta) * dt);

    veh_next.sx = sx_next;

    // Compute change in 2D y-position (this is y_dot)
    double sy_next = sy + (v * sin(theta) * dt);

    veh_next.sy = sy_next;

    // Compute change in 2D orientation (this is theta_dot)
    double theta_next = theta + (v * kappa * dt);

    veh_next.theta = theta_next;

    // Get curvature command
    double kappa_next = getCurvatureCommand(curvature, dt , veh.v, t);

    veh_next.kappa = kappa_next;

    // Get velocity command
    double v_next = getVelocityCommand(v_goal, v);

    veh_next.v = v_next;

    // Get acceleration command
    // double a_next_cmd = 0;

    // Update the next vehicle state
    veh_next = responseToControlInputs(veh, veh_next, dt);

    // Update veh_temp
    veh_temp=veh_next;

    return veh_next;
}

//------------------MAIN FUNCTION AND HELPER FOR STANDALONE OPERATION------------------------//

#ifdef STANDALONE

union Spline trajectoryGenerator(double sx, double sy, double theta, double v, double kappa)
{

#ifdef DEBUG
    cout << "Function: main()"<< endl;
#endif

// Set velocity
    double v_0=0.10;

// Initialize convergence criteria to false
    bool convergence = FALSE;

// Set goal vector
    union State goal;
    goal.sx = sx;
    goal.sy = sy;
    goal.theta = theta;
    goal.v = v;
    goal.kappa = kappa;

// Initialize and set current state vector
    union State veh;
    veh.sx = 0.0;
    veh.sy = 0.0;
    veh.theta = 0.0;
    veh.v = v;
    veh.kappa = 0.0;

// Initialize next state
    union State veh_next;
    veh_next.sx = 0.0;
    veh_next.sy = 0.0;
    veh_next.theta = 0.0;
    veh_next.v = 0.00;
    veh_next.kappa = 0.0;

// Initialize parameters as a function of init and goal vectors
    union Spline curvature;
    curvature = initParams(veh, goal);

 // Initialize iteration counter
    int iteration = 0;

// Set timestep
    double dt = step_size;

// While loop for computing trajectory parameters
    while(convergence == FALSE && iteration<10)
    {

    // Set time horizon
        double horizon = curvature.s/v_0;

    // Run motion model
        veh_next = motionModel(veh, goal, curvature, dt, horizon, 0);

    // Determine convergence criteria
        convergence = checkConvergence(veh_next, goal);

    // If the motion model doesn't get us to the goal compute new parameters
        if(convergence==FALSE)
        {
        // Update parameters
            curvature = generateCorrection(veh, veh_next, goal, curvature, dt, horizon);
            iteration++;
            if(curvature.success==FALSE)
            {
                break;
            }

        #ifdef ONE_ITER
            convergence = TRUE;
        #endif
        }
    }

    if(convergence==FALSE)
    {
        cout << "Not Converged"<< endl;
    }

    else
    {
        cout << "Converged in "<<iteration<<" iterations"<< endl;

    #ifdef LOG_OUTPUT
    // Set time horizon
        double horizon = curvature.s/v_0;
    // Run motion model and log data for plotting
        veh_next = motionModel(veh, goal, curvature, 0.1, horizon, 1);
        fmm_sx<<"0.0 \n";
        fmm_sy<<"0.0 \n";
    #endif
    }

    return curvature;
}

int main(void)
{
// Open files for logging
    fmm_sx.open("fmm_sx.dat");
    fmm_sy.open("fmm_sy.dat");
    fmm_v.open("fmm_v.dat");
    fmm_theta.open("fmm_theta.dat");
    fmm_kappa.open("fmm_kappa.dat");

    int i;
    int j;
    int k=0;
    double sx = 1;
    double sy = 1;
    union Spline curvature;

    #ifdef GEN_PLOT_FILES
    for(i=5; i<20; i++)
    {
        sx=i;
        for(j=-50; j<50; j++)
        {
            if(abs(j)<i+10)
            {
                cout<<"Trajectory Number: "<<k<<endl;
                sy=j;
                double k_goal;
                if(j<0)
                {
                    k_goal=0.15;
                }
                else
                {
                    k_goal=-0.15;
                }
                curvature = trajectoryGenerator(sx, sy, 0.0, 0.1, 0.0);
                k++;
            }
        }
    }
    #endif

trajectoryGenerator(33.7553, 2.04623, 0.0726169, 11.1111, 0.0);

//Close files
    fmm_sx.close();
    fmm_sy.close();
    fmm_v.close();
    fmm_theta.close();
    fmm_kappa.close();

    return 1;
}

#endif