//
// Created by christoph on 02.10.18.
//

#include <fstream>
#include <iostream>
#include <chrono>

#include <boost/algorithm/string.hpp>
#include <boost/algorithm/string/split.hpp>

#include <GL/glew.h>

#include <Utils/Convert.hpp>
#include <Utils/File/Logfile.hpp>
#include <Math/Math.hpp>
#include <Graphics/Renderer.hpp>
#include <Graphics/Shader/ShaderManager.hpp>
#include <Graphics/Texture/TextureManager.hpp>
#include <Graphics/OpenGL/GeometryBuffer.hpp>
#include <Graphics/OpenGL/Texture.hpp>

#include "Utils/HairLoader.hpp"
#include "Utils/TrajectoryFile.hpp"
#include "VoxelCurveDiscretizer.hpp"

#define BIAS 0.001

/**
 * Helper function for rayBoxIntersection (see below).
 */
bool rayBoxPlaneIntersection(float rayOriginX, float rayDirectionX, float lowerX, float upperX,
                             float &tNear, float &tFar)
{
    if (std::abs(rayDirectionX) < BIAS) {
        // Ray is parallel to the x planes
        if (rayOriginX < lowerX || rayOriginX > upperX) {
            return false;
        }
    } else {
        // Not parallel to the x planes. Compute the intersection distance to the planes.
        float t0 = (lowerX - rayOriginX) / rayDirectionX;
        float t1 = (upperX - rayOriginX) / rayDirectionX;
        if (t0 > t1) {
            // Since t0 intersection with near plane
            float tmp = t0;
            t0 = t1;
            t1 = tmp;
        }

        if (t0 > tNear) {
            // We want the largest tNear
            tNear = t0;
        }
        if (t1 < tFar) {
            // We want the smallest tFar
            tFar = t1;
        }
        if (tNear > tFar) {
            // Box is missed
            return false;
        }
        if (tFar < 0) {
            // Box is behind ray
            return false;
        }
    }
    return true;
}

/**
 * Implementation of ray-box intersection (idea from A. Glassner et al., "An Introduction to Ray Tracing").
 * For more details see: https://www.siggraph.org//education/materials/HyperGraph/raytrace/rtinter3.htm
 */
bool rayBoxIntersection(glm::vec3 rayOrigin, glm::vec3 rayDirection, glm::vec3 lower, glm::vec3 upper,
                        float &tNear, float &tFar)
{
    tNear = -1e7;
    tFar = 1e7;
    for (int i = 0; i < 3; i++) {
        if (!rayBoxPlaneIntersection(rayOrigin[i], rayDirection[i], lower[i], upper[i], tNear, tFar)) {
            return false;
        }
    }

    //entrancePoint = rayOrigin + tNear * rayDirection;
    //exitPoint = rayOrigin + tFar * rayDirection;
    return true;
}

bool VoxelDiscretizer::addPossibleIntersections(const glm::vec3 &v1, const glm::vec3 &v2, float a1, float a2)
{
    float tNear, tFar;
    glm::vec3 voxelLower = glm::vec3(index);
    glm::vec3 voxelUpper = glm::vec3(index + glm::ivec3(1,1,1));
    if (rayBoxIntersection(v1, (v2 - v1), voxelLower, voxelUpper, tNear, tFar)) {
        bool intersectionNear = 0.0f <= tNear && tNear <= 1.0f;
        bool intersectionFar = 0.0f <= tFar && tFar <= 1.0f;
        if (intersectionNear) {
            glm::vec3 entrancePoint = v1 + tNear * (v2 - v1);
            float interpolatedAttribute = a1 + tNear * (a2 - a1);
            currentCurveIntersections.emplace_back(entrancePoint, interpolatedAttribute);
        }
        if (intersectionFar) {
            glm::vec3 exitPoint = v1 + tFar * (v2 - v1);
            float interpolatedAttribute = a1 + tFar * (a2 - a1);
            currentCurveIntersections.emplace_back(exitPoint, interpolatedAttribute);
        }
        if (intersectionNear || intersectionFar) {
            return true; // Intersection found
        }
    }

    return false;
}

void VoxelDiscretizer::setIndex(glm::ivec3 index)
{
    this->index = index;
}

float VoxelDiscretizer::computeDensity(float maxVorticity)
{
    float density = 0.0f;
    for (LineSegment &line : lines) {
        density += line.length() * line.avgOpacity(maxVorticity);
    }
    return density;
}

float VoxelDiscretizer::computeDensityHair(float opacity)
{
    float density = 0.0f;
    for (LineSegment &line : lines) {
        density += line.length() * opacity;
    }
    return density;
}

VoxelCurveDiscretizer::VoxelCurveDiscretizer(const glm::ivec3 &gridResolution, const glm::ivec3 &quantizationResolution)
        : gridResolution(gridResolution), quantizationResolution(quantizationResolution)
{
    voxels = nullptr;
}

VoxelCurveDiscretizer::~VoxelCurveDiscretizer()
{
    delete[] voxels;
}

void VoxelCurveDiscretizer::setVoxelGrid(const sgl::AABB3 &aabb)
{
    glm::vec3 gridDimensions = aabb.getDimensions();
    glm::vec3 gridResolutionCubic = gridResolution;

    float maxDimensionLength = 0.0f;
    for (int i = 0; i < 3; i++) {
        maxDimensionLength = std::max(maxDimensionLength, gridDimensions[i]);
    }
    for (int i = 0; i < 3; i++) {
        float sideLengthFactor = gridDimensions[i] / maxDimensionLength;
        gridResolution[i] = (int)std::ceil(gridResolution[i] * sideLengthFactor);
    }

    voxels = new VoxelDiscretizer[gridResolution.x * gridResolution.y * gridResolution.z];
    for (int z = 0; z < gridResolution.z; z++) {
        for (int y = 0; y < gridResolution.y; y++) {
            for (int x = 0; x < gridResolution.x; x++) {
                int index = x + y*gridResolution.x + z*gridResolution.x*gridResolution.y;
                voxels[index].setIndex(glm::ivec3(x, y, z));
            }
        }
    }
}

VoxelGridDataCompressed VoxelCurveDiscretizer::createFromTrajectoryDataset(const std::string &filename,
        TrajectoryType trajectoryType, std::vector<float> &attributes, float &_maxVorticity,
        unsigned int maxNumLinesPerVoxel, bool useGPU)
{
    linesBoundingBox = sgl::AABB3();
    std::vector<Curve> curves;
    Curve currentCurve;
    maxVorticity = 0.0f;
    isHairDataset = false;
    uint64_t numLineSegments = 0;
    uint64_t numLines = 0;

    bool isRings = boost::starts_with(filename, "Data/Rings");
    bool isConvectionRolls = boost::starts_with(filename, "Data/ConvectionRolls/output");
    bool isConvectionRollsSmall = boost::starts_with(filename, "Data/ConvectionRolls/turbulence20000");

    Trajectories trajectories = loadTrajectoriesFromFile(filename, trajectoryType);

    for (size_t i = 0; i < trajectories.size(); i++) {
        Trajectory &trajectory = trajectories.at(i);

        currentCurve = Curve();
        for (size_t j = 0; j < trajectory.positions.size(); j++) {
            glm::vec3 &position = trajectory.positions.at(j);
            linesBoundingBox.combine(position);
            currentCurve.points.push_back(position);
            currentCurve.attributes.push_back(trajectory.attributes.at(0).at(j));
        }

        currentCurve.lineID = numLines;
        curves.push_back(currentCurve);
        numLineSegments += currentCurve.points.size() - 1;
        numLines++;
    }

    std::cout << "Num Lines: " << numLines << std::endl;
    std::cout << "Num LineSegments: " << numLineSegments << std::endl << std::flush;

    /*if (isRings) {
        linesBoundingBox = sgl::AABB3();
        linesBoundingBox.combine(glm::vec3(-2));
        linesBoundingBox.combine(glm::vec3(2));
    } else if (isConvectionRolls) {
        linesBoundingBox = sgl::AABB3();
        //linesBoundingBox.combine(glm::vec3(-1.5, -0.5, -1.5));
        //linesBoundingBox.combine(glm::vec3( 1.5, 0.5, 1.5));
        linesBoundingBox.combine(glm::vec3(-1.5, -1.5, -1.5));
        linesBoundingBox.combine(glm::vec3( 1.5, 1.5, 1.5));
    } else if (isConvectionRollsSmall) {
        linesBoundingBox = sgl::AABB3();
        linesBoundingBox.combine(glm::vec3(0.0, 0.0, 0.0));
        linesBoundingBox.combine(glm::vec3(1.0, 1.0, 1.0)); // 1.0, 1.0, 0.03
    }*/

    _maxVorticity = maxVorticity;
    this->attributes = attributes;

    // Move to origin and scale to range from (0, 0, 0) to (rx, ry, rz).
    setVoxelGrid(linesBoundingBox);
    linesToVoxel = sgl::matrixScaling(1.0f / linesBoundingBox.getDimensions() * glm::vec3(gridResolution))
            * sgl::matrixTranslation(-linesBoundingBox.getMinimum());
    voxelToLines = glm::inverse(linesToVoxel);

    std::cout << "Grid resolution: " << gridResolution.x << " " << gridResolution.y
              << " " << gridResolution.z << std::endl;
    std::cout << "Bounding box: " << linesBoundingBox.getMaximum().x << " " << linesBoundingBox.getMaximum().y
              << " " << linesBoundingBox.getMaximum().z << std::endl << std::flush;

    // Transform curves to voxel grid space
    for (Curve &curve : curves) {
        for (glm::vec3 &v : curve.points) {
            v = sgl::transformPoint(linesToVoxel, v);
        }
    }

    if (!useGPU) {
        // Insert lines into voxel representation
        //int lineNum = 0;
        for (const Curve &curve : curves) {
            /*if (lineNum == 1000) {
                break;
            }*/
            nextStreamline(curve);
            //lineNum++;
        }
        return compressData();
    } else {
        return createVoxelGridGPU(curves, maxNumLinesPerVoxel);
    }
}

VoxelGridDataCompressed VoxelCurveDiscretizer::createFromHairDataset(const std::string &filename, float &lineRadius,
        glm::vec4 &hairStrandColor, unsigned int maxNumLinesPerVoxel, bool useGPU)
{
    HairData hairData;
    loadHairFile(filename, hairData);
    downscaleHairData(hairData, HAIR_MODEL_SCALING_FACTOR);

    // Assume default thickness, opacity and color for now to simplify the implementation
    lineRadius = hairData.defaultThickness;
    hairStrandColor = glm::vec4(hairData.defaultColor, hairData.defaultOpacity);
    this->hairThickness = lineRadius;
    this->hairStrandColor = hairStrandColor;
    this->hairOpacity = hairData.defaultOpacity;

    linesBoundingBox = sgl::AABB3();
    std::vector<Curve> curves;
    Curve currentCurve;
    maxVorticity = 0.0f;
    int lineCounter = 0;
    isHairDataset = true;

    // Process all strands and convert them to curves
    for (HairStrand &strand : hairData.strands) {
        lineCounter++;
        if (lineCounter % 1000 == 999) {
            sgl::Logfile::get()->writeInfo(std::string() + "Parsing hair strand " + sgl::toString(lineCounter) + "...");
        }

        for (glm::vec3 point : strand.points) {
            glm::vec3 scaledPoint = point;
            currentCurve.points.push_back(scaledPoint);
            currentCurve.attributes.push_back(this->hairOpacity);
            linesBoundingBox.combine(scaledPoint);
        }

        curves.push_back(currentCurve);
        currentCurve = Curve();
        currentCurve.lineID = lineCounter;
    }

    // Move to origin and scale to range from (0, 0, 0) to (rx, ry, rz).
    setVoxelGrid(linesBoundingBox);
    linesToVoxel = sgl::matrixScaling(1.0f / linesBoundingBox.getDimensions() * glm::vec3(gridResolution))
                   * sgl::matrixTranslation(-linesBoundingBox.getMinimum());
    voxelToLines = glm::inverse(linesToVoxel);

    // Transform curves to voxel grid space
    for (Curve &curve : curves) {
        for (glm::vec3 &v : curve.points) {
            v = sgl::transformPoint(linesToVoxel, v);
        }
    }

    if (!useGPU) {
        // Insert lines into voxel representation
        for (const Curve &curve : curves) {
            nextStreamline(curve);
        }
        return compressData();
    } else {
        return createVoxelGridGPU(curves, maxNumLinesPerVoxel);
    }
}

VoxelGridDataCompressed VoxelCurveDiscretizer::compressData()
{
    VoxelGridDataCompressed dataCompressed;
    dataCompressed.gridResolution = gridResolution;
    dataCompressed.quantizationResolution = quantizationResolution;
    dataCompressed.worldToVoxelGridMatrix = this->getWorldToVoxelGridMatrix();
    dataCompressed.dataType = isHairDataset ? 1u : 0u;

    if (isHairDataset) {
        dataCompressed.hairStrandColor = hairStrandColor;
        dataCompressed.hairThickness = hairThickness;
    } else {
        dataCompressed.attributes = attributes;
        dataCompressed.maxVorticity = maxVorticity;
    }

    int n = gridResolution.x * gridResolution.y * gridResolution.z;
    std::vector<float> voxelDensities;
    std::vector<uint32_t> usedVoxels;
    voxelDensities.reserve(n);
    usedVoxels.reserve(n);

    size_t lineOffset = 0;
    dataCompressed.voxelLineListOffsets.reserve(n);
    dataCompressed.numLinesInVoxel.reserve(n);
    dataCompressed.lineSegments.clear();

    for (int i = 0; i < n; i++) {
        dataCompressed.voxelLineListOffsets.push_back(lineOffset);
        size_t numLines = voxels[i].lines.size();
        dataCompressed.numLinesInVoxel.push_back(numLines);
        if (isHairDataset) {
            voxelDensities.push_back(voxels[i].computeDensityHair(hairOpacity));
        } else {
            voxelDensities.push_back(voxels[i].computeDensity(maxVorticity));
        }
        usedVoxels.push_back(numLines > 0 ? 1 : 0);

#ifdef PACK_LINES
        std::vector<LineSegmentCompressed> lineSegments;
        lineSegments.resize(voxels[i].lines.size());
        for (size_t j = 0; j < voxels[i].lines.size(); j++) {
            compressLine(voxels[i].getIndex(), voxels[i].lines[j], lineSegments[j]);

            // Test
            /*LineSegment originalLine = voxels[i].lines[j];
            LineSegment decompressedLine;
            decompressLine(glm::vec3(voxels[i].getIndex()), lineSegments[j], decompressedLine);
            if (!checkLinesEqual(originalLine, decompressedLine)) {
                compressLine(voxels[i].getIndex(), voxels[i].lines[j], lineSegments[j]);
                decompressLine(glm::vec3(voxels[i].getIndex()), lineSegments[j], decompressedLine);
            }*/
        }
        dataCompressed.lineSegments.insert(dataCompressed.lineSegments.end(),
                lineSegments.begin(), lineSegments.end());
#else
        dataCompressed.lineSegments.insert(dataCompressed.lineSegments.end(),
                voxels[i].lines.begin(), voxels[i].lines.end());
#endif

        lineOffset += numLines;
    }

    std::vector<float> voxelAOFactors;
    voxelAOFactors.resize(n);
    generateVoxelAOFactorsFromDensity(voxelDensities, voxelAOFactors, gridResolution, isHairDataset);

    dataCompressed.voxelDensities = voxelDensities;
    dataCompressed.voxelAOFactors = voxelAOFactors;
    return dataCompressed;
}

std::vector<VoxelDiscretizer*> VoxelCurveDiscretizer::getVoxelsInAABB(const sgl::AABB3 &aabb)
{
    std::vector<VoxelDiscretizer*> voxelsInAABB;
    glm::vec3 minimum = aabb.getMinimum();
    glm::vec3 maximum = aabb.getMaximum();

    glm::ivec3 lower = glm::ivec3(minimum); // Round down
    glm::ivec3 upper = glm::ivec3(ceil(maximum.x), ceil(maximum.y), ceil(maximum.z)); // Round up
    lower = glm::max(lower, glm::ivec3(0));
    upper = glm::min(upper, gridResolution - glm::ivec3(1));

    for (int z = lower.z; z <= upper.z; z++) {
        for (int y = lower.y; y <= upper.y; y++) {
            for (int x = lower.x; x <= upper.x; x++) {
                int index = x + y*gridResolution.x + z*gridResolution.x*gridResolution.y;
                voxelsInAABB.push_back(voxels + index);
            }
        }
    }
    return voxelsInAABB;
}

void VoxelCurveDiscretizer::nextStreamline(const Curve &line)
{
    int N = line.points.size();

    // Add intersections to voxels
    std::set<VoxelDiscretizer*> usedVoxels;
    for (int i = 0; i < N-1; i++) {
        // Get line segment
        glm::vec3 v1 = line.points.at(i);
        glm::vec3 v2 = line.points.at(i+1);
        float a1 = line.attributes.at(i);
        float a2 = line.attributes.at(i+1);

        // Remove invalid line points (used in many scientific datasets to indicate invalid lines).
        const float MAX_VAL = 1e10;
        if (std::fabs(v1.x) > MAX_VAL || std::fabs(v1.y) > MAX_VAL || std::fabs(v1.z) > MAX_VAL
                || std::fabs(v2.x) > MAX_VAL || std::fabs(v2.y) > MAX_VAL || std::fabs(v2.z) > MAX_VAL) {
            continue;
        }

        // Compute AABB of current segment
        sgl::AABB3 segmentAABB = sgl::AABB3();
        segmentAABB.combine(v1);
        segmentAABB.combine(v2);

        // Iterate over all voxels with possible intersections
        std::vector<VoxelDiscretizer*> voxelsInAABB = getVoxelsInAABB(segmentAABB);

        for (VoxelDiscretizer *voxel : voxelsInAABB) {
            // Line-voxel intersection
            if (voxel->addPossibleIntersections(v1, v2, a1, a2)) {
                // Intersection(s) added to "currentLineIntersections", voxel used
                usedVoxels.insert(voxel);
            }
        }
    }

    // Convert intersections to clipped line segments
    for (VoxelDiscretizer *voxel : usedVoxels) {
        if (voxel->currentCurveIntersections.size() < 2) {
            voxel->currentCurveIntersections.clear();
            continue;
        }
        auto it1 = voxel->currentCurveIntersections.begin();
        auto it2 = voxel->currentCurveIntersections.begin();
        it2++;
        while (it2 != voxel->currentCurveIntersections.end()) {
            voxel->lines.push_back(LineSegment(it1->v, it1->a, it2->v, it2->a, line.lineID));
            it1++; it1++;
            if (it1 == voxel->currentCurveIntersections.end()) break;
            it2++; it2++;
        }
        voxel->currentCurveIntersections.clear();
    }
}

template<typename T>
T clamp(T x, T a, T b) {
    if (x < a) {
        return a;
    } else if (x > b) {
        return b;
    } else {
        return x;
    }
}

void VoxelCurveDiscretizer::quantizeLine(const glm::vec3 &voxelPos, const LineSegment &line,
        LineSegmentQuantized &lineQuantized, int faceIndex1, int faceIndex2)
{
    lineQuantized.a1 = line.a1;
    lineQuantized.a2 = line.a2;

    glm::ivec2 facePosition3D1, facePosition3D2;
    quantizePoint(line.v1 - voxelPos, facePosition3D1, faceIndex1);
    quantizePoint(line.v2 - voxelPos, facePosition3D2, faceIndex2);
    lineQuantized.lineID = line.lineID;
    lineQuantized.faceIndex1 = faceIndex1;
    lineQuantized.faceIndex2 = faceIndex2;
    lineQuantized.facePositionQuantized1 = facePosition3D1.x + facePosition3D1.y*quantizationResolution.x;
    lineQuantized.facePositionQuantized2 = facePosition3D2.x + facePosition3D2.y*quantizationResolution.x;
}

int intlog2(int x) {
    int exponent = 0;
    while (x > 1) {
        x /= 2;
        exponent++;
    }
    return exponent;
}

void VoxelCurveDiscretizer::compressLine(const glm::ivec3 &voxelIndex, const LineSegment &line,
        LineSegmentCompressed &lineCompressed)
{
    LineSegmentQuantized lineQuantized;
    int faceIndex1 = computeFaceIndex(line.v1, voxelIndex);
    int faceIndex2 = computeFaceIndex(line.v2, voxelIndex);
    quantizeLine(glm::vec3(voxelIndex), line, lineQuantized, faceIndex1, faceIndex2);

    uint8_t attr1Unorm = std::round(lineQuantized.a1*255.0f);
    uint8_t attr2Unorm = std::round(lineQuantized.a2*255.0f);

    int c = round(2*intlog2(quantizationResolution.x));
    lineCompressed.linePosition = lineQuantized.faceIndex1;
    lineCompressed.linePosition |= lineQuantized.faceIndex2 << 3;
    lineCompressed.linePosition |= lineQuantized.facePositionQuantized1 << 6;
    lineCompressed.linePosition |= lineQuantized.facePositionQuantized2 << (6 + c);
    lineCompressed.attributes = 0;
    if (c > 12) {
        // Quantization resolution of 128 or 256
        lineCompressed.attributes |= lineQuantized.facePositionQuantized2 >> (c - (6 + 2*c - 32));
    }
    lineCompressed.attributes |= (lineQuantized.lineID & 31u) << 11;
    lineCompressed.attributes |= attr1Unorm << 16;
    lineCompressed.attributes |= attr2Unorm << 24;
}

void VoxelCurveDiscretizer::quantizePoint(const glm::vec3 &v, glm::ivec2 &qv, int faceIndex)
{
    int dimensions[2];
    if (faceIndex == 0 || faceIndex == 1) {
        // x face
        dimensions[0] = 1;
        dimensions[1] = 2;
    } else if (faceIndex == 2 || faceIndex == 3) {
        // y face
        dimensions[0] = 0;
        dimensions[1] = 2;
    } else {
        // z face
        dimensions[0] = 0;
        dimensions[1] = 1;
    }

    // Iterate over all dimensions
    for (int i = 0; i < 2; i++) {
        int quantizationPos = std::floor(v[dimensions[i]] * quantizationResolution[dimensions[i]]);
        qv[i] = glm::clamp(quantizationPos, 0, quantizationResolution[dimensions[i]]-1);
    }
}

int VoxelCurveDiscretizer::computeFaceIndex(const glm::vec3 &v, const glm::ivec3 &voxelIndex)
{
    glm::ivec3 lower = voxelIndex, upper = voxelIndex + glm::ivec3(1);
    for (int i = 0; i < 3; i++) {
        if (std::abs(v[i] - lower[i]) < 0.00001f) {
            return 2*i;
        }
        if (std::abs(v[i] - upper[i]) < 0.00001f) {
            return 2*i+1;
        }
    }
    sgl::Logfile::get()->writeError(std::string() + "Error in VoxelCurveDiscretizer::computeFaceIndex: "
            + "Invalid position.");
    return 0;
}

glm::vec3 VoxelCurveDiscretizer::getQuantizedPositionOffset(uint32_t faceIndex, uint32_t quantizedPos1D)
{
    glm::vec2 quantizedFacePosition = glm::vec2(
            float(quantizedPos1D % quantizationResolution.x),
            float(quantizedPos1D / quantizationResolution.x))
                    / float(quantizationResolution.x);

    // Whether the face is the face in x/y/z direction with greater dimensions (offset factor)
    float face0or1 = float(faceIndex % 2);

    glm::vec3 offset;
    if (faceIndex <= 1) {
        offset = glm::vec3(face0or1, quantizedFacePosition.x, quantizedFacePosition.y);
    } else if (faceIndex <= 3) {
        offset = glm::vec3(quantizedFacePosition.x, face0or1, quantizedFacePosition.y);
    } else if (faceIndex <= 5) {
        offset = glm::vec3(quantizedFacePosition.x, quantizedFacePosition.y, face0or1);
    }
    return offset;
}

void VoxelCurveDiscretizer::decompressLine(const glm::vec3 &voxelPosition, const LineSegmentCompressed &compressedLine,
        LineSegment &decompressedLine)
{
    const uint32_t c = 2*intlog2(quantizationResolution.x);
    const uint32_t bitmaskQuantizedPos = quantizationResolution.x*quantizationResolution.x-1;
    uint32_t faceStartIndex = compressedLine.linePosition & 0x7u;
    uint32_t faceEndIndex = (compressedLine.linePosition >> 3) & 0x7u;
    uint32_t quantizedStartPos1D = (compressedLine.linePosition >> 6) & bitmaskQuantizedPos;
    uint32_t quantizedEndPos1D = (compressedLine.linePosition >> 6+c) & bitmaskQuantizedPos;
    if (c > 12) {
        quantizedEndPos1D |= (compressedLine.attributes << (c - (6 + 2*c - 32))) & bitmaskQuantizedPos;
    }
    uint32_t lineID = (compressedLine.attributes >> 11) & 31u;
    uint32_t attr1 = (compressedLine.attributes >> 16) & 0xFFu;
    uint32_t attr2 = (compressedLine.attributes >> 24) & 0xFFu;

    decompressedLine.v1 = voxelPosition + getQuantizedPositionOffset(faceStartIndex, quantizedStartPos1D);
    decompressedLine.v2 = voxelPosition + getQuantizedPositionOffset(faceEndIndex, quantizedEndPos1D);
    decompressedLine.a1 = float(attr1) / 255.0f;
    decompressedLine.a2 = float(attr2) / 255.0f;
    decompressedLine.lineID = lineID;
}

bool VoxelCurveDiscretizer::checkLinesEqual(const LineSegment &originalLine, const LineSegment &decompressedLine)
{
    bool linesEqual = true;
    if (originalLine.lineID % 256 != decompressedLine.lineID) {
        linesEqual = false;
        sgl::Logfile::get()->writeError("VoxelCurveDiscretizer::checkLinesEqual: lineID");
    }

    if (std::abs(originalLine.a1 - decompressedLine.a1) > 0.01f
            || std::abs(originalLine.a2 - decompressedLine.a2) > 0.01f) {
        linesEqual = false;
        sgl::Logfile::get()->writeError("VoxelCurveDiscretizer::checkLinesEqual: attribute");
    }

    if (glm::length(originalLine.v1 - decompressedLine.v1) > 0.5f) {
        linesEqual = false;
        sgl::Logfile::get()->writeError(std::string() + "VoxelCurveDiscretizer::checkLinesEqual: position2, error: "
                + sgl::toString(glm::length(originalLine.v2 - decompressedLine.v2)));
    }

    if (glm::length(originalLine.v2 - decompressedLine.v2) > 0.5f) {
        linesEqual = false;
        sgl::Logfile::get()->writeError(std::string() + "VoxelCurveDiscretizer::checkLinesEqual: position1, error: "
                                        + sgl::toString(glm::length(originalLine.v2 - decompressedLine.v2)));
    }

    return linesEqual;
}

struct LinePoint {
    LinePoint(glm::vec3 linePoint, float lineAttribute) : linePoint(linePoint), lineAttribute(lineAttribute) {}
    glm::vec3 linePoint;
    float lineAttribute;
};

std::string ivec3ToString(const glm::ivec3 &v) {
    return std::string() + "ivec3(" + sgl::toString(v.x) + ", " + sgl::toString(v.y) + ", " + sgl::toString(v.z) + ")";
}

void VoxelCurveDiscretizer::recreateDensityAndAOFactors(VoxelGridDataCompressed &dataCompressed,
        VoxelGridDataGPU &dataGPU, unsigned int maxNumLinesPerVoxel)
{
    glm::ivec3 numWorkGroupsVoxel = glm::ivec3(sgl::iceil(gridResolution.x, 64), sgl::iceil(gridResolution.y, 4),
                                               gridResolution.z);
    uint32_t gridSize1D = gridResolution.x *gridResolution.y *gridResolution.z;
    uint32_t zeroData = 0u;
    void *bufferMemory;

    // Set preprocessor defines for the shaders.
    sgl::ShaderManager->addPreprocessorDefine("MAX_NUM_LINES_PER_VOXEL", maxNumLinesPerVoxel);
    sgl::ShaderManager->addPreprocessorDefine("gridResolution", ivec3ToString(gridResolution));
    sgl::ShaderManager->addPreprocessorDefine(
            "GRID_RESOLUTION_LOG2", sgl::toString(sgl::intlog2(gridResolution.x)));
    sgl::ShaderManager->addPreprocessorDefine("GRID_RESOLUTION", gridResolution.x);
    sgl::ShaderManager->addPreprocessorDefine("quantizationResolution", ivec3ToString(quantizationResolution));
    sgl::ShaderManager->addPreprocessorDefine("QUANTIZATION_RESOLUTION", sgl::toString(quantizationResolution.x));
    sgl::ShaderManager->addPreprocessorDefine(
            "QUANTIZATION_RESOLUTION_LOG2", sgl::toString(sgl::intlog2(quantizationResolution.x)));

    // PART 3: Compute the densities
    auto startDensity = std::chrono::system_clock::now();

    sgl::ShaderProgramPtr computeDensityShader = sgl::ShaderManager->getShaderProgram({"RecomputeDensity.Compute"});
    computeDensityShader->setUniformImageTexture(0, dataGPU.densityTexture, GL_R32F, GL_READ_WRITE, 0, true, 0);
    computeDensityShader->dispatchCompute(numWorkGroupsVoxel.x, numWorkGroupsVoxel.y, numWorkGroupsVoxel.z);
    glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT);

    auto endDensity = std::chrono::system_clock::now();
    auto elapsedDensity = std::chrono::duration_cast<std::chrono::milliseconds>(endDensity - startDensity);
    sgl::Logfile::get()->writeInfo(std::string() + "Computational time to compute the densities: "
                                   + std::to_string(elapsedDensity.count()));

    // PART 5: Compute the ambient occlusion factors on the GPU using the density texture.
    auto startAO_GPU = std::chrono::system_clock::now();

    const int FILTER_SIZE = 7;
    const int FILTER_EXTENT = (FILTER_SIZE - 1) / 2;
    const int FILTER_NUM_FIELDS = FILTER_SIZE*FILTER_SIZE*FILTER_SIZE;
    float blurKernel[FILTER_NUM_FIELDS];
    generateGaussianBlurKernel(blurKernel, FILTER_SIZE, FILTER_EXTENT);
    sgl::GeometryBufferPtr gaussianKernelBuffer = sgl::Renderer->createGeometryBuffer(
            FILTER_NUM_FIELDS * sizeof(float), &blurKernel,
            sgl::UNIFORM_BUFFER, sgl::BUFFER_STATIC);

    sgl::ShaderProgramPtr computeAOShader = sgl::ShaderManager->getShaderProgram({"ComputeAO.Compute"});
    computeAOShader->setUniformImageTexture(0, dataGPU.densityTexture, GL_R32F, GL_READ_WRITE, 0, true, 0);
    computeAOShader->setUniform("densityTexture", dataGPU.densityTexture, 0);
    sgl::ShaderManager->bindShaderStorageBuffer(6, gaussianKernelBuffer);
    computeAOShader->dispatchCompute(numWorkGroupsVoxel.x, numWorkGroupsVoxel.y, numWorkGroupsVoxel.z);
    glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT);

    auto endAO_GPU = std::chrono::system_clock::now();
    auto elapsedAO_GPU = std::chrono::duration_cast<std::chrono::milliseconds>(endAO_GPU - startAO_GPU);
    sgl::Logfile::get()->writeInfo(std::string() + "Computational time to compute the ambient occlusion factors (GPU): "
                                   + std::to_string(elapsedAO_GPU.count()));

    glUseProgram(0); // For ImGui to stop complaining when binding last_program...
}

VoxelGridDataCompressed VoxelCurveDiscretizer::createVoxelGridGPU(
        std::vector<Curve> &curves, unsigned int maxNumLinesPerVoxel)
{
    glm::ivec3 numWorkGroupsVoxel = glm::ivec3(sgl::iceil(gridResolution.x, 64), sgl::iceil(gridResolution.y, 4),
            gridResolution.z);
    uint32_t gridSize1D = gridResolution.x *gridResolution.y *gridResolution.z;
    uint32_t zeroData = 0u;
    void *bufferMemory;

    // Set preprocessor defines for the shaders.
    sgl::ShaderManager->addPreprocessorDefine("MAX_NUM_LINES_PER_VOXEL", maxNumLinesPerVoxel);
    sgl::ShaderManager->addPreprocessorDefine("gridResolution", ivec3ToString(gridResolution));
    sgl::ShaderManager->addPreprocessorDefine(
            "GRID_RESOLUTION_LOG2", sgl::toString(sgl::intlog2(gridResolution.x)));
    sgl::ShaderManager->addPreprocessorDefine("GRID_RESOLUTION", gridResolution.x);
    sgl::ShaderManager->addPreprocessorDefine("quantizationResolution", ivec3ToString(quantizationResolution));
    sgl::ShaderManager->addPreprocessorDefine("QUANTIZATION_RESOLUTION", sgl::toString(quantizationResolution.x));
    sgl::ShaderManager->addPreprocessorDefine(
            "QUANTIZATION_RESOLUTION_LOG2", sgl::toString(sgl::intlog2(quantizationResolution.x)));
    if (isHairDataset) {
        sgl::ShaderManager->addPreprocessorDefine("HAIR_RENDERING", maxNumLinesPerVoxel);
    } else {
        sgl::ShaderManager->removePreprocessorDefine("HAIR_RENDERING");
    }

    // PART 1: Create the LinePointBuffer, LineOffsetBuffer, NumSegmentsBuffer (empty) and LineSegmentsBuffer.
    auto startBuffers = std::chrono::system_clock::now();
    std::vector<LinePoint> linePoints;
    std::vector<uint32_t> lineOffsets;
    lineOffsets.push_back(0);
    uint32_t offsetCounter = 0;
    for (Curve &curve : curves) {
        size_t curveNumPoints = curve.points.size();
        for (size_t i = 0; i < curveNumPoints; i++) {
            linePoints.push_back(LinePoint(curve.points.at(i), curve.attributes.at(i)));
        }
        offsetCounter += curveNumPoints;
        lineOffsets.push_back(offsetCounter);
    }
    sgl::GeometryBufferPtr linePointBuffer = sgl::Renderer->createGeometryBuffer(
            (linePoints.size()+1) * sizeof(LinePoint), &linePoints.front(),
            sgl::SHADER_STORAGE_BUFFER, sgl::BUFFER_STATIC);
    sgl::GeometryBufferPtr lineOffsetBuffer = sgl::Renderer->createGeometryBuffer(
            (curves.size()+1) * sizeof(uint32_t), &lineOffsets.front(),
            sgl::SHADER_STORAGE_BUFFER, sgl::BUFFER_STATIC);
    sgl::GeometryBufferPtr numSegmentsBuffer = sgl::Renderer->createGeometryBuffer(
            gridSize1D * sizeof(uint32_t),
            sgl::SHADER_STORAGE_BUFFER, sgl::BUFFER_STATIC);
    GLuint bufferID = ((sgl::GeometryBufferGL*)numSegmentsBuffer.get())->getBuffer();
    glClearNamedBufferData(bufferID, GL_R32UI, GL_RED_INTEGER, GL_UNSIGNED_INT, (const void*)&zeroData);
    sgl::GeometryBufferPtr lineSegmentsBuffer = sgl::Renderer->createGeometryBuffer(
            maxNumLinesPerVoxel * gridSize1D * sizeof(LineSegmentCompressed),
            sgl::SHADER_STORAGE_BUFFER, sgl::BUFFER_STATIC);

    auto endBuffers = std::chrono::system_clock::now();
    auto elapsedBuffers = std::chrono::duration_cast<std::chrono::milliseconds>(endBuffers - startBuffers);
    sgl::Logfile::get()->writeInfo(std::string() + "Computational time to create the buffers: "
                                   + std::to_string(elapsedBuffers.count()));

    // PART 2: Discretize, quantize and voxelize the lines.
    auto startVoxelize = std::chrono::system_clock::now();
    unsigned int numWorkGroupsLines = sgl::iceil(curves.size(), 256);
    sgl::ShaderProgramPtr discretizeLinesShader = sgl::ShaderManager->getShaderProgram({"DiscretizeLines.Compute"});
    sgl::ShaderManager->bindShaderStorageBuffer(2, linePointBuffer);
    sgl::ShaderManager->bindShaderStorageBuffer(3, lineOffsetBuffer);
    sgl::ShaderManager->bindShaderStorageBuffer(4, numSegmentsBuffer);
    sgl::ShaderManager->bindShaderStorageBuffer(5, lineSegmentsBuffer);
    discretizeLinesShader->setUniform("numLines", (unsigned int)curves.size());
    discretizeLinesShader->dispatchCompute(numWorkGroupsLines);
    glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
    glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT);

    // End of PART 2: Read the line segment buffer & number of line segments per voxel buffer back from the GPU.
    std::vector<LineSegmentCompressed> compressedLineSegments;
    bufferMemory = lineSegmentsBuffer->mapBuffer(sgl::BUFFER_MAP_READ_ONLY);
    compressedLineSegments.resize(maxNumLinesPerVoxel * gridSize1D);
    memcpy(&compressedLineSegments.front(), bufferMemory, compressedLineSegments.size() * sizeof(LineSegmentCompressed));
    lineSegmentsBuffer->unmapBuffer();

    std::vector<uint32_t> numSegmentsPerVoxel;
    bufferMemory = numSegmentsBuffer->mapBuffer(sgl::BUFFER_MAP_READ_ONLY);
    numSegmentsPerVoxel.resize(gridSize1D);
    memcpy(&numSegmentsPerVoxel.front(), bufferMemory, numSegmentsPerVoxel.size() * sizeof(uint32_t));
    numSegmentsBuffer->unmapBuffer();

    auto endVoxelize = std::chrono::system_clock::now();
    auto elapsedVoxelize = std::chrono::duration_cast<std::chrono::milliseconds>(endVoxelize - startVoxelize);
    sgl::Logfile::get()->writeInfo(std::string() + "Computational time to voxelize the lines: "
                                   + std::to_string(elapsedVoxelize.count()));

    // PART 3: Compute the densities
    auto startDensity = std::chrono::system_clock::now();

    sgl::TextureSettings densityTextureSettings = sgl::TextureSettings();
    densityTextureSettings.type = sgl::TEXTURE_3D;
    densityTextureSettings.internalFormat = GL_R32F;
    sgl::TexturePtr densityTexture = sgl::TextureManager->createEmptyTexture(
            gridResolution.x, gridResolution.y, gridResolution.z, densityTextureSettings);
    sgl::ShaderProgramPtr computeDensityShader = sgl::ShaderManager->getShaderProgram({"ComputeDensity.Compute"});
    computeDensityShader->setUniformImageTexture(0, densityTexture, GL_R32F, GL_READ_WRITE, 0, true, 0);
    computeDensityShader->dispatchCompute(numWorkGroupsVoxel.x, numWorkGroupsVoxel.y, numWorkGroupsVoxel.z);
    glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT);

    // End of PART 3: Read the density values back from the GPU.
    std::vector<float> voxelDensities;
    voxelDensities.resize(gridSize1D);
    sgl::TextureGL *densityTextureGL = (sgl::TextureGL*)densityTexture.get();
    glGetTextureImage(densityTextureGL->getTexture(), 0, GL_RED, GL_FLOAT,
            sizeof(float) * gridSize1D, (void*)&voxelDensities.front());

    auto endDensity = std::chrono::system_clock::now();
    auto elapsedDensity = std::chrono::duration_cast<std::chrono::milliseconds>(endDensity - startDensity);
    sgl::Logfile::get()->writeInfo(std::string() + "Computational time to compute the densities: "
                                   + std::to_string(elapsedDensity.count()));

    // PART 4: Reduce the size of the buffer using a prefix sum (on the CPU for now).
    auto startPrefixSum = std::chrono::system_clock::now();

    uint32_t lineSegmentOffset = 0;
    std::vector<uint32_t> lineSegmentOffsets;
    std::vector<LineSegmentCompressed> reducedLineSegmentBuffer;
    for (size_t i = 0; i < numSegmentsPerVoxel.size(); i++) {
        lineSegmentOffsets.push_back(lineSegmentOffset);
        size_t numSegmentsCurrentVoxel = numSegmentsPerVoxel.at(i);
        for (size_t j = 0; j < numSegmentsCurrentVoxel; j++) {
            reducedLineSegmentBuffer.push_back(compressedLineSegments.at(i*maxNumLinesPerVoxel+j));
        }
        lineSegmentOffset += numSegmentsCurrentVoxel;
    }

    uint32_t offsetTest = lineSegmentOffsets.back();
    LineSegmentCompressed lineSegmentTest = reducedLineSegmentBuffer.back();
    std::vector<LineSegmentCompressed> testArray = reducedLineSegmentBuffer;
    std::reverse(testArray.begin(), testArray.end());

    auto endPrefixSum = std::chrono::system_clock::now();
    auto elapsedPrefixSum = std::chrono::duration_cast<std::chrono::milliseconds>(endPrefixSum - startPrefixSum);
    sgl::Logfile::get()->writeInfo(std::string() + "Computational time to reduce the buffers: "
                                   + std::to_string(elapsedPrefixSum.count()));

    // PART 5: Compute the ambient occlusion factors. For now, do this on CPU (legacy).
    /*auto startAO_CPU = std::chrono::system_clock::now();

    std::vector<float> voxelAOFactors;
    voxelAOFactors.resize(gridSize1D);
    generateVoxelAOFactorsFromDensity(voxelDensities, voxelAOFactors, gridResolution, isHairDataset);

    auto endAO_CPU = std::chrono::system_clock::now();
    auto elapsedAO_CPU = std::chrono::duration_cast<std::chrono::milliseconds>(endAO_CPU - startAO_CPU);
    sgl::Logfile::get()->writeInfo(std::string() + "Computational time to compute the ambient occlusion factors (CPU): "
                                   + std::to_string(elapsedAO_CPU.count()));*/

    // PART 5: Compute the ambient occlusion factors on the GPU using the density texture.
    auto startAO_GPU = std::chrono::system_clock::now();
    const int FILTER_SIZE = 7;
    const int FILTER_EXTENT = (FILTER_SIZE - 1) / 2;
    const int FILTER_NUM_FIELDS = FILTER_SIZE*FILTER_SIZE*FILTER_SIZE;
    float blurKernel[FILTER_NUM_FIELDS];
    generateGaussianBlurKernel(blurKernel, FILTER_SIZE, FILTER_EXTENT);
    sgl::GeometryBufferPtr gaussianKernelBuffer = sgl::Renderer->createGeometryBuffer(
            FILTER_NUM_FIELDS * sizeof(float), &blurKernel,
            sgl::SHADER_STORAGE_BUFFER, sgl::BUFFER_STATIC);

    sgl::TextureSettings aoTextureSettings = sgl::TextureSettings();
    aoTextureSettings.type = sgl::TEXTURE_3D;
    aoTextureSettings.internalFormat = GL_R32F;
    sgl::TexturePtr aoTexture = sgl::TextureManager->createEmptyTexture(
            gridResolution.x, gridResolution.y, gridResolution.z, aoTextureSettings);
    sgl::ShaderProgramPtr computeAOShader = sgl::ShaderManager->getShaderProgram({"ComputeAO.Compute"});
    computeAOShader->setUniformImageTexture(0, aoTexture, GL_R32F, GL_READ_WRITE, 0, true, 0);
    computeAOShader->setUniform("densityTexture", densityTexture, 0);
    sgl::ShaderManager->bindShaderStorageBuffer(6, gaussianKernelBuffer);
    computeAOShader->dispatchCompute(numWorkGroupsVoxel.x, numWorkGroupsVoxel.y, numWorkGroupsVoxel.z);
    glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT);

    // End of PART 5: Read the AO factors back from the GPU.
    std::vector<float> voxelAOFactors;
    voxelAOFactors.resize(gridSize1D);
    sgl::TextureGL *aoTextureGL = (sgl::TextureGL*)aoTexture.get();
    glGetTextureImage(aoTextureGL->getTexture(), 0, GL_RED, GL_FLOAT,
                      sizeof(float) * gridSize1D, (void*)&voxelAOFactors.front());
    normalizeVoxelAOFactors(voxelAOFactors, gridResolution, isHairDataset);

    auto endAO_GPU = std::chrono::system_clock::now();
    auto elapsedAO_GPU = std::chrono::duration_cast<std::chrono::milliseconds>(endAO_GPU - startAO_GPU);
    sgl::Logfile::get()->writeInfo(std::string() + "Computational time to compute the ambient occlusion factors (GPU): "
                                   + std::to_string(elapsedAO_GPU.count()));

    glUseProgram(0); // For ImGui to stop complaining when binding last_program...

    // FINAL STEP: Now, write the data to the struct.
    VoxelGridDataCompressed dataCompressed;
    dataCompressed.gridResolution = gridResolution;
    dataCompressed.quantizationResolution = quantizationResolution;
    dataCompressed.worldToVoxelGridMatrix = this->getWorldToVoxelGridMatrix();
    dataCompressed.dataType = isHairDataset ? 1u : 0u;

    if (isHairDataset) {
        dataCompressed.hairStrandColor = hairStrandColor;
        dataCompressed.hairThickness = hairThickness;
    } else {
        dataCompressed.attributes = attributes;
        dataCompressed.maxVorticity = maxVorticity;
    }

    dataCompressed.voxelLineListOffsets = lineSegmentOffsets;
    dataCompressed.numLinesInVoxel = numSegmentsPerVoxel;
    dataCompressed.lineSegments = reducedLineSegmentBuffer;

    dataCompressed.voxelDensities = voxelDensities;
    dataCompressed.voxelAOFactors = voxelAOFactors;
    return dataCompressed;
}


#include "rnnlm.h"

#include "math-utils.h"

#include <iostream>
#include <fstream>
#include <sstream>
#include <algorithm>

#include <boost/program_options/parsers.hpp>
#include <boost/program_options/variables_map.hpp>
#include <boost/algorithm/string.hpp>

using namespace std;
using namespace dynet;
using namespace boost::program_options;

unsigned LAYERS = 2;
unsigned INPUT_DIM = 256;
unsigned HIDDEN_DIM = 512;
unsigned BATCH_SIZE = 1;
unsigned VOCAB_SIZE = 0;
unsigned ETA_EPOCH = 0;
unsigned MAX_EPOCH = 20;
dynet::real LAMBDA;
dynet::real ETA = 0.1f;
dynet::real ETA_DECAY = 2;
unsigned TREPORT = 20;
unsigned DREPORT = 200;
double MAX_LOSS = 9e+99;

dynet::Dict d;// vocabulary

bool verbose = false;

typedef std::vector<int> Sentence;
typedef std::vector<Sentence> Corpus;

Corpus Read_Corpus(const string &fp, bool dir=true);

inline size_t CalcSize(const Sentence & src, int trg);
inline void Create_MiniBatches(const Corpus& traincor, size_t max_size,
	std::vector<std::vector<Sentence> > & traincor_minibatch);

template <class RNNLM_t>
void TrainModel(Model &model, RNNLM_t &rnn
	, Corpus &traincor, Corpus &devcor
	, Trainer* sgd
	, const string& param_file, float dropout_p);
template <class RNNLM_t>
void TrainModel_Batch1(Model &model, RNNLM_t &rnn
	, Corpus &traincor, Corpus &devcor
	, Trainer* sgd
	, const string& param_file, float dropout_p);
template <class RNNLM_t>
void TrainModel_Batch2(Model &model, RNNLM_t &rnn
	, Corpus &traincor, Corpus &devcor
	, Trainer* sgd
	, const string& param_file, float dropout_p);

template <class RNNLM_t>
void TestModel(Model &model, RNNLM_t &rnn
	, const Corpus &testcor);

template <class Builder>
int main_body(variables_map vm);

// Sort in descending order of length
struct CompareLen {
	bool operator()(const std::vector<int>& first, const std::vector<int>& second) {
		return first.size() > second.size();
	}
};

int main(int argc, char** argv) {
	dynet::initialize(argc, argv);

	// command line processing
	using namespace boost::program_options;
	variables_map vm;
	options_description opts("Allowed options");
	opts.add_options()
		("help", "print help message")
		("train,t", value<string>(), "file containing training sentences")
		("dev,d", value<string>(), "file containing development sentences")
		("test,T", value<string>(), "file containing testing source sentences for computing perplexity scores")
		("lang", value<bool>()->default_value(true), "use source or target language from a corpus, e.g., source ||| target")
		("initialise,l", value<string>(), "load initial parameters from file")
		("parameters,s", value<string>(), "save best parameters to this file")
		("layers", value<unsigned>()->default_value(LAYERS), "use <num> layers for RNN components")
		("hidden_dim,h", value<unsigned>()->default_value(HIDDEN_DIM), "use <num> dimensions for recurrent hidden states")
		("input_dim,i", value<unsigned>()->default_value(INPUT_DIM), "use <num> dimensions for input embedding")
		("batch_size,b", value<unsigned>()->default_value(BATCH_SIZE), "impose mini-batch size for training")
		("gru", "use Gated Recurrent Unit (GRU) for recurrent structure; default RNN")
		("lstm", "use Long Short Term Memory (LSTM) for recurrent structure; default RNN")
		("vlstm", "use Vanilla Long Short Term Memory (VLSTM) for recurrent structure; default RNN")
		("dglstm", "use Depth-Gated Long Short Term Memory (DGLSTM) for recurrent structure; default RNN")
		("dropout", value<float>(), "apply dropout technique (Gal et al., 2015) for RNN")
		("sgd_trainer", value<unsigned>()->default_value(0), "use specific SGD trainer (0: vanilla SGD; 1: momentum SGD; 2: Adagrad; 3: AdaDelta; 4: Adam)")
		("sparse_updates", value<bool>()->default_value(true), "enable/disable sparse update(s) for lookup parameter(s)")
		("r2l_target", "use right-to-left direction for target during training; default not")
		("max_loss", value<double>()->default_value(MAX_LOSS), "impose max loss if possible (for interruped running jobs)")
		("treport", value<unsigned>()->default_value(TREPORT), "impose instance reporting value (during training); default 50")
		("dreport", value<unsigned>()->default_value(DREPORT), "impose instance reporting value (during training) on development/test sets; default 500")
		("eta_epoch", value<unsigned>()->default_value(ETA_EPOCH), "limit number of epochs for scheduling the learning rate, e.g., decreasing by a factor of eta_decay; default 4")
		("max_epoch", value<unsigned>()->default_value(MAX_EPOCH), "limit number of epochs for training; default 20")
		("lambda", value<dynet::real>()->default_value(1e-6), "the L2 regularization coefficient; default 1e-6.")
		("eta", value<dynet::real>()->default_value(ETA), "the learning rate of neural network model; default 0.1")
		("eta_decay", value<dynet::real>()->default_value(ETA_DECAY), "decay value for decreasing the learning rate; default 2.")
		("seed", value<unsigned>()->default_value(1), "value for random initialization; default 1")
		("verbose,v", "be extremely noisy");

	store(parse_command_line(argc, argv, opts), vm);

	notify(vm);

	verbose = vm.count("verbose");
	ETA_EPOCH = vm["eta_epoch"].as<unsigned>();
	MAX_EPOCH = vm["max_epoch"].as<unsigned>();
	MAX_LOSS = vm["max_loss"].as<double>();
	TREPORT = vm["treport"].as<unsigned>();
	DREPORT = vm["dreport"].as<unsigned>();
	LAMBDA = vm["lambda"].as<dynet::real>();
	ETA_DECAY = vm["eta_decay"].as<dynet::real>();
	ETA = vm["eta"].as<dynet::real>();
	LAYERS = vm["layers"].as<unsigned>();
	INPUT_DIM = vm["input_dim"].as<unsigned>();
	HIDDEN_DIM = vm["hidden_dim"].as<unsigned>();

	if (vm.count("help")
		|| vm.count("train") != 1
		|| (vm.count("devel") != 1
			&& !(vm.count("test") == 0))) {
		cout << opts << "\n";
			return 1;
	}

	cerr << "PID=" << ::getpid() << endl;

	if (vm.count("lstm"))
		return main_body<LSTMBuilder>(vm);
	else if (vm.count("vlstm"))
		return main_body<VanillaLSTMBuilder>(vm);
	else if (vm.count("dglstm"))
		return main_body<DGLSTMBuilder>(vm);
   	else if (vm.count("gru"))
		return main_body<GRUBuilder>(vm);
   	else
		return main_body<SimpleRNNBuilder>(vm);
}

template <class Builder>
int main_body(variables_map vm)
{
	Corpus traincor, devcor, testcor;
	cerr << "Reading training data from " << vm["train"].as<string>() << "...\n";
	kSOS = d.convert("<s>");
	kEOS = d.convert("</s>");
	kUNK = d.convert("<unk>");
	traincor = Read_Corpus(vm["train"].as<string>(), vm["lang"].as<bool>());
	d.freeze(); // no new word types allowed

	VOCAB_SIZE = d.size();

	if (vm.count("dev")) {
		cerr << "Reading development data from " << vm["dev"].as<string>() << "...\n";
		devcor = Read_Corpus(vm["dev"].as<string>(), vm["lang"].as<bool>());
	}

	if (vm.count("test")){
		cerr << "Reading testing data from " << vm["test"].as<string>() << "...\n";
		testcor = Read_Corpus(vm["test"].as<string>(), vm["lang"].as<bool>());
	}

	float dropout_p = 0.0f;
	if (vm.count("dropout"))//apply dropout technique
	{
		dropout_p = vm["dropout"].as<float>();
		cerr << "Applying dropout technique with dropout probability p=" << dropout_p << "..." << endl;
	}

	string param_file = "";
	if (vm.count("parameters"))
		param_file = vm["parameters"].as<string>();
	else if (vm.count("initialise"))
		param_file = vm["initialise"].as<string>();
	else {
		ostringstream os;
		os << "rnnlm" << '_' << LAYERS << '_' << HIDDEN_DIM << '_' << INPUT_DIM << '_' << VOCAB_SIZE  << "_pid" << getpid() << ".params";
		param_file = os.str();
	}
	cerr << "Parameters will be updated to: " << param_file << endl;

	BATCH_SIZE = vm["batch_size"].as<unsigned>();

	Model model;
	unsigned sgd_trainer_type = vm["sgd_trainer"].as<unsigned>();
	Trainer* sgd_trainer = nullptr;
	if (sgd_trainer_type == 1)
		sgd_trainer = new MomentumSGDTrainer(model, vm["eta"].as<float>());
	else if (sgd_trainer_type == 2)
		sgd_trainer = new AdagradTrainer(model, vm["eta"].as<float>());
	else if (sgd_trainer_type == 3)
		sgd_trainer = new AdadeltaTrainer(model);
	else if (sgd_trainer_type == 4)
		sgd_trainer = new AdamTrainer(model, vm["eta"].as<float>());
	else if (sgd_trainer_type == 0)//Vanilla SGD trainer
		sgd_trainer = new SimpleSGDTrainer(model, vm["eta"].as<float>());
	else
		assert("Unknown SGD trainer type! (0: vanilla SGD; 1: momentum SGD; 2: Adagrad; 3: AdaDelta; 4: Adam)");
	sgd_trainer->eta_decay = ETA_DECAY;
	sgd_trainer->sparse_updates_enabled = vm["sparse_updates"].as<bool>();
	if (!sgd_trainer->sparse_updates_enabled)
		cerr << "Sparse updates for lookup parameter(s) to be disabled!" << endl;
	sgd_trainer->clip_threshold *= BATCH_SIZE;

	if (vm.count("r2l_target")) cerr << "Using right-to-left direction..." << endl;

	RNNLanguageModel<Builder> lm;
	if (vm.count("initialise"))// load from pre-trained model
		lm.LoadModel(&model, vm["initialise"].as<string>());
	else// create new model
		lm.CreateModel(&model, LAYERS, INPUT_DIM, HIDDEN_DIM, VOCAB_SIZE, vm.count("r2l_target"));

	if (vm.count("test")){
		TestModel(model, lm, testcor);
	}// otherwise, do training.
	else
		TrainModel_Batch2(model, lm, traincor, devcor, sgd_trainer, param_file, dropout_p);

	// cleaning up
	delete sgd_trainer;

	return EXIT_SUCCESS;
}

//====================================== FUNCTION IMPLEMENTATION ======================================

template <class RNNLM_t>
void TestModel(Model &model, RNNLM_t &rnn
	, const Corpus &testcor)
{
	rnn.DisableDropout();

	double dloss = 0;
	unsigned dtokens = 0;
	unsigned id = 0;
	unsigned unk_dtokens = 0;
	for (auto& sent: testcor)
	{
		ComputationGraph cg;
		unsigned tokens = 0, unk_tokens = 0;
		Expression i_xent = rnn.BuildLMGraph(sent, tokens, unk_tokens, cg);

		//more score details for each sentence
		double loss = as_scalar(cg.forward(i_xent));
		//int tokens = sent.size() - 1;
		cerr << id++ << "\t" << loss << "\t" << tokens << "\t" << exp(loss / tokens) << endl;

		dloss += loss;
		dtokens += tokens;
		unk_dtokens += unk_tokens;
	}

	cerr << "-------------------------------------------------------------------------" << endl;
	cerr << "***TEST " << "sentences=" << testcor.size() << " unks=" << unk_dtokens << " E=" << (dloss / dtokens) << " ppl=" << exp(dloss / dtokens) << ' ';
	cerr << "\n-------------------------------------------------------------------------\n" << endl;
}

template <class RNNLM_t>
void TrainModel(Model &model, RNNLM_t &rnn
	, Corpus &traincor, Corpus &devcor
	, Trainer* sgd_trainer
	, const string& param_file, float dropout_p)
{
	unsigned report_every_i = TREPORT;
	unsigned devcor_every_i_reports = DREPORT;
	double best = 9e+99;

	vector<unsigned> order((traincor.size() + BATCH_SIZE - 1) / BATCH_SIZE);
	for (unsigned i = 0; i < order.size(); ++i) order[i] = i * BATCH_SIZE;

	cerr << "**SHUFFLE\n";
	shuffle(order.begin(), order.end(), *rndeng);

	unsigned si = 0;//order.size();
	Timer timer_epoch("completed in"), timer_iteration("completed in");

	int report = 0;
	unsigned lines = 0;
	while (sgd_trainer->epoch < MAX_EPOCH) {
		rnn.EnableDropout(dropout_p);

		double loss = 0;
		unsigned tokens = 0, unk_tokens = 0;
		for (unsigned i = 0; i < report_every_i; ++i, ++si) {
			if (si == order.size()) {
				//timing
				cerr << "***Epoch " << sgd_trainer->epoch << " is finished. ";
				timer_epoch.show();

				si = 0;

				if (ETA_EPOCH == 0)
					sgd_trainer->update_epoch();
				else sgd_trainer->update_epoch(1, ETA_EPOCH); // @vhoang2: learning rate annealing (after lr_epochs, for every next epoch, the learning rate will be decreased by a factor of eta_decay.

				if (sgd_trainer->epoch >= MAX_EPOCH) break;

				cerr << "**SHUFFLE\n";
				shuffle(order.begin(), order.end(), *rndeng);

				timer_epoch.reset();
			}

			// build graph for this instance
			ComputationGraph cg;
			unsigned c1 = 0, c2 = 0;
			Expression i_xent = rnn.BuildLMGraph(traincor[order[si]], c1/*tokens*/, c2/*unk_tokens*/, cg);

			float closs = as_scalar(cg.forward(i_xent));// consume the loss
			if (!is_valid(closs)){
				std::cerr << "***Warning***: nan or -nan values occurred!" << std::endl;
				continue;
			}

			loss += closs;
			tokens += c1;
			unk_tokens += c2;

			cg.backward(i_xent);
			sgd_trainer->update();

			lines++;
		}

		if (sgd_trainer->epoch >= MAX_EPOCH) continue;

		sgd_trainer->status();
		cerr << "sents=" << lines << " unks=" << unk_tokens << " E=" << (loss / tokens) << " ppl=" << exp(loss / tokens) << ' ';
		double elapsed = timer_iteration.elapsed();
		cerr << "[time_elapsed=" << elapsed << "(msec)" << " (" << tokens * 1000.f / elapsed << " words/sec)]" << endl;
		timer_iteration.reset();

		//rnn.RandomSample(); // uncomment this to see some randomly-generated sentences

		// show score on devcor data?
		//report++;
		report += report_every_i;
		if (report % devcor_every_i_reports == 0) {
			rnn.DisableDropout();

			double dloss = 0;
			unsigned dtokens = 0, unk_dtokens = 0;
			for (auto& sent: devcor){
				ComputationGraph cg;
				Expression i_xent = rnn.BuildLMGraph(sent, dtokens, unk_dtokens, cg);
				dloss += as_scalar(cg.forward(i_xent));
			}

			if (dloss < best) {
				best = dloss;
				//ofstream out(param_file);
				//boost::archive::text_oarchive oa(out);
				//oa << lm;
				rnn.SaveModel(param_file);
			}
			//else
			//	sgd_trainer->eta *= 0.5;

			cerr << "\n***DEV [epoch=" << (lines / (double)traincor.size()) << " eta=" << sgd_trainer->eta << "]" << " sents=" << devcor.size() << " unks=" << unk_dtokens << " E=" << (dloss / dtokens) << " ppl=" << exp(dloss / dtokens) << ' ';
			timer_iteration.show();
			timer_iteration.reset();
		}
	}

	cerr << endl << "Training completed!" << endl;
}

//Note: This version works increasingly worse when batch_size > 16.
template <class RNNLM_t>
void TrainModel_Batch1(Model &model, RNNLM_t &rnn
	, Corpus &traincor, Corpus &devcor
	, Trainer* sgd_trainer
	, const string& param_file, float dropout_p)
{
	if (BATCH_SIZE == 1){
		TrainModel(model, rnn, traincor, devcor, sgd_trainer, param_file, dropout_p);
		return;
	}

	// Sort the traincor sentences in descending order of length
	CompareLen comp;
	sort(traincor.begin(), traincor.end(), comp);
	// Pad the sentences in the same batch with EOS so they are the same length
	// This modifies the traincor objective a bit by making it necessary to
	// predict EOS multiple times, but it's easy and not so harmful
	// Potential risk: if there is a very long sentence, computation will be significantly increased!
	for(size_t i = 0; i < traincor.size(); i += BATCH_SIZE)
		for(size_t j = 1; j < BATCH_SIZE; ++j)
			while(traincor[i+j].size() < traincor[i].size())
				traincor[i+j].push_back(kEOS);

	unsigned report_every_i = TREPORT;
	unsigned devcor_every_i_reports = DREPORT;
	double best = 9e+99;

	vector<unsigned> order((traincor.size() + BATCH_SIZE - 1) / BATCH_SIZE);
	for (unsigned i = 0; i < order.size(); ++i) order[i] = i * BATCH_SIZE;

	cerr << "**SHUFFLE\n";
	shuffle(order.begin(), order.end(), *rndeng);

	unsigned si = 0;//order.size();
	Timer timer_epoch("completed in"), timer_iteration("completed in");

	int report = 0;
	unsigned lines = 0;
	while (sgd_trainer->epoch < MAX_EPOCH) {
		rnn.EnableDropout(dropout_p);

		double loss = 0;
		unsigned tokens = 0, unk_tokens = 0;
		for (unsigned i = 0; i < report_every_i; /*++i,*/ ++si) {
			if (si == order.size()) {
				//timing
				cerr << "***Epoch " << sgd_trainer->epoch << " is finished. ";
				timer_epoch.show();

				si = 0;

				if (ETA_EPOCH == 0)
					sgd_trainer->update_epoch();
				else sgd_trainer->update_epoch(1, ETA_EPOCH); // @vhoang2: learning rate annealing (after lr_epochs, for every next epoch, the learning rate will be decreased by a factor of eta_decay.

				if (sgd_trainer->epoch >= MAX_EPOCH) break;

				cerr << "**SHUFFLE\n";
				shuffle(order.begin(), order.end(), *rndeng);

				timer_epoch.reset();
			}

			// build graph for this instance
			ComputationGraph cg;
			unsigned bsize = std::min((unsigned)traincor.size()-order[si], BATCH_SIZE); // batch size
			unsigned c1 = 0, c2 = 0;
			Expression i_xent = rnn.BuildLMGraph(traincor, order[si], bsize, c1/*tokens*/, c2/*unk_tokens*/, cg);

			float closs = as_scalar(cg.forward(i_xent));// consume the loss
			if (!is_valid(closs)){
				std::cerr << "***Warning***: nan or -nan values occurred!" << std::endl;
				continue;
			}

			loss += closs;
			tokens += c1;
			unk_tokens += c2;

			cg.backward(i_xent);
			sgd_trainer->update();

			lines += bsize;
			i += bsize;
		}

		if (sgd_trainer->epoch >= MAX_EPOCH) continue;

		sgd_trainer->status();
		cerr << "sents=" << lines << " unks=" << unk_tokens << " E=" << (loss / tokens) << " ppl=" << exp(loss / tokens) << ' ';
		double elapsed = timer_iteration.elapsed();
		cerr << "[time_elapsed=" << elapsed << "(msec)" << " (" << tokens * 1000.f / elapsed << " words/sec)]" << endl;
		timer_iteration.reset();

		//rnn.RandomSample(); // uncomment this to see some randomly-generated sentences

		// show score on devcor data?
		//report++;
		report += report_every_i;
		if (report % devcor_every_i_reports == 0) {
			rnn.DisableDropout();

			double dloss = 0;
			unsigned dtokens = 0, unk_dtokens = 0;
			for (auto& sent: devcor){
				ComputationGraph cg;
				Expression i_xent = rnn.BuildLMGraph(sent, dtokens, unk_dtokens, cg);
				dloss += as_scalar(cg.forward(i_xent));
			}

			if (dloss < best) {
				best = dloss;
				//ofstream out(param_file);
				//boost::archive::text_oarchive oa(out);
				//oa << lm;
				rnn.SaveModel(param_file);
			}
			//else
			//	sgd_trainer->eta *= 0.5;

			cerr << "\n***DEV [epoch=" << (lines / (double)traincor.size()) << " eta=" << sgd_trainer->eta << "]" << " sents=" << devcor.size() << " unks=" << unk_dtokens << " E=" << (dloss / dtokens) << " ppl=" << exp(dloss / dtokens) << ' ';
			timer_iteration.show();
			timer_iteration.reset();
		}
	}

	cerr << endl << "Training completed!" << endl;
}

// --------------------------------------------------------------------------------------------------------------------------------
// The following codes are referenced from lamtram toolkit (https://github.com/neubig/lamtram).
struct SingleLength
{
	SingleLength(const vector<Sentence> & v) : vec(v) { }
	inline bool operator() (int i1, int i2)
	{
		return (vec[i2].size() < vec[i1].size());
	}
	const vector<Sentence> & vec;
};

inline size_t CalcSize(const Sentence & src, int trg) {
	return src.size()+1;
}

inline void Create_MiniBatches(const Corpus& traincor,
	size_t max_size,
	std::vector<std::vector<Sentence> > & traincor_minibatch)
{
	std::vector<int> train_ids(traincor.size());
	std::iota(train_ids.begin(), train_ids.end(), 0);

	if(max_size > 1)
		sort(train_ids.begin(), train_ids.end(), SingleLength(traincor));

	std::vector<Sentence> traincor_next;
	size_t first_size = 0;
	for(size_t i = 0; i < train_ids.size(); i++) {
		if (traincor_next.size() == 0)
			first_size = traincor[train_ids[i]].size();

		traincor_next.push_back(traincor[train_ids[i]]);

		if ((traincor_next.size()+1) * first_size > max_size) {
			traincor_minibatch.push_back(traincor_next);
			traincor_next.clear();
		}
	}

	if (traincor_next.size()) traincor_minibatch.push_back(traincor_next);
}
// --------------------------------------------------------------------------------------------------------------------------------

template <class RNNLM_t>
void TrainModel_Batch2(Model &model, RNNLM_t &rnn
	, Corpus &traincor, Corpus &devcor
	, Trainer* sgd_trainer
	, const string& param_file, float dropout_p)
{
	if (BATCH_SIZE == 1){
		TrainModel(model, rnn, traincor, devcor, sgd_trainer, param_file, dropout_p);
		return;
	}

	// create minibatches
	std::vector<std::vector<Sentence> > train_minibatch;
	size_t minibatch_size = BATCH_SIZE;
	Create_MiniBatches(traincor, minibatch_size, train_minibatch);

	std::vector<int> train_ids_minibatch(train_minibatch.size());
	std::iota(train_ids_minibatch.begin(), train_ids_minibatch.end(), 0);
	cerr << "**SHUFFLE\n";
	std::shuffle(train_ids_minibatch.begin(), train_ids_minibatch.end(), *dynet::rndeng);

	unsigned report_every_i = TREPORT;
	unsigned dev_every_i_reports = DREPORT;
	double best = 9e+99;
	unsigned si = 0, last_print = 0, lines = 0;
	Timer timer_epoch("completed in"), timer_iteration("completed in");
	while (sgd_trainer->epoch < MAX_EPOCH) {
		rnn.EnableDropout(dropout_p);

		double loss = 0;
		unsigned tokens = 0, unk_tokens = 0;
		for (unsigned i = 0; i < dev_every_i_reports; ++si) {
			if (si == train_ids_minibatch.size()) {
				//timing
				cerr << "***Epoch " << sgd_trainer->epoch << " is finished. ";
				timer_epoch.show();

				si = 0;
				last_print = 0;
				lines = 0;

				if (ETA_EPOCH == 0)
					sgd_trainer->update_epoch();
				else sgd_trainer->update_epoch(1, ETA_EPOCH); // @vhoang2: learning rate annealing (after lr_epochs, for every next epoch, the learning rate will be decreased by a factor of eta_decay.

				if (sgd_trainer->epoch >= MAX_EPOCH) break;

				cerr << "**SHUFFLE\n";
				shuffle(train_ids_minibatch.begin(), train_ids_minibatch.end(), *dynet::rndeng);

				timer_epoch.reset();
			}

			// build graph for this instance
			ComputationGraph cg;
			unsigned c1 = 0, c2 = 0;// intermediate variables
			Expression i_xent = rnn.BuildLMGraph(train_minibatch[train_ids_minibatch[si]], c1/*tokens*/, c2/*unk_tokens*/, cg);

			cg.forward(i_xent);// forward step
			float closs = as_scalar(cg.get_value(i_xent.i));// consume the loss
			if (!is_valid(closs)){
				std::cerr << "***Warning***: nan or -nan values occurred!" << std::endl;
				continue;
			}

			// update returned values
			loss += closs;
			tokens += c1;
			unk_tokens += c2;

			cg.backward(i_xent);// backward step
			sgd_trainer->update();// SGD update step

			lines += train_minibatch[train_ids_minibatch[si]].size();
			i += train_minibatch[train_ids_minibatch[si]].size();

			if (lines / report_every_i != last_print
					|| i >= dev_every_i_reports /*|| sgd.epoch >= max_epochs*/
					|| si + 1 == train_ids_minibatch.size()){
				last_print = lines / report_every_i;

				sgd_trainer->status();
				cerr << "sents=" << lines << " unks=" << unk_tokens << " E=" << (loss / tokens) << " ppl=" << exp(loss / tokens) << ' ';
				double elapsed = timer_iteration.elapsed();
				cerr << "[time_elapsed=" << elapsed << "(msec)" << " (" << tokens * 1000.f / elapsed << " words/sec)]" << endl;
			}
		}

		timer_iteration.reset();

		if (sgd_trainer->epoch >= MAX_EPOCH) continue;

		//rnn.RandomSample(); // uncomment this to see some randomly-generated sentences

		// show score on devcor data?
		rnn.DisableDropout();

		double dloss = 0;
		unsigned dtokens = 0, unk_dtokens = 0;
		for (auto& sent: devcor){
			ComputationGraph cg;
			Expression i_xent = rnn.BuildLMGraph(sent, dtokens, unk_dtokens, cg);
			dloss += as_scalar(cg.forward(i_xent));
		}

		if (dloss < best) {
			best = dloss;
			//ofstream out(param_file);
			//boost::archive::text_oarchive oa(out);
			//oa << lm;
			rnn.SaveModel(param_file);
		}
		//else
		//	sgd_trainer->eta *= 0.5;

		cerr << "--------------------------------------------------------------------------------------------------------" << endl;
		cerr << "***DEV [epoch=" << sgd_trainer->epoch + lines / (double)traincor.size() << " eta=" << sgd_trainer->eta << "]" << " sents=" << devcor.size() << " unks=" << unk_dtokens << " E=" << (dloss / dtokens) << " ppl=" << exp(dloss / dtokens) << ' ';
		timer_iteration.show();
		cerr << "--------------------------------------------------------------------------------------------------------" << endl;
		timer_iteration.reset();
	}

	cerr << endl << "Training completed!" << endl;
}

Corpus Read_Corpus(const string &fp, bool dir)
{
	Corpus cor;

	std::ifstream f(fp);

	string line;
	int lc = 0, toks = 0;
	while (std::getline(f, line))
	{
		// parse the line which can be: source ||| target
		size_t pos_sep = line.find("|||");
		if (pos_sep != std::string::npos){
			if (dir)// source
				line = line.substr(0, pos_sep);
			else// target
				line = line.substr(pos_sep + 3);
		}

		boost::trim(line);

		Sentence source;
		source = read_sentence(line, d);
		cor.push_back(Sentence(source));

		toks += source.size();

		if ((source.front() != kSOS && source.back() != kEOS))
		{
			cerr << "Sentence in " << fp << ":" << lc << " didn't start or end with <s> and </s>\n";
			abort();
		}

		lc++;
	}

	cerr << lc << " lines, " << toks << " tokens, " << d.size() << " types\n";

	return cor;
}