// Copyright (c) 2013 Vasili Baranau
// Distributed under the MIT software license
// See the accompanying file License.txt or http://opensource.org/licenses/MIT

#include "GenerationManager.h"

#include <algorithm>
#include <boost/shared_ptr.hpp>
#include "Core/Headers/Path.h"
#include "Core/Headers/StlUtilities.h"
#include "Core/Headers/Math.h"
#include "Model/Headers/Config.h"

#include "PackingServices/Headers/MathService.h"
#include "PackingServices/Headers/PackingSerializer.h"
#include "PackingServices/Headers/GeometryService.h"
#include "PackingServices/Headers/ImmobileParticlesService.h"

#include "PackingServices/PostProcessing/Headers/InsertionRadiiGenerator.h"
#include "PackingServices/PostProcessing/Headers/HessianService.h"
#include "PackingServices/PostProcessing/Headers/PressureService.h"
#include "PackingServices/PostProcessing/Headers/RattlerRemovalService.h"
#include "PackingServices/PostProcessing/Headers/MolecularDynamicsService.h"

#include "PackingServices/EnergyServices/Headers/HarmonicPotential.h"
#include "PackingServices/EnergyServices/Headers/IEnergyService.h"

#include "PackingServices/DistanceServices/Headers/DistanceService.h"

#include "Geometries/Headers/BulkGeometry.h"
#include "Geometries/Headers/CircleGeometry.h"
#include "Geometries/Headers/RectangleGeometry.h"
#include "Geometries/Headers/TrapezoidGeometry.h"

#include "PackingGenerators/InitialGenerators/Headers/BulkPoissonGenerator.h"
#include "PackingGenerators/InitialGenerators/Headers/BulkPoissonInCellsGenerator.h"
#include "PackingGenerators/InitialGenerators/Headers/HcpGenerator.h"

#include "PackingGenerators/Headers/IPackingGenerator.h"

using namespace std;
using namespace Geometries;
using namespace PackingServices;
using namespace PackingGenerators;
using namespace Model;
using namespace Core;

namespace Generation
{
    GenerationManager::GenerationManager(PackingSerializer* packingSerializer,
            IPackingGenerator* packingGenerator,
            InsertionRadiiGenerator* insertionRadiiGenerator,
            DistanceService* distanceService,
            OrderService* orderService,
            IEnergyService* contractionEnergyService,
            HessianService* hessianService,
            PressureService* pressureService,
            MolecularDynamicsService* molecularDynamicsService,
            RattlerRemovalService* rattlerRemovalService,
            ImmobileParticlesService* immobileParticlesService)
    {
        this->packingSerializer = packingSerializer;
        this->packingGenerator = packingGenerator;
        this->insertionRadiiGenerator = insertionRadiiGenerator;
        this->distanceService = distanceService;
        this->orderService = orderService;
        this->contractionEnergyService = contractionEnergyService;
        this->hessianService = hessianService;
        this->pressureService = pressureService;
        this->molecularDynamicsService = molecularDynamicsService;
        this->rattlerRemovalService = rattlerRemovalService;
        this->immobileParticlesService = immobileParticlesService;

        innerDiameterRatio = 1.0;
    }

    void GenerationManager::GeneratePacking(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, PACKING_FILE_NAME, false, false, &GenerationManager::GeneratePacking);
    }

    void GenerationManager::GenerateInsertionRadii(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, INSERTION_RADII_FILE_NAME, true, true, &GenerationManager::GenerateInsertionRadii);
    }

    void GenerationManager::CalculateDistancesToClosestSurfaces(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, DISTANCES_TO_CLOSEST_SURFACES_FOLDER_NAME, true, true, &GenerationManager::CalculateDistancesToClosestSurfaces);
    }

    void GenerationManager::CalculateContactNumberDistribution(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, CONTACT_NUMBER_DISTRIBUTION_FILE_NAME, true, true, &GenerationManager::CalculateContactNumberDistribution);
    }

    void GenerationManager::CalculateEntropy(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, ENTROPY_FILE_NAME, true, true, &GenerationManager::CalculateEntropy);
    }

    void GenerationManager::CalculateDirections(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, PARTICLE_DIRECTIONS_FILE_NAME, true, true, &GenerationManager::CalculateDirections);
    }

    void GenerationManager::CalculateContractionEnergies(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, CONTRACTION_ENERGIES_FILE_NAME, true, true, &GenerationManager::CalculateContractionEnergies);
    }

    void GenerationManager::GenerateOrder(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, ORDER_FILE_NAME, true, true, &GenerationManager::GenerateOrder);
    }

    void GenerationManager::CalculateHessianEigenvalues(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, HESSIAN_EIGENVALUES_FILE_NAME, true, true, &GenerationManager::CalculateHessianEigenvalues);
    }

    void GenerationManager::CalculatePressures(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, PRESSURES_FILE_NAME, true, true, &GenerationManager::CalculatePressures);
    }

    void GenerationManager::CalculateMolecularDynamicsStatistics(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, MOLECULAR_DYNAMICS_STATISTICS_FILE_NAME, true, true, &GenerationManager::CalculateMolecularDynamicsStatistics);
    }

    void GenerationManager::RemoveRattlers(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, "", false, true, &GenerationManager::RemoveRattlers);
    }

    void GenerationManager::CalculatePairCorrelationFunction(const Model::ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, PAIR_CORRELATION_FUNCTION_FILE_NAME, true, true, &GenerationManager::CalculatePairCorrelationFunction);
    }

    void GenerationManager::CalculateStructureFactor(const Model::ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, STRUCTURE_FACTOR_FILE_NAME, true, true, &GenerationManager::CalculateStructureFactor);
    }

    void GenerationManager::GenerateLocalOrientationalDisorder(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, LOCAL_ORIENTATIONAL_DISORDER_FILE_NAME, true, true, &GenerationManager::GenerateLocalOrientationalDisorder);
    }

    void GenerationManager::CalculateImmediateMolecularDynamicsStatistics(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, IMMEDIATE_MOLECULAR_DYNAMICS_STATISTICS_FILE_NAME, true, true, &GenerationManager::CalculateImmediateMolecularDynamicsStatistics);
    }

    void GenerationManager::CalculateNearestNeighbors(const ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, NEAREST_NEIGHBORS_FILE_NAME, true, true, &GenerationManager::CalculateNearestNeighbors);
    }

    void GenerationManager::CalculateActiveGeometry(const Model::ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, ACTIVE_GEOMETRY_FILE_NAME, true, true, &GenerationManager::CalculateActiveGeometry);
    }

    void GenerationManager::CalculateSuccessfulPermutationProbability(const Model::ExecutionConfig& userConfig)
    {
        ExecuteAlgorithm(userConfig, SUCCESSFUL_PERMUTATION_PROBABILITY_FILE_NAME, true, true, &GenerationManager::CalculateSuccessfulPermutationProbability);
    }

    void GenerationManager::ExecuteAlgorithm(const ExecutionConfig& userConfig, string targetFileName,
            bool shouldExitIfTargetFileExists, bool shouldAlwaysReadPacking, Action algorithm)
    {
        string targetFilePath = Path::Append(userConfig.generationConfig.baseFolder, targetFileName);
        if (shouldExitIfTargetFileExists && Path::Exists(targetFilePath))
        {
            return;
        }

        ExecutionConfig fullConfig;
        FillFullConfig(userConfig, &fullConfig);

        Math::SetSeed(fullConfig.generationConfig.seed);
        boost::shared_ptr<IGeometry> geometry = CreateGeometry(fullConfig.systemConfig);

        SystemConfig activeConfig;
        boost::shared_ptr<IGeometry> activeGeometry;
        string activeConfigPath = Path::Append(fullConfig.generationConfig.baseFolder, ACTIVE_GEOMETRY_FILE_NAME);
        if (Path::Exists(activeConfigPath))
        {
            SpatialVector shift;
            packingSerializer->ReadActiveConfig(activeConfigPath, &activeConfig, &shift);
            activeConfig.MergeWith(fullConfig.systemConfig);

            activeGeometry = CreateGeometry(activeConfig, shift);
        }

        ModellingContext context(&fullConfig.systemConfig, geometry.get(), activeGeometry.get());

        packingGenerator->SetContext(context);
        insertionRadiiGenerator->SetContext(context);
        distanceService->SetContext(context);
        orderService->SetContext(context);
        contractionEnergyService->SetContext(context);
        hessianService->SetContext(context);
        pressureService->SetContext(context);
        molecularDynamicsService->SetContext(context);
        rattlerRemovalService->SetContext(context);

        Packing particles;
        ReadOrCreatePacking(fullConfig, context, shouldAlwaysReadPacking, &particles);

        (this->*algorithm)(fullConfig, context, targetFilePath, &particles);
    }

    // Actions for corresponding algorithms

    void GenerationManager::GeneratePacking(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        string infoFile = Path::Append(fullConfig.generationConfig.baseFolder, PACKING_FILE_NAME_NFO);
        if (Path::Exists(infoFile))
        {
            return;
        }

//        // A quick fix for Matthias packings
//        Packing& particlesRef = *particles;
//        string nearestNeighbotsPath = Path::Append(fullConfig.generationConfig.baseFolder, NEAREST_NEIGHBORS_FILE_NAME);
//        if (Path::Exists(nearestNeighbotsPath))
//        {
//            vector<ParticlePair> closestPairs;
//            packingSerializer->ReadNearestNeighbors(nearestNeighbotsPath, &closestPairs);
//            for (size_t i = 0; i < closestPairs.size(); ++i)
//            {
//                bool radiusCorrect = closestPairs[i].normalizedDistanceSquare > 0.95;
//                particlesRef[closestPairs[i].firstParticleIndex].isImmobile = radiusCorrect;
//                particlesRef[closestPairs[i].secondParticleIndex].isImmobile = radiusCorrect;
//            }
//        }

        packingGenerator->SetGenerationConfig(fullConfig.generationConfig);
        packingGenerator->ArrangePacking(particles);

        // TODO: Think about writing info file here.
        // TODO: Think on renaming the folder, updating the generation config, rescaling the packing, updating the info file according to the innerDiameterRatio here.
        packingSerializer->SerializePacking(targetFilePath, *particles);

        innerDiameterRatio = packingGenerator->GetFinalInnerDiameterRatio();
        string contractionEnergiesFilePath = Path::Append(fullConfig.generationConfig.baseFolder, CONTRACTION_ENERGIES_FILE_NAME);
        CalculateContractionEnergies(fullConfig, context, contractionEnergiesFilePath, particles);

        innerDiameterRatio = 1.0;
    }

    void GenerationManager::GenerateInsertionRadii(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        vector<FLOAT_TYPE> insertionRadii;
        insertionRadiiGenerator->FillInsertionRadii(*particles, fullConfig.generationConfig.insertionRadiiCount, &insertionRadii);
        packingSerializer->SerializeInsertionRadii(targetFilePath, insertionRadii);
    }

    void GenerationManager::CalculateDistancesToClosestSurfaces(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFolderPath, Packing* particles)
    {
        vector<int> surfaceIndexes;
        vector<vector<FLOAT_TYPE> > distancesToSurfaces;

        int minNeighborsCount = 2; // Now: some number to ensure that we have all the numbers consecutively. Before: contractionEnergyService->GetMinNeighborsCount(); // average neighbors count in mechanically stable packings with infinite friction
        int maxNeighborsCount = 14; // max number of contacts in monodisperse packings
        int minSurfaceIndex = minNeighborsCount - 1;
        int surfaceIndexesCount = maxNeighborsCount - minNeighborsCount + 1;
        surfaceIndexes.resize(surfaceIndexesCount);
        VectorUtilities::FillLinearScale(minSurfaceIndex, &surfaceIndexes);
        surfaceIndexes.insert(surfaceIndexes.begin(), 0);

        insertionRadiiGenerator->FillDistancesToSurfaces(*particles, fullConfig.generationConfig.insertionRadiiCount, surfaceIndexes, targetFolderPath, *packingSerializer);
    }

    void GenerationManager::CalculateContactNumberDistribution(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        printf("Calculating contact number distributions\n");
        Packing& packingToUse = *particles;

        // just make sure that the closest pair touches
        //ClosestPairProvider
        distanceService->SetParticles(*particles);
        ParticlePair closestPair = distanceService->FindClosestPair();
        FLOAT_TYPE closestNormalizedDistance = std::sqrt(closestPair.normalizedDistanceSquare);
        printf("Normalized distance of a closest pair is %f\n", closestNormalizedDistance);
        if (closestNormalizedDistance > 1.0001)
        {
            printf("Min normalized distance is too high, parcking was probably not rescaled to the final density. Rescaling before contacts calculation...\n");
            Packing rescaledParticles = *particles;
            for (Particle& p : rescaledParticles)
            {
                p.diameter *= closestNormalizedDistance;
            }
            packingToUse = rescaledParticles;
        }

        FLOAT_TYPE contractionRate = 1.0 - 1e-4;
        vector<int> neighborCounts;
        vector<int> neighborCountFrequencies;
        vector<vector<int>> touchingParticleIndexes;
        contractionEnergyService->SetParticles(packingToUse);
        FLOAT_TYPE estimatedCoordinationNumber = insertionRadiiGenerator->GetContactNumberDistribution(packingToUse, contractionEnergyService,
                contractionRate, &neighborCounts, &neighborCountFrequencies, &touchingParticleIndexes);

        printf("Estimated coordination number is %f\n", estimatedCoordinationNumber);

        packingSerializer->SerializeContactNumberDistribution(targetFilePath, neighborCounts, neighborCountFrequencies);

        string targetFolder = Core::Path::GetParentPath(targetFilePath);
        string contactNumbersFilePath = Core::Path::Append(targetFolder, CONTACTING_NEIGHBORS_FILE_NAME);
        packingSerializer->SerializeContactingNeighborIndexes(contactNumbersFilePath, touchingParticleIndexes);

        vector<FLOAT_TYPE> normalizedContactingNeighborDistances;
        insertionRadiiGenerator->FillNormalizedContactingNeighborDistances(packingToUse, touchingParticleIndexes, &normalizedContactingNeighborDistances);

        string contactingNeighborDistancesFilePath = Core::Path::Append(targetFolder, CONTACTING_NEIGHBOR_DISTANCES_FILE_NAME);
        packingSerializer->SerializeContactingNeighborDistances(contactingNeighborDistancesFilePath, normalizedContactingNeighborDistances);

//        FLOAT_TYPE expectedCoordinationNumber = GetExpectedCoordinationNumber(fullConfig, context, targetFilePath, particles);
//        // Find the best contraction rate to get this coordination number
//        FLOAT_TYPE contractionRate = insertionRadiiGenerator->GetContractionRateForCoordinationNumber(contractionEnergyService, expectedCoordinationNumber);
//
//        // Find coordination numbers distribution
//        vector<int> neighborCounts;
//        vector<int> neighborCountFrequencies;
//        FLOAT_TYPE estimatedCoordinationNumber = insertionRadiiGenerator->GetContactNumberDistribution(*particles, contractionEnergyService, contractionRate, &neighborCounts, &neighborCountFrequencies);
//        if (std::abs(estimatedCoordinationNumber - expectedCoordinationNumber) > 1e-2)
//        {
//            printf("Coordination number estimated incorrectly\n");
//        }
//
//        packingSerializer->SerializeContactNumberDistribution(targetFilePath, neighborCounts, neighborCountFrequencies);
    }

    FLOAT_TYPE GenerationManager::GetExpectedCoordinationNumber(const Model::ExecutionConfig& fullConfig, const Model::ModellingContext& context, std::string targetFilePath, Model::Packing* particles) const
    {
        string energyFilePath = Path::Append(Path::GetParentPath(targetFilePath), CONTRACTION_ENERGIES_FILE_NAME);
        if (!Path::Exists(energyFilePath))
        {
            throw InvalidOperationException("Contraction energy file does not exist.");
        }
        contractionEnergyService->SetParticles(*particles);

        // Read coordination numbers from file
        vector<FLOAT_TYPE> contractionRatios;
        vector<FLOAT_TYPE> energyPowers;
        vector<FLOAT_TYPE> contractionEnergies;
        vector<int> nonRattlersCounts;
        packingSerializer->ReadContractionEnergies(energyFilePath, &contractionRatios, &energyPowers, &contractionEnergies, &nonRattlersCounts);

        // Find expected coordination number
        FLOAT_TYPE expectedCoordinationNumber = 0;
        for (size_t i = 0; i < contractionRatios.size(); ++i)
        {
            if (energyPowers[i] > 0)
            {
                expectedCoordinationNumber = contractionEnergies[i - 1] / nonRattlersCounts[i - 1];
                break;
            }
        }

        return expectedCoordinationNumber;
    }

    void GenerationManager::CalculateEntropy(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        FLOAT_TYPE entropy = insertionRadiiGenerator->CalculateEntropy(*particles, fullConfig.generationConfig.insertionRadiiCount);
        packingSerializer->SerializeEntropy(targetFilePath, entropy);

//        vector<FLOAT_TYPE> localEntropies;
//        FLOAT_TYPE entropy = insertionRadiiGenerator->CalculateLocalEntropies(particles, fullConfig->generationConfig.insertionRadiiCount, &localEntropies);
//        packingSerializer->SerializeLocalEntropy(targetFilePath, entropy, &localEntropies);
    }

    void GenerationManager::CalculateDirections(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        vector<OrderService::NeighborDirections> particleDirections;
        orderService->SetParticles(*particles);
        orderService->FillParticleDirections(&particleDirections);
        packingSerializer->SerializeParticleDirections(targetFilePath, fullConfig.systemConfig.particlesCount, particleDirections);
    }

    void GenerationManager::CalculateContractionEnergies(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        printf("Calculating contraction energies\n");

        vector<FLOAT_TYPE> contractionRatios;
        FillContractionRatios(&contractionRatios);
        int count = contractionRatios.size();
        vector<FLOAT_TYPE> fullContractionRatios;

        // Copy contraction rates twice
        StlUtilities::Append(contractionRatios, &fullContractionRatios);
        StlUtilities::Append(contractionRatios, &fullContractionRatios);

        const HarmonicPotential zeroPotential(0.0);
        const HarmonicPotential secondPotential(2.0);

        vector<const IPairPotential*> fullPotentials(count, &zeroPotential);
        vector<const IPairPotential*> potentials(count, &secondPotential);
        StlUtilities::Append(potentials, &fullPotentials);

        // Rescale the contraction ratios according to the final innerDiameterRatio (this method may be called after packing generation).
        vector<FLOAT_TYPE> rescaledContractionRatios;
        for (vector<FLOAT_TYPE>::iterator it = fullContractionRatios.begin(); it != fullContractionRatios.end(); ++it)
        {
            FLOAT_TYPE contractionRatio = *it;
            rescaledContractionRatios.push_back(contractionRatio / innerDiameterRatio);
        }

        contractionEnergyService->SetParticles(*particles);
        IEnergyService::EnergiesResult result = contractionEnergyService->GetContractionEnergies(rescaledContractionRatios, fullPotentials);

        // Use non-rescaled contraction ratios.
        vector<FLOAT_TYPE> fullEnergyPowers(count, 0.0);
        vector<FLOAT_TYPE> energyPowers(count, 2.0);
        StlUtilities::Append(energyPowers, &fullEnergyPowers);

        packingSerializer->SerializeContractionEnergies(targetFilePath, fullContractionRatios, fullEnergyPowers, result.contractionEnergies, result.nonRattlersCounts);
    }

    void GenerationManager::GenerateOrder(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        orderService->SetParticles(*particles);
        OrderService::Order order = orderService->GetOrder(6); // 6 is the best order, as is maximum for FCC packings
        packingSerializer->SerializeOrder(targetFilePath, order);
    }

    void GenerationManager::CalculateHessianEigenvalues(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        hessianService->SetParticles(*particles);
        vector<FLOAT_TYPE> hessianEigenvalues;
        hessianService->FillHessianEigenvalues(&hessianEigenvalues);
        packingSerializer->SerializeHessianEigenvalues(targetFilePath, hessianEigenvalues);
    }

    void GenerationManager::CalculatePressures(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        vector<FLOAT_TYPE> contractionRatios;
        FillContractionRatios(&contractionRatios);
        int count = contractionRatios.size();
        vector<FLOAT_TYPE> fullContractionRatios;

        // No sense in calculating pressure for power = 0
        StlUtilities::Append(contractionRatios, &fullContractionRatios);
        vector<FLOAT_TYPE> fullEnergyPowers(count, 2.0);

        vector<FLOAT_TYPE> pressures;
        pressureService->SetParticles(*particles);
        pressureService->FillPressures(fullContractionRatios, fullEnergyPowers, &pressures);
        packingSerializer->SerializePressures(targetFilePath, fullContractionRatios, fullEnergyPowers, pressures);
    }

    void GenerationManager::CalculateMolecularDynamicsStatistics(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        printf("Calculating molecular dynamics statistics\n");

        molecularDynamicsService->SetGenerationConfig(fullConfig.generationConfig);
        molecularDynamicsService->SetParticles(*particles);

        MolecularDynamicsStatistics statistics = molecularDynamicsService->CalculateStationaryStatistics();
        packingSerializer->SerializeMolecularDynamicsStatistics(targetFilePath, statistics);
    }

    void GenerationManager::RemoveRattlers(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        vector<bool> rattlerMask(context.config->particlesCount);
        rattlerRemovalService->SetParticles(*particles);
//        const FLOAT_TYPE contractionRate = 1.0 - 1.0e-7;
//        const FLOAT_TYPE contractionRate = 1.0 - 1.0e-5;
        const FLOAT_TYPE contractionRate = 1.0 / 1.5; //particles that have one radius between surfaces will become in contact
        rattlerRemovalService->FillRattlerMask(contractionRate, &rattlerMask);
        int nonRattlersCount = rattlerRemovalService->FindNonRattlersCount(rattlerMask);
        Packing nonRattlerParticles;
        nonRattlerParticles.resize(nonRattlersCount);
        rattlerRemovalService->FillNonRattlerPacking(rattlerMask, &nonRattlerParticles);

        printf("Removing rattlers. Non-rattlers count: %d\n", nonRattlersCount);

        PackingInfo oldInfo;
        PackingInfo nonRattlerInfo;

        ExecutionConfig nonRattlerConfig;

        string packingInfoFilePath = Path::Append(fullConfig.generationConfig.baseFolder, PACKING_FILE_NAME_NFO);
        packingSerializer->ReadPackingInfo(packingInfoFilePath, &oldInfo);

        rattlerRemovalService->FillNonRattlerConfig(nonRattlersCount, fullConfig, &nonRattlerConfig);
        rattlerRemovalService->FillNonRattlerPackingInfo(nonRattlersCount, nonRattlerParticles, nonRattlerConfig, oldInfo, &nonRattlerInfo);

        // Serialize non-rattler packing
        packingSerializer->SerializePackingInfo(packingInfoFilePath, nonRattlerConfig.systemConfig, nonRattlerInfo);

        string packingFilePath = Path::Append(fullConfig.generationConfig.baseFolder, PACKING_FILE_NAME);
        packingSerializer->SerializePacking(packingFilePath, nonRattlerParticles);

        string configFilePath = Path::Append(fullConfig.generationConfig.baseFolder, CONFIG_FILE_NAME);
        packingSerializer->SerializeConfig(configFilePath, nonRattlerConfig);
    }

    void GenerationManager::CalculatePairCorrelationFunction(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        printf("Calculating pair correlation function\n");

        distanceService->SetParticles(*particles);
        PairCorrelationFunction pairCorrelationFunction;
        distanceService->FillPairCorrelationFunction(&pairCorrelationFunction);

        packingSerializer->SerializePairCorrelationFunction(targetFilePath, pairCorrelationFunction);
    }

    void GenerationManager::CalculateStructureFactor(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        printf("Calculating structure factor\n");

        Packing& particlesRef = *particles;

        std::vector<int> particleIndexesOfInterest;
        if (fullConfig.generationConfig.particlesToKeepForStructureFactor > 0)
        {
            int particlesToChooseCount = fullConfig.generationConfig.particlesToKeepForStructureFactor;
            bool keepSmallest = fullConfig.generationConfig.keepSmallParticlesForStructureFactor.value;

            // choosing that many largest particles
            // NOTE: maybe use nth element
            std::vector<double> diameters;
            for (const DomainParticle& particle : particlesRef)
            {
                diameters.push_back(particle.diameter);
            }
            std::vector<int> permutation;
            StlUtilities::SortPermutation(diameters, &permutation);

            if (keepSmallest)
            {
                permutation.erase(permutation.begin() + particlesToChooseCount, permutation.end());
            }
            else
            {
                int particlesToRemove = permutation.size() - particlesToChooseCount;
                permutation.erase(permutation.begin(), permutation.begin() + particlesToRemove);
            }

            particleIndexesOfInterest.swap(permutation);

            double minSelectedDiameter = 1000; // TODO: use numeric_limits
            double maxSelectedDiameter = 0;
            for (int i : particleIndexesOfInterest)
            {
                double diameter = particlesRef[i].diameter;
                if (diameter > maxSelectedDiameter)
                {
                    maxSelectedDiameter = diameter;
                }
                if (diameter < minSelectedDiameter)
                {
                    minSelectedDiameter = diameter;
                }
            }

            if (keepSmallest)
            {
                printf("Expected to calculate structure factor for %d smallest particles. Selected %d smallest particles. Max selected diameter: %f\n",
                    particlesToChooseCount,
                    particleIndexesOfInterest.size(),
                    maxSelectedDiameter);
            }
            else
            {
                printf("Expected to calculate structure factor for %d largest particles. Selected %d largest particles. Min selected diameter: %f\n",
                    particlesToChooseCount,
                    particleIndexesOfInterest.size(),
                    minSelectedDiameter);
            }
        }

        distanceService->SetParticles(*particles);
        StructureFactor structureFactor;
        distanceService->FillStructureFactor(particleIndexesOfInterest, &structureFactor);

        packingSerializer->SerializeStructureFactor(targetFilePath, structureFactor);
    }

    void GenerationManager::GenerateLocalOrientationalDisorder(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        printf("Calculating local orientational disorder\n");

        orderService->SetParticles(*particles);
        OrderService::LocalOrientationalDisorder localOrientationalDisorder;
        orderService->FillLocalOrientationalDisorder(&localOrientationalDisorder);
        packingSerializer->SerializeLocalOrientationalDisorder(targetFilePath, localOrientationalDisorder);
        packingSerializer->SerializeCloseNeighbors(Path::Append(fullConfig.generationConfig.baseFolder, "close_neighbors.txt"), localOrientationalDisorder);
    }

    void GenerationManager::CalculateImmediateMolecularDynamicsStatistics(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        printf("Calculating immediate molecular dynamics statistics\n");

        molecularDynamicsService->SetGenerationConfig(fullConfig.generationConfig);
        molecularDynamicsService->SetParticles(*particles);

        MolecularDynamicsStatistics statistics = molecularDynamicsService->CalculateImmediateStatistics();
        packingSerializer->SerializeMolecularDynamicsStatistics(targetFilePath, statistics);
    }

    void GenerationManager::CalculateNearestNeighbors(const Model::ExecutionConfig& fullConfig, const Model::ModellingContext& context, std::string targetFilePath, Model::Packing* particles)
    {
        printf("Calculating nearest neighbors\n");

        distanceService->SetParticles(*particles);
        vector<ParticlePair> closestPairs;
        distanceService->FillClosestPairs(&closestPairs);

        vector<bool> isImmobileMask(fullConfig.systemConfig.particlesCount);
        vector<int> isImmobileIntMask(fullConfig.systemConfig.particlesCount);

        VectorUtilities::InitializeWith(&isImmobileMask, false);
        VectorUtilities::InitializeWith(&isImmobileIntMask, 0);

        Model::Packing& particlesRef = *particles;
        for (ParticleIndex i = 0; i < fullConfig.systemConfig.particlesCount; ++i)
        {
            isImmobileMask[i] = particlesRef[i].isImmobile;
            isImmobileIntMask[i] = particlesRef[i].isImmobile ? 1 : 0;
        }

        vector<int> permutation;
        StlUtilities::SortPermutation(isImmobileIntMask, &permutation);

        vector<ParticlePair> sortedClosestPairs;
        StlUtilities::Permute(closestPairs, permutation, &sortedClosestPairs);

        vector<bool> sortedIsImmobileMask;
        StlUtilities::Permute(isImmobileMask, permutation, &sortedIsImmobileMask);

        packingSerializer->SerializeNearestNeighbors(targetFilePath, closestPairs, sortedIsImmobileMask);
    }

    void GenerationManager::CreateActiveConfig(const ExecutionConfig& fullConfig, FLOAT_TYPE contractionFactorByParticleCenters, SystemConfig* activeConfig, SpatialVector* shift) const
    {
        SystemConfig& activeConfigRef = *activeConfig;
        activeConfigRef.Reset();
        activeConfigRef.MergeWith(fullConfig.systemConfig);

        FLOAT_TYPE meanParticleRadius = 0.5; // TODO: determine better
        // contractionFactor is measured by the true radius of the bounding cylinder (i.e., by particle surfaces)
        FLOAT_TYPE contractionFactor = ((activeConfigRef.packingSize[Axis::X] - 2.0 * meanParticleRadius) * contractionFactorByParticleCenters + 2.0 * meanParticleRadius) / activeConfigRef.packingSize[Axis::X];
        SpatialVector contractionFactors;
        VectorUtilities::InitializeWith(&contractionFactors, contractionFactor);

//            contractionFactors[DIMENSIONS - 1] = 1.0; // For packings that are periodic by z. TODO: periodicity shall be specified in generation.conf. IGeometry shall include this information (as done in LatticeGeoetry)
        contractionFactors[DIMENSIONS - 1] = (activeConfigRef.packingSize[DIMENSIONS - 1] - 8.0 * meanParticleRadius) / activeConfigRef.packingSize[DIMENSIONS - 1]; // Remove 2 diameters from each side
        VectorUtilities::Multiply(activeConfigRef.packingSize, contractionFactors, &activeConfigRef.packingSize);

        SpatialVector margins;
        VectorUtilities::Subtract(fullConfig.systemConfig.packingSize, activeConfigRef.packingSize, &margins);
        VectorUtilities::MultiplyByValue(margins, 0.5, shift);
    }

    void GenerationManager::GetActiveConfigWithParticlesCount(const ExecutionConfig& fullConfig, FLOAT_TYPE contractionFactorByParticleCenters, Packing* particles, SystemConfig* activeConfig, SpatialVector* shift) const
    {
        SystemConfig& activeConfigRef = *activeConfig;
        // the default active area is 0.5 of the cylinder, if cylinder radius is measured by particle centers
        CreateActiveConfig(fullConfig, contractionFactorByParticleCenters, activeConfig, shift);
        boost::shared_ptr<IGeometry> activeGeometry = CreateGeometry(activeConfigRef, *shift);

        ////////////////////
        // Set active geometry

        // NOTE: here I use a dirty hack and do not call SetContext for all the services, as I know, that non of the services caches the active geometry.
        // TODO: extract setting context from ExecuteAlgorithm to a separate method, call here.

        // Set correct isImmobile mask
        SetActiveParticlesByActiveGeometry(fullConfig.systemConfig, *activeGeometry.get(), particles);

        ////////////////////
        // Determine the amount of active non-rattler particles

        // Find the best contraction rate to get the expected coordination number
//        FLOAT_TYPE expectedCoordinationNumber = GetExpectedCoordinationNumber(fullConfig, context, targetFilePath, particles);
//        FLOAT_TYPE contractionRate = insertionRadiiGenerator->GetContractionRateForCoordinationNumber(contractionEnergyService, expectedCoordinationNumber);
        FLOAT_TYPE contractionRate = 1.0 - 1e-7;

        vector<bool> rattlerMask;
        rattlerRemovalService->SetParticles(*particles);
        rattlerRemovalService->FillRattlerMask(contractionRate, &rattlerMask);

        int activeNonRattlersCount = 0;
        Packing& particlesRef = *particles;
        for (ParticleIndex i = 0; i < fullConfig.systemConfig.particlesCount; ++i)
        {
//            if (!particlesRef[i].isImmobile && !rattlerMask[i])
            if (!particlesRef[i].isImmobile)
            {
                activeNonRattlersCount++;
            }
        }

        activeConfigRef.particlesCount = activeNonRattlersCount;
    }

    void GenerationManager::FindActiveConfigForActiveParticlesCount(const ExecutionConfig& fullConfig, int activeParticlesCount, Packing* particles, SystemConfig* activeConfig, SpatialVector* shift) const
    {
        GetActiveParticlesCountFunctor getActiveParticlesCountFunctor(*this, fullConfig, particles, activeConfig, shift);
        StlUtilities::DoBinarySearch(0.0, 1.0, activeParticlesCount, 0, 1e-7, getActiveParticlesCountFunctor);
    }

    void GenerationManager::CalculateActiveGeometry(const ExecutionConfig& fullConfig, const ModellingContext& context, string targetFilePath, Packing* particles)
    {
        printf("Calculating active geometry\n");

        SpatialVector shift;
        SystemConfig activeConfig;

        // NOTE: this is a dirty hack. I'm using insertionRadiiCount parameter to store expectedActiveParticlesCount
        // TODO: add a separate config parameter
        int expectedActiveParticlesCount = fullConfig.generationConfig.insertionRadiiCount;

        // The usual workflow is the following:
        // 1. run CalculateActiveGeometry without an expectedActiveParticlesCount to determine approximate amount of particles inside the active geometry.
        // 2. select the min amount of activeParticlesCount
        // 3. rename the old active configs
        // 4. run CalculateActiveGeometry with expectedActiveParticlesCount = this min number
        // It ensures that all the active geometries have equal amount of active particles.
        // It will make the entropy and compactivity calculations easier
        // (actually, i'm not sure how to calculate compactivity and how to account for rattlers if the amount of particles varies).
        // TODO: extract this entire procedure to a separate method (it's hard currently,
        // as requires processing of several packings at once, which is completely unsupported by the architecture)
        if (expectedActiveParticlesCount == -1)
        {
            // the default active area is 0.5 of the cylinder, if cylinder radius is measured by particle centers
            FLOAT_TYPE contractionFactorByParticleCenters = 0.5;
            GetActiveConfigWithParticlesCount(fullConfig, contractionFactorByParticleCenters, particles, &activeConfig, &shift);
        }
        else
        {
            FindActiveConfigForActiveParticlesCount(fullConfig, expectedActiveParticlesCount, particles, &activeConfig, &shift);
        }

        packingSerializer->SerializeActiveConfig(targetFilePath, activeConfig, shift);
    }

    void GenerationManager::CalculateSuccessfulPermutationProbability(const Model::ExecutionConfig& fullConfig, const Model::ModellingContext& context, std::string targetFilePath, Model::Packing* particles)
    {
        printf("Calculating successful permutation probability\n");
        Packing& packingToUse = *particles;

        int maxPermutations = 5000;
        FLOAT_TYPE successfulPermutationProbability = insertionRadiiGenerator->GetSuccessfulPermutationProbability(particles, maxPermutations);

        printf("Estimated successful permutation probability is %f\n", successfulPermutationProbability);

        //packingSerializer->SerializeContactNumberDistribution(targetFilePath, neighborCounts, neighborCountFrequencies);
    }

    void GenerationManager::FillContractionRatios(vector<FLOAT_TYPE>* contractionRatios) const
    {
        contractionRatios->clear();
        AddLogContractionRatios(contractionRatios);
//        AddLinearContractionRatios(contractionRatios);
        contractionRatios->push_back(0.9999999); // predefined value

        StlUtilities::SortAndResizeToUnique(contractionRatios);
    }

    void GenerationManager::AddLogContractionRatios(vector<FLOAT_TYPE>* contractionRatios) const
    {
        const int pointsCount = 15;
        const FLOAT_TYPE minContractionRatio = 0.9999999;
        const FLOAT_TYPE maxContractionRatio = 0.99995;

        vector<FLOAT_TYPE> contractionRatiosArray(pointsCount);

        VectorUtilities::FillLogScale(1. - minContractionRatio, 1. - maxContractionRatio, &contractionRatiosArray);
        VectorUtilities::MultiplyByValue(contractionRatiosArray, -1.0, &contractionRatiosArray);
        VectorUtilities::AddValue(contractionRatiosArray, 1.0, &contractionRatiosArray);

        StlUtilities::Append(contractionRatiosArray, contractionRatios);
    }

    void GenerationManager::AddLinearContractionRatios(vector<FLOAT_TYPE>* contractionRatios) const
    {
        // 200 points per 0.01
        // maxParticleDiameter / maxContractionRatio * innerDiameterRatio <= maxParticleDiameter * VerletListNeighborProvider::MAX_EXPECTED_OUTER_DIAMETER_RATIO + cutoffDistance
        // Approximately 1.0 / maxContractionRatio * innerDiameterRatio <= VerletListNeighborProvider::MAX_EXPECTED_OUTER_DIAMETER_RATIO
        const int pointsCount = 401;
        const FLOAT_TYPE minContractionRatio = 1.0;
        const FLOAT_TYPE maxContractionRatio = 0.98;

        vector<FLOAT_TYPE> contractionRatiosArray(pointsCount);

        VectorUtilities::FillLinearScale(minContractionRatio, maxContractionRatio, &contractionRatiosArray);
        StlUtilities::Append(contractionRatiosArray, contractionRatios);
    }

    void GenerationManager::FillFullConfig(const ExecutionConfig& userConfig, ExecutionConfig* fullConfig) const
    {
        ExecutionConfig fileConfig;
        packingSerializer->ReadConfig(userConfig.generationConfig.baseFolder, &fileConfig);

        fullConfig->MergeWith(userConfig);
        fullConfig->MergeWith(fileConfig);
    }

    boost::shared_ptr<IGeometry> GenerationManager::CreateGeometry(const SystemConfig& config, const Core::SpatialVector& shift) const
    {
        // TODO: use shift in all non-bulk geometries and may be add a template parameter to all of them <TSupportsShift>.
        boost::shared_ptr<IGeometry> geometry;
        switch (config.boundariesMode)
        {
        case BoundariesMode::Bulk:
            geometry.reset(new BulkGeometry(config));
            break;
        case BoundariesMode::Ellipse:
            geometry.reset(new CircleGeometry(config, shift));
            break;
        case BoundariesMode::Rectangle:
            geometry.reset(new RectangleGeometry(config));
            break;
        case BoundariesMode::Trapezoid:
            geometry.reset(new TrapezoidGeometry(config));
            break;
        case BoundariesMode::Unknown:
            break;
        }

        return geometry;
    }

    boost::shared_ptr<IGeometry> GenerationManager::CreateGeometry(const SystemConfig& config) const
    {
        SpatialVector shift;
        VectorUtilities::InitializeWith(&shift, 0.0);
        return CreateGeometry(config, shift);
    }

    void GenerationManager::ReadOrCreatePacking(const ExecutionConfig& fullConfig, const ModellingContext& context, bool shouldAlwaysReadPacking, Packing* particles) const
    {
        Packing& particlesRef = *particles;
        particlesRef.resize(fullConfig.systemConfig.particlesCount);

        for (ParticleIndex i = 0; i < fullConfig.systemConfig.particlesCount; ++i)
        {
            particlesRef[i].index = i;
        }

        if (shouldAlwaysReadPacking)
        {
            // Always read the packing, even if the generation mode is "start"
            packingSerializer->ReadPacking(Path::Append(fullConfig.generationConfig.baseFolder, PACKING_FILE_NAME), particles);
        }
        else
        {
            ReadOrCreatePacking(fullConfig, context, particles);
        }

        // Setting immobile particles in different ways

        Nullable<bool> shouldSuppressCrystallization = fullConfig.generationConfig.shouldSuppressCrystallization;
        if (shouldSuppressCrystallization.hasValue && shouldSuppressCrystallization.value)
        {
            immobileParticlesService->SetContext(context);
            immobileParticlesService->SetAndArrangeImmobileParticles(particles);
        }

        string immobileParticlesPath = Path::Append(fullConfig.generationConfig.baseFolder, IMMOBILE_PARTICLES_FILE_NAME);
        if (Path::Exists(immobileParticlesPath))
        {
            vector<ParticleIndex> immobileParticleIndexes;
            packingSerializer->ReadImmobileParticleIndexes(immobileParticlesPath, &immobileParticleIndexes);
            for (std::size_t immobileParticleIndexIndex = 0; immobileParticleIndexIndex < immobileParticleIndexes.size(); ++immobileParticleIndexIndex)
            {
                ParticleIndex immobileParticleIndex = immobileParticleIndexes[immobileParticleIndexIndex];
                particlesRef[immobileParticleIndex].isImmobile = true;
            }
        }

        if (context.activeGeometry != NULL)
        {
            SetActiveParticlesByActiveGeometry(fullConfig.systemConfig, *context.activeGeometry, particles);
        }
    }

    void GenerationManager::SetActiveParticlesByActiveGeometry(const SystemConfig& systemConfig, const IGeometry& activeGeometry, Packing* particles) const
    {
        Packing& particlesRef = *particles;

        int immobileParticlesCount = 0;
        for (ParticleIndex i = 0; i < systemConfig.particlesCount; ++i)
        {
            particlesRef[i].isImmobile = !activeGeometry.IsSphereInside(particlesRef[i].coordinates, particlesRef[i].diameter * 0.5);
            if (particlesRef[i].isImmobile)
            {
                immobileParticlesCount++;
            }
        }

        if (immobileParticlesCount > 0)
        {
            printf("mobileParticlesCount = %d, immobileParticlesCount = %d\n", systemConfig.particlesCount - immobileParticlesCount, immobileParticlesCount);
        }
    }

    void GenerationManager::ReadOrCreatePacking(const ExecutionConfig& fullConfig, const ModellingContext& context, Packing* particles) const
    {
        string baseFolder = fullConfig.generationConfig.baseFolder;
        int particlesCount = fullConfig.systemConfig.particlesCount;
        particles->resize(particlesCount);

        if (fullConfig.generationConfig.shouldStartGeneration.hasValue && fullConfig.generationConfig.shouldStartGeneration.value)
        {
            packingSerializer->ReadParticleDiameters(Path::Append(baseFolder, DIAMETERS_FILE_NAME), particles);

            if (fullConfig.generationConfig.initialParticleDistribution == InitialParticleDistribution::Poisson)
            {
                // TODO: create PoissonGenerator PoissonInCellsGenerator for all geometries, just call IsSphereInside for the given geometry until the sphere is in the geometry.
                BulkPoissonGenerator initialGenerator;
                CreateInitialPacking(fullConfig, context, particles, &initialGenerator);
            }
            else
            {
                BulkPoissonInCellsGenerator initialGenerator;
                CreateInitialPacking(fullConfig, context, particles, &initialGenerator);
            }

//            HcpGenerator initialGenerator;
//            CreateInitialPacking(fullConfig, context, particles, &initialGenerator);

            packingSerializer->SerializePacking(Path::Append(baseFolder, INIT_PACKING_FILE_NAME), *particles);
        }
        else
        {
            packingSerializer->ReadPacking(Path::Append(baseFolder, PACKING_FILE_NAME), particles);
        }
    }

    void GenerationManager::CreateInitialPacking(const ExecutionConfig& fullConfig, const ModellingContext& context, Packing* particles, IPackingGenerator* initialGenerator) const
    {
        initialGenerator->SetContext(context);
        initialGenerator->SetGenerationConfig(fullConfig.generationConfig);
        initialGenerator->ArrangePacking(particles);
    }

    GenerationManager::~GenerationManager()
    {

    }
}


#include <GL/freeglut.h>
#include <Eigen/Eigen>
#include <vector>
#include <iostream>
#include <iomanip>
#include <fstream>

#include "structures.h"
#include "transformations.h"
#include "cauchy.h"
#include "../../python-scripts/constraints/relative_pose_tait_bryan_wc_jacobian.h"
#include "../../python-scripts/constraints/relative_pose_rodrigues_wc_jacobian.h"
#include "../../python-scripts/constraints/relative_pose_quaternion_wc_jacobian.h"
#include "../../python-scripts/constraints/quaternion_constraint_jacobian.h"
#include "../../python-scripts/constraints/relative_pose_wc_jacobian.h"

#include "manif/SE2.h"

using manif::SE2d;
using manif::SE2Tangentd;

const unsigned int window_width = 1920;
const unsigned int window_height = 1080;
int mouse_old_x, mouse_old_y;
int mouse_buttons = 0;
float rotate_x = 0.0, rotate_y = 0.0;
float translate_z = -10.0;
float translate_x, translate_y = 0.0;

bool initGL(int *argc, char **argv);
void display();
void keyboard(unsigned char key, int x, int y);
void mouse(int button, int state, int x, int y);
void motion(int x, int y);
void reshape(int w, int h);
void printHelp();

std::vector<Eigen::Affine3d> m_poses;
std::vector<std::tuple<int, int, Eigen::Affine3d>> edges_g2o;
std::vector<std::vector<double>> edges_g2o_w;

int main(int argc, char *argv[]){
	if(argc != 2){
		std::cout << "USAGE: " << argv[0] << " file.g2o" << std::endl;
		return 1;
	}

	std::ifstream g2o_file(argv[1]);

	std::string line;
    while (std::getline(g2o_file, line)) {
        std::stringstream line_stream(line);
        std::string class_element;
        line_stream >> class_element;
        if (class_element == "VERTEX_SE2"){
            int id=0;
            TaitBryanPose p;
            line_stream >> id;
            line_stream >> p.px;
            line_stream >> p.py;
            line_stream >> p.ka;

            Eigen::Affine3d m = affine_matrix_from_pose_tait_bryan(p);
            m_poses.push_back(m);
        }
        else if(class_element == "EDGE_SE2"){
            int pose_id1,pose_id2;
            line_stream>>pose_id1;
            line_stream>>pose_id2;
            TaitBryanPose p;
            line_stream >> p.px;
            line_stream >> p.py;
            line_stream >> p.ka;

            Eigen::Affine3d m = affine_matrix_from_pose_tait_bryan(p);

            edges_g2o.emplace_back(std::make_tuple(pose_id1,pose_id2, m));

            std::vector<double> w(6);
			for(size_t i = 0 ; i < 6; i++){
				line_stream >> w[i];
			}
			std::vector<double> ww(3);
			ww[0] = w[0];
			ww[1] = w[3];
			ww[2] = w[5];
			edges_g2o_w.push_back(ww);

        }
    }
    g2o_file.close();

	if (false == initGL(&argc, argv)) {
		return 4;
	}

	printHelp();
	glutDisplayFunc(display);
	glutKeyboardFunc(keyboard);
	glutMouseFunc(mouse);
	glutMotionFunc(motion);
	glutMainLoop();

	return 0;
}

bool initGL(int *argc, char **argv) {
	glutInit(argc, argv);
	glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE);
	glutInitWindowSize(window_width, window_height);
	glutCreateWindow("pose2d_graph_slam");
	glutDisplayFunc(display);
	glutKeyboardFunc(keyboard);
	glutMotionFunc(motion);

	// default initialization
	glClearColor(1.0, 1.0, 1.0, 1.0);
	glEnable(GL_DEPTH_TEST);

	// viewport
	glViewport(0, 0, window_width, window_height);

	// projection
	glMatrixMode(GL_PROJECTION);
	glLoadIdentity();
	gluPerspective(60.0, (GLfloat) window_width / (GLfloat) window_height, 0.01,
			10000.0);
	glutReshapeFunc(reshape);

	return true;
}

void display() {
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();
	glTranslatef(translate_x, translate_y, translate_z);
	glRotatef(rotate_x, 1.0, 0.0, 0.0);
	glRotatef(rotate_y, 0.0, 0.0, 1.0);

	glBegin(GL_LINES);
	glColor3f(1.0f, 0.0f, 0.0f);
	glVertex3f(0.0f, 0.0f, 0.0f);
	glVertex3f(1.0f, 0.0f, 0.0f);

	glColor3f(0.0f, 1.0f, 0.0f);
	glVertex3f(0.0f, 0.0f, 0.0f);
	glVertex3f(0.0f, 1.0f, 0.0f);

	glColor3f(0.0f, 0.0f, 1.0f);
	glVertex3f(0.0f, 0.0f, 0.0f);
	glVertex3f(0.0f, 0.0f, 1.0f);
	glEnd();

    glBegin(GL_LINES);
    for(size_t i = 0; i < edges_g2o.size(); i++){
        const int pose_1 = std::get<0>(edges_g2o[i]);
        const int pose_2 = std::get<1>(edges_g2o[i]);
        glVertex3f(m_poses[pose_1](0,3), m_poses[pose_1](1,3), m_poses[pose_1](2,3) );
        glVertex3f(m_poses[pose_2](0,3), m_poses[pose_2](1,3), m_poses[pose_2](2,3) );
    }
    glEnd();

    glColor3f(1,0,1);
    glPointSize(5);
    glBegin(GL_POINTS);
    for(size_t i = 0; i < m_poses.size(); i++){
        glVertex3f(m_poses[i](0,3), m_poses[i](1,3), m_poses[i](2,3) );
    }
    glEnd();

    /*glBegin(GL_LINES);
    for(size_t i = 0; i < m_poses.size(); i++){

		glColor3f(1.0f, 0.0f, 0.0f);
		glVertex3f(m_poses[i](0,3), m_poses[i](1,3), m_poses[i](2,3));
		glVertex3f(m_poses[i](0,3) + m_poses[i](0,0), m_poses[i](1,3) + m_poses[i](1,0), m_poses[i](2,3) + + m_poses[i](2,0));

		glColor3f(0.0f, 1.0f, 0.0f);
		glVertex3f(m_poses[i](0,3), m_poses[i](1,3), m_poses[i](2,3));
		glVertex3f(m_poses[i](0,3) + m_poses[i](0,1), m_poses[i](1,3) + m_poses[i](1,1), m_poses[i](2,3) + + m_poses[i](2,1));

		glColor3f(0.0f, 0.0f, 1.0f);
		glVertex3f(m_poses[i](0,3), m_poses[i](1,3), m_poses[i](2,3));
		glVertex3f(m_poses[i](0,3) + m_poses[i](0,2), m_poses[i](1,3) + m_poses[i](1,2), m_poses[i](2,3) + + m_poses[i](2,2));

	}
    glEnd();*/

	glutSwapBuffers();
}

void keyboard(unsigned char key, int /*x*/, int /*y*/) {
	switch (key) {
		case (27): {
			glutDestroyWindow(glutGetWindow());
			return;
		}
		case 'n':{
			for(size_t i = 0 ; i < m_poses.size(); i++){
				TaitBryanPose pose = pose_tait_bryan_from_affine_matrix(m_poses[i]);
				pose.px += ((float(rand()%1000000))/1000000.0f - 0.5) * 1.1;
				pose.py += ((float(rand()%1000000))/1000000.0f - 0.5) * 1.1;
				pose.pz += ((float(rand()%1000000))/1000000.0f - 0.5) * 1.1;
				pose.om += ((float(rand()%1000000))/1000000.0f - 0.5) * 0.11;
				pose.fi += ((float(rand()%1000000))/1000000.0f - 0.5) * 0.11;
				pose.ka += ((float(rand()%1000000))/1000000.0f - 0.5) * 0.11;
				m_poses[i] = affine_matrix_from_pose_tait_bryan(pose);
			}
			break;
		}
		case 't':{
			std::vector<Eigen::Triplet<double>> tripletListA;
			std::vector<Eigen::Triplet<double>> tripletListP;
			std::vector<Eigen::Triplet<double>> tripletListB;

			std::vector<TaitBryanPose> poses;

			for(size_t i = 0 ; i < m_poses.size(); i++){
				poses.push_back(pose_tait_bryan_from_affine_matrix(m_poses[i]));
			}

			for(size_t i = 0 ; i < edges_g2o.size(); i++){
				const int first = std::get<0>(edges_g2o[i]);
				const int second = std::get<1>(edges_g2o[i]);
				const Eigen::Affine3d& rel = std::get<2>(edges_g2o[i]);
				TaitBryanPose pose_rel = pose_tait_bryan_from_affine_matrix(rel);

				TaitBryanPose from = poses[first];
				TaitBryanPose to = poses[second];

				Eigen::Matrix<double, 6, 1> delta;
				relative_pose_obs_eq_tait_bryan_wc_case1(
					delta,
					poses[first].px,
					poses[first].py,
					poses[first].pz,
					poses[first].om,
					poses[first].fi,
					poses[first].ka,
					poses[second].px,
					poses[second].py,
					poses[second].pz,
					poses[second].om,
					poses[second].fi,
					poses[second].ka,
					pose_rel.px,
					pose_rel.py,
					pose_rel.pz,
					pose_rel.om,
					pose_rel.fi,
					pose_rel.ka);

				Eigen::Matrix<double, 6, 12, Eigen::RowMajor> jacobian;
				relative_pose_obs_eq_tait_bryan_wc_case1_jacobian(jacobian,
					poses[first].px,
					poses[first].py,
					poses[first].pz,
					poses[first].om,
					poses[first].fi,
					poses[first].ka,
					poses[second].px,
					poses[second].py,
					poses[second].pz,
					poses[second].om,
					poses[second].fi,
					poses[second].ka);
				int ir = tripletListB.size();

				int ic_1 = first * 3;
				int ic_2 = second * 3;

				tripletListA.emplace_back(ir + 0, ic_1    , -jacobian(0,0));
				tripletListA.emplace_back(ir + 0, ic_1 + 1, -jacobian(0,1));
				tripletListA.emplace_back(ir + 0, ic_1 + 2, -jacobian(0,5));

				tripletListA.emplace_back(ir + 0, ic_2    , -jacobian(0,6));
				tripletListA.emplace_back(ir + 0, ic_2 + 1, -jacobian(0,7));
				tripletListA.emplace_back(ir + 0, ic_2 + 2, -jacobian(0,11));

				tripletListA.emplace_back(ir + 1, ic_1    , -jacobian(1,0));
				tripletListA.emplace_back(ir + 1, ic_1 + 1, -jacobian(1,1));
				tripletListA.emplace_back(ir + 1, ic_1 + 2, -jacobian(1,5));

				tripletListA.emplace_back(ir + 1, ic_2    , -jacobian(1,6));
				tripletListA.emplace_back(ir + 1, ic_2 + 1, -jacobian(1,7));
				tripletListA.emplace_back(ir + 1, ic_2 + 2, -jacobian(1,11));

				tripletListA.emplace_back(ir + 2, ic_1    , -jacobian(5,0));
				tripletListA.emplace_back(ir + 2, ic_1 + 1, -jacobian(5,1));
				tripletListA.emplace_back(ir + 2, ic_1 + 2, -jacobian(5,5));

				tripletListA.emplace_back(ir + 2, ic_2    , -jacobian(5,6));
				tripletListA.emplace_back(ir + 2, ic_2 + 1, -jacobian(5,7));
				tripletListA.emplace_back(ir + 2, ic_2 + 2, -jacobian(5,11));

				float angle_diff = delta(5,0);
				if(fabs(angle_diff) > M_PI){
					angle_diff -= 2.0*M_PI;
				}
				if(fabs(angle_diff)< -M_PI){
					angle_diff += 2.0*M_PI;
				}

				tripletListB.emplace_back(ir,     0, delta(0,0));
				tripletListB.emplace_back(ir + 1, 0, delta(1,0));
				tripletListB.emplace_back(ir + 2, 0, angle_diff);

				/*if(abs(first-second)==1){
					tripletListP.emplace_back(ir ,    ir,     1000);
					tripletListP.emplace_back(ir + 1, ir + 1, 1000);
					tripletListP.emplace_back(ir + 2, ir + 2, 1000);
				}else{
					tripletListP.emplace_back(ir ,    ir,     1);
					tripletListP.emplace_back(ir + 1, ir + 1, 1);
					tripletListP.emplace_back(ir + 2, ir + 2, 0.000001);
				}*/
				tripletListP.emplace_back(ir ,    ir,     cauchy(delta(0,0),1) * edges_g2o_w[i][0]);
				tripletListP.emplace_back(ir + 1, ir + 1, cauchy(delta(1,0),1) * edges_g2o_w[i][1]);
				tripletListP.emplace_back(ir + 2, ir + 2, cauchy(angle_diff,1) * edges_g2o_w[i][2]);

				//tripletListP.emplace_back(ir ,    ir,     edges_g2o_w[i][0]);
				//tripletListP.emplace_back(ir + 1, ir + 1, edges_g2o_w[i][1]);
				//tripletListP.emplace_back(ir + 2, ir + 2, edges_g2o_w[i][2]);
			}

			int ir = tripletListB.size();
			tripletListA.emplace_back(ir     , 0, 1);
			tripletListA.emplace_back(ir + 1 , 1, 1);
			tripletListA.emplace_back(ir + 2 , 2, 1);

			tripletListP.emplace_back(ir     , ir,     1000000);
			tripletListP.emplace_back(ir + 1 , ir + 1, 1000000);
			tripletListP.emplace_back(ir + 2 , ir + 2, 1000000);

			tripletListB.emplace_back(ir     , 0, 0);
			tripletListB.emplace_back(ir + 1 , 0, 0);
			tripletListB.emplace_back(ir + 2 , 0, 0);

			Eigen::SparseMatrix<double> matA(tripletListB.size(), m_poses.size() * 3);
			Eigen::SparseMatrix<double> matP(tripletListB.size(), tripletListB.size());
			Eigen::SparseMatrix<double> matB(tripletListB.size(), 1);

			matA.setFromTriplets(tripletListA.begin(), tripletListA.end());
			matP.setFromTriplets(tripletListP.begin(), tripletListP.end());
			matB.setFromTriplets(tripletListB.begin(), tripletListB.end());

			Eigen::SparseMatrix<double> AtPA(m_poses.size() * 3 , m_poses.size() * 3);
			Eigen::SparseMatrix<double> AtPB(m_poses.size() * 3 , 1);

			{
			Eigen::SparseMatrix<double> AtP = matA.transpose() * matP;
			AtPA = (AtP) * matA;
			AtPB = (AtP) * matB;
			}

			tripletListA.clear();
			tripletListP.clear();
			tripletListB.clear();
			//LM
			/*{
				for(size_t i = 0 ; i < m_poses.size() * 3; i++){
					tripletListA.emplace_back(i, i, 10);
				}
				Eigen::SparseMatrix<double> matLM(m_poses.size() * 3, m_poses.size() * 3);
				matLM.setFromTriplets(tripletListA.begin(), tripletListA.end());

				AtPA = AtPA + matLM;
			}*/

			std::cout << "AtPA.size: " << AtPA.size() << std::endl;
			std::cout << "AtPB.size: " << AtPB.size() << std::endl;

			std::cout << "start solving AtPA=AtPB" << std::endl;
			Eigen::SimplicialCholesky<Eigen::SparseMatrix<double>> solver(AtPA);

			std::cout << "x = solver.solve(AtPB)" << std::endl;
			Eigen::SparseMatrix<double> x = solver.solve(AtPB);

			std::vector<double> h_x;

			for (int k=0; k<x.outerSize(); ++k){
				for (Eigen::SparseMatrix<double>::InnerIterator it(x,k); it; ++it){
					h_x.push_back(it.value());
				}
			}

			std::cout << "h_x.size(): " << h_x.size() << std::endl;
			std::cout << "AtPA=AtPB SOLVED" << std::endl;

			if(h_x.size() == 3 * m_poses.size()){
				int counter = 0;

				for(size_t i = 0; i < m_poses.size(); i++){
					TaitBryanPose pose = pose_tait_bryan_from_affine_matrix(m_poses[i]);
					pose.px += h_x[counter++];
					pose.py += h_x[counter++];
					pose.ka += h_x[counter++];
					m_poses[i] = affine_matrix_from_pose_tait_bryan(pose);
				}
				std::cout << "optimizing with tait bryan finished" << std::endl;
			}else{
				std::cout << "optimizing with tait bryan FAILED" << std::endl;
			}
			break;
		}
		case 'y':{
			std::vector<Eigen::Triplet<double>> tripletListA;
			std::vector<Eigen::Triplet<double>> tripletListP;
			std::vector<Eigen::Triplet<double>> tripletListB;

			std::vector<SE2d> X;

			for(size_t i = 0 ; i < m_poses.size(); i++){
				TaitBryanPose p = pose_tait_bryan_from_affine_matrix(m_poses[i]);
				X.push_back(SE2d(p.px, p.py, p.ka));
			}

			for(size_t i = 0 ; i < edges_g2o.size(); i++){
				const int first = std::get<0>(edges_g2o[i]);
				const int second = std::get<1>(edges_g2o[i]);
				const Eigen::Affine3d& rel = std::get<2>(edges_g2o[i]);

				TaitBryanPose p = pose_tait_bryan_from_affine_matrix(rel);
				SE2d U = SE2d(p.px, p.py, p.ka);

				SE2Tangentd     d;
				SE2Tangentd     u;
				Eigen::Matrix<double, 3, 3>         J_d_xi, J_d_xj;

				SE2d         Xi,
							 Xj;

				Xi = X[first];
				Xj = X[second];

				d  = Xj.rminus(Xi, J_d_xj, J_d_xi);
				u = U.log();

				int ir = tripletListB.size();
				int ic_1 = first * 3;
				int ic_2 = second * 3;

				for(size_t row = 0 ; row < 3; row ++){
					tripletListA.emplace_back(ir + row, ic_1    , -J_d_xi(row,0));
					tripletListA.emplace_back(ir + row, ic_1 + 1, -J_d_xi(row,1));
					tripletListA.emplace_back(ir + row, ic_1 + 2, -J_d_xi(row,2));

					tripletListA.emplace_back(ir + row, ic_2    , -J_d_xj(row,0));
					tripletListA.emplace_back(ir + row, ic_2 + 1, -J_d_xj(row,1));
					tripletListA.emplace_back(ir + row, ic_2 + 2, -J_d_xj(row,2));
				}

				SE2Tangentd delta = d - u;

				tripletListB.emplace_back(ir,     0, delta.coeffs()(0));
				tripletListB.emplace_back(ir + 1, 0, delta.coeffs()(1));
				tripletListB.emplace_back(ir + 2, 0, delta.coeffs()(2));

				/*if(abs(first - second) == 1){
					tripletListP.emplace_back(ir ,    ir,     1000);
					tripletListP.emplace_back(ir + 1, ir + 1, 1000);
					tripletListP.emplace_back(ir + 2, ir + 2, 1000);
				}else{
					tripletListP.emplace_back(ir ,    ir,     1);
					tripletListP.emplace_back(ir + 1, ir + 1, 1);
					tripletListP.emplace_back(ir + 2, ir + 2, 0.000001);
				}*/

				tripletListP.emplace_back(ir ,    ir,     cauchy(delta.coeffs()(0),1) * edges_g2o_w[i][0]);
				tripletListP.emplace_back(ir + 1, ir + 1, cauchy(delta.coeffs()(1),1) * edges_g2o_w[i][1]);
				tripletListP.emplace_back(ir + 2, ir + 2, cauchy(delta.coeffs()(2),1) * edges_g2o_w[i][2]);
			}

			int ir = tripletListB.size();
			tripletListA.emplace_back(ir     , 0, 1);
			tripletListA.emplace_back(ir + 1 , 1, 1);
			tripletListA.emplace_back(ir + 2 , 2, 1);

			tripletListP.emplace_back(ir     , ir,     1000000);
			tripletListP.emplace_back(ir + 1 , ir + 1, 1000000);
			tripletListP.emplace_back(ir + 2 , ir + 2, 1000000);

			tripletListB.emplace_back(ir     , 0, 0);
			tripletListB.emplace_back(ir + 1 , 0, 0);
			tripletListB.emplace_back(ir + 2 , 0, 0);

			Eigen::SparseMatrix<double> matA(tripletListB.size(), m_poses.size() * 3);
			Eigen::SparseMatrix<double> matP(tripletListB.size(), tripletListB.size());
			Eigen::SparseMatrix<double> matB(tripletListB.size(), 1);

			matA.setFromTriplets(tripletListA.begin(), tripletListA.end());
			matP.setFromTriplets(tripletListP.begin(), tripletListP.end());
			matB.setFromTriplets(tripletListB.begin(), tripletListB.end());

			Eigen::SparseMatrix<double> AtPA(m_poses.size() * 3 , m_poses.size() * 3);
			Eigen::SparseMatrix<double> AtPB(m_poses.size() * 3 , 1);

			{
			Eigen::SparseMatrix<double> AtP = matA.transpose() * matP;
			AtPA = (AtP) * matA;
			AtPB = (AtP) * matB;
			}

			tripletListA.clear();
			tripletListP.clear();
			tripletListB.clear();

			std::cout << "start solving AtPA=AtPB" << std::endl;
			Eigen::SimplicialCholesky<Eigen::SparseMatrix<double>> solver(AtPA);

			std::cout << "x = solver.solve(AtPB)" << std::endl;
			Eigen::SparseMatrix<double> x = solver.solve(AtPB);

			std::vector<double> h_x;

			for (int k=0; k<x.outerSize(); ++k){
				for (Eigen::SparseMatrix<double>::InnerIterator it(x,k); it; ++it){
					h_x.push_back(it.value());
				}
			}

			if(X.size() * 3 == h_x.size()){

				int counter = 0;
				for(size_t i = 0 ; i < X.size(); i++){
					SE2Tangentd     dx;
					dx.coeffs()(0) = h_x[counter++];
					dx.coeffs()(1) = h_x[counter++];
					dx.coeffs()(2) = h_x[counter++];
					X[i] = X[i] +  dx;
				}

				for (int i = 0 ; i < m_poses.size(); i++){
					TaitBryanPose p;
					p.px = X[i].translation()(0);
					p.py = X[i].translation()(1);
					p.ka = X[i].angle();

					m_poses[i] = affine_matrix_from_pose_tait_bryan(p);
				}
				std::cout << "h_x.size(): " << h_x.size() << std::endl;
				std::cout << "AtPA=AtPB SOLVED" << std::endl;
			}else{
				std::cout << "h_x.size(): " << h_x.size() << std::endl;
				std::cout << "AtPA=AtPB FAILED" << std::endl;
			}

			break;
		}
	}
	printHelp();
	glutPostRedisplay();
}

void mouse(int button, int state, int x, int y) {
	if (state == GLUT_DOWN) {
		mouse_buttons |= 1 << button;
	} else if (state == GLUT_UP) {
		mouse_buttons = 0;
	}

	mouse_old_x = x;
	mouse_old_y = y;
}

void motion(int x, int y) {
	float dx, dy;
	dx = (float) (x - mouse_old_x);
	dy = (float) (y - mouse_old_y);

	if (mouse_buttons & 1) {
		rotate_x += dy * 0.2f;
		rotate_y += dx * 0.2f;

	} else if (mouse_buttons & 4) {
		translate_z += dy * 0.05f;
	} else if (mouse_buttons & 3) {
		translate_x += dx * 0.05f;
		translate_y -= dy * 0.05f;
	}

	mouse_old_x = x;
	mouse_old_y = y;

	glutPostRedisplay();
}

void reshape(int w, int h) {
	glViewport(0, 0, (GLsizei) w, (GLsizei) h);
	glMatrixMode(GL_PROJECTION);
	glLoadIdentity();
	gluPerspective(60.0, (GLfloat) w / (GLfloat) h, 0.01, 10000.0);
	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();
}

void printHelp() {
	std::cout << "-------help-------" << std::endl;
	std::cout << "n: add noise to poses" << std::endl;
	std::cout << "t: optimize (Tait-Bryan)" << std::endl;
	std::cout << "y: optimize (Lie Algebra: manif lib)" << std::endl;
}


#include "SCEngine.hpp"

#include <string>
#include <glog/logging.h>
#include <algorithm>
#include <fstream>
#include <cmath>
#include <mutex>
#include <boost/date_time/posix_time/posix_time.hpp>
#include <petuum_ps_common/include/petuum_ps.hpp>
#include <io/general_fstream.hpp>

#include "util/Eigen/Dense"
#include "util/context.hpp"

namespace sparsecoding {

    // Constructor
    SCEngine::SCEngine(): thread_counter_(0) {
        // timer
        initT_ = boost::posix_time::microsec_clock::local_time();

        /* context */
        lda::Context & context = lda::Context::get_instance();
        // input and output
        data_file_ = context.get_string("data_file");
        input_data_format_ = context.get_string("input_data_format");
        is_partitioned_ = context.get_int32("is_partitioned");
        output_path_ = context.get_string("output_path");
        output_data_format_ = context.get_string("output_data_format");
        maximum_running_time_ = context.get_double("maximum_running_time");
        load_cache_ = context.get_int32("load_cache");
        cache_path_ = context.get_string("cache_path");

        // objective function parameters
        int m = context.get_int32("m");
        int n = context.get_int32("n");
        dictionary_size_ = context.get_int32("dictionary_size");
        C_ = context.get_double("c");
        lambda_ = context.get_double("lambda");

        // petuum parameters
        client_id_ = context.get_int32("client_id");
        num_clients_ = context.get_int32("num_clients");
        num_worker_threads_ = context.get_int32("num_worker_threads");

        // optimization parameters
        num_epochs_ = context.get_int32("num_epochs");
        minibatch_size_ = context.get_int32("minibatch_size");
        num_eval_minibatch_ = context.get_int32("num_eval_minibatch");
        num_eval_samples_ = context.get_int32("num_eval_samples");
        init_step_size_B_ = context.get_double("init_step_size_B");
        step_size_offset_B_ = context.get_double("step_size_offset_B");
        step_size_pow_B_ = context.get_double("step_size_pow_B");
        num_iter_S_per_minibatch_ =
            context.get_int32("num_iter_S_per_minibatch");
        init_step_size_S_ = context.get_double("init_step_size_S");
        step_size_offset_S_ = context.get_double("step_size_offset_S");
        step_size_pow_S_ = context.get_double("step_size_pow_S");
        // default step size
        if (init_step_size_B_ < FLT_MIN) {
            init_step_size_B_ = context.get_double("init_step_size")
                / num_clients_ / num_worker_threads_;
            step_size_offset_B_ = context.get_double("step_size_offset");
            step_size_pow_B_ = context.get_double("step_size_pow");
        }
        if (init_step_size_S_ < FLT_MIN) {
            init_step_size_S_ = context.get_double("init_step_size") / m;
            step_size_offset_S_ = context.get_double("step_size_offset");
            step_size_pow_S_ = context.get_double("step_size_pow");
        }
        init_S_low_ = context.get_double("init_S_low");
        init_S_high_ = context.get_double("init_S_high");
        init_B_low_ = context.get_double("init_B_low");
        init_B_high_ = context.get_double("init_B_high");

        /* Init matrices */
        // Partition by column id mod num_clients_
        int client_n = (n - (n / num_clients_) * num_clients_ > client_id_)?
            n / num_clients_ + 1: n / num_clients_;
        // Init matrix loader of data matrix X
        if (is_partitioned_) {
            X_matrix_loader_.Init(data_file_, input_data_format_, m, client_n);
        } else {
            X_matrix_loader_.Init(data_file_, input_data_format_, m, n,
                    client_id_, num_clients_);
        }

        // Init matrix loader of coefficients S
        if (dictionary_size_ == 0)
            dictionary_size_ = n;
        S_matrix_loader_.Init(dictionary_size_, client_n,
            init_S_low_, init_S_high_);

        int max_client_n = ceil(float(n) / num_clients_);
        int iter_minibatch =
            ceil(float(max_client_n) / num_worker_threads_ / minibatch_size_);
        num_eval_per_client_ =
            (num_epochs_ * iter_minibatch - 1)
              / num_eval_minibatch_ + 1;

        // Ouput parameters
        LOG(INFO) << "Sparse Coding running on "
            << num_clients_
            << ((num_clients_ > 1)? " clients, each with ": " client, with ")
            << num_worker_threads_ << " threads";
        LOG(INFO) << "Input file: " << data_file_;
        LOG(INFO) << "Matrix size: " << n << " by " << m;
        LOG(INFO) << "Dictionary size: " << dictionary_size_;
        LOG(INFO) << "C: " << C_ << ", lambda: " << lambda_;
        LOG(INFO) << "Minibatch size: " << minibatch_size_;
        LOG(INFO) << "Epochs: " << num_epochs_
            << " (Max iter: " << iter_minibatch  * num_epochs_ << ")";
        LOG(INFO) << "Evaluate loss per " << num_eval_minibatch_ << " iter "
            << "by evaluating " << num_eval_samples_ << " samples per thread"
            << " (*" << num_clients_ * num_worker_threads_<< " = "
            << num_eval_samples_ * num_clients_ * num_worker_threads_ << ")";
        LOG(INFO) << "step size of B at ith iteration: " << init_step_size_B_
            << " * (i + " << step_size_offset_B_ << ")^(-" << step_size_pow_B_
            << ")";
        LOG(INFO) << "step size of S at ith iteration: " << init_step_size_S_
            << " * (i + " << step_size_offset_S_ << ")^(-" << step_size_pow_S_
            << ")";
        LOG(INFO) << "Iter of S in each S loop: " << num_iter_S_per_minibatch_;
    }

    // Helper function, regularize a vector vec
    // such that its l2-norm is smaller than C
    inline void RegVec(std::vector<float> & vec, float C,
            std::vector<float> & vec_result) {
        float sum = 0.0;
        int len = vec.size();
        for (int i = 0; i < len; ++i) {
            sum += vec[i] * vec[i];
        }
        float ratio = (sum > C? sqrt(C / sum): 1.0);
        for (int i = 0; i < len; ++i) {
            vec_result[i] = vec[i] * ratio;
        }
    }

    // Save results: dicitonary B, coefficients S, loss evaluated on different
    // machines, time between evaluations to disk.
    // Shall be called after calling petuum::PSTableGroup::GlobalBarrier()
    void SCEngine::SaveResults(int thread_id, petuum::Table<float> & B_table,
            petuum::Table<float> & loss_table) {
        // size of matrices
        int m = X_matrix_loader_.GetM();
        int client_n = X_matrix_loader_.GetClientN();

        // Caches
        std::vector<float> B_row_cache(m), S_cache(dictionary_size_);

        // Output files
        //std::ofstream fout_loss, fout_B, fout_S, fout_time;

        // Only thread 0 of client 0 write dictionary B, loss and time to disk
        if (client_id_ == 0 && thread_id == 0) {
            // Write loss to disk
            std::string loss_filename = output_path_ + "/loss.txt";
	    petuum::io::ofstream fout_loss(loss_filename);

            //fout_loss.open(loss_filename.c_str());
            LOG(INFO) << "Writing loss result to directory: " << output_path_;
            petuum::RowAccessor row_acc;
            std::vector<float> petuum_row_cache(m);
            for (int iter = 0; iter < num_eval_per_client_; ++iter) {
                for (int client = 0; client < num_clients_; ++client) {
                    int row_id = client * num_eval_per_client_ + iter;
                    const auto & row = loss_table.
                        Get<petuum::DenseRow<float> >(row_id, &row_acc);
                    row.CopyToVector(&petuum_row_cache);
                    if (std::abs(petuum_row_cache[0]) > FLT_MIN) {
                        fout_loss << petuum_row_cache[0] << "\t";
                    } else {
                        fout_loss << "N/A" << "\t";
                    }
                }
                fout_loss << "\n";
            }
            fout_loss.close();
            // Write time to disk
            std::string time_filename = output_path_ + "/time.txt";
	    petuum::io::ofstream fout_time(time_filename);

            //fout_time.open(time_filename.c_str());
            for (int iter = 0; iter < num_eval_per_client_; ++iter) {
                for (int client = 0; client < num_clients_; ++client) {
                    int row_id =
                        (client + num_clients_) * num_eval_per_client_ + iter;
                    const auto & row = loss_table.
                        Get<petuum::DenseRow<float> >(row_id, &row_acc);
                    row.CopyToVector(&petuum_row_cache);
                    if (std::abs(petuum_row_cache[0]) > FLT_MIN) {
                        fout_time << petuum_row_cache[0] << "\t";
                    } else {
                        fout_time << "N/A" << "\t";
                    }
                }
                fout_time << "\n";
            }
            fout_time.close();
            // Write dictionary B to disk
            // with filename output_path_/B.[txt|bin]
            //if (output_data_format_ == "text") {
                std::string B_filename = output_path_ + "/B.txt";
                //fout_B.open(B_filename.c_str());
		petuum::io::ofstream fout_B(B_filename);

		for (int row_id = 0; row_id < dictionary_size_; ++row_id) {
                  const auto & row = B_table.
                    Get<petuum::DenseRow<float> >(row_id, &row_acc);
                  row.CopyToVector(&petuum_row_cache);
                  // Regularize by C_
                  RegVec(petuum_row_cache, C_, B_row_cache);
                  for (int col_id = 0; col_id < m; ++col_id) {
                    if (std::abs(B_row_cache[col_id]) < INFINITESIMAL) {
                        B_row_cache[col_id] = 0.0;
                    }
                    fout_B << B_row_cache[col_id] << "\t";

                    fout_B << "\n";
                  }
                }
                fout_B.close();

        }
        // Thread 0 of each client save that client's part of S
        // to output_path_/S.[txt|bin].client_id_
        if (thread_id == 0) {

                std::string S_filename = output_path_ + "/S.txt."
                    + std::to_string(client_id_);
                //fout_S.open(S_filename.c_str());
		petuum::io::ofstream fout_S(S_filename);

            for (int col_id_client = 0; col_id_client < client_n;
                    ++col_id_client) {
                if (S_matrix_loader_.GetCol(col_id_client, S_cache)) {
                    for (int row_id = 0; row_id < dictionary_size_; ++row_id) {
                        if (std::abs(S_cache[row_id]) < INFINITESIMAL) {
                            S_cache[row_id] = 0.0;
                        }
                        fout_S << S_cache[row_id] << "\t";

                    }
                                                    fout_S << "\n";

                }
            }
            fout_S.close();
        }
    }

    // Init B and S from cache file
    void SCEngine::LoadCache(int thread_id, petuum::Table<float> & B_table) {
        // size of matrices
        int m = X_matrix_loader_.GetM();
        int client_n = X_matrix_loader_.GetClientN();
        if (client_id_ == 0 && thread_id == 0) {
            // Load B

            std::vector<float> B_row_cache(m);

            std::string B_filename = cache_path_ + "/B.txt";
            //fout_B.open(B_filename.c_str());
	    petuum::io::ifstream fout_B(B_filename);

            CHECK(fout_B.good())
                << "Cache file " << B_filename << " does not exist!";
            for (int row_id = 0; row_id < dictionary_size_; ++row_id) {
                petuum::UpdateBatch<float> B_update;
                for (int col_id = 0; col_id < m; ++col_id) {
                    if (input_data_format_ == "text") {
                        fout_B >> B_row_cache[col_id];
                    } else if (input_data_format_ == "binary") {
                        fout_B.read(reinterpret_cast<char*> (
                                &(B_row_cache[col_id])), 4);
                    }
                    B_update.Update(col_id, B_row_cache[col_id]);
                }
                B_table.BatchInc(row_id, B_update);
            }
            fout_B.close();
        }
        // Load S
        if (thread_id == 0) {

            std::vector<float> S_cache(dictionary_size_),
                S_inc_cache(dictionary_size_);

            std::string S_filename;
                S_filename = cache_path_ + "/S.txt." +
                    std::to_string(client_id_);
                   //fout_S.open(S_filename.c_str());
	    petuum::io::ifstream fout_S(S_filename);

            CHECK(fout_S.good())
                << "Cache file " << S_filename << " does not exist!";
               for (int col_id_client = 0; col_id_client < client_n;
                    ++col_id_client) {
                   if (S_matrix_loader_.GetCol(col_id_client, S_cache)) {
                       for (int row_id = 0; row_id < dictionary_size_;
                            ++row_id) {
                            fout_S >> S_inc_cache[row_id];
                                                S_inc_cache[row_id] =
                            S_inc_cache[row_id] - S_cache[row_id];
                       }
                    S_matrix_loader_.IncCol(col_id_client, S_inc_cache);
                   }
                S_matrix_loader_.GetCol(col_id_client, S_cache);
               }
               fout_S.close();
        }
    }

    // Init B table
    void SCEngine::InitRand(int thread_id, petuum::Table<float> & B_table) {
        if (thread_id != 0 || client_id_ != 0)
            return;
        // size of matrices
        int m = X_matrix_loader_.GetM();
        srand((unsigned)time(NULL));
        std::vector<float> B_row_cache(m);
        // petuum row accessor
        petuum::RowAccessor row_acc;
        for (int row_id = 0; row_id < dictionary_size_; ++row_id) {
            petuum::UpdateBatch<float> B_update;
            for (int col_id = 0; col_id < m; ++col_id) {
                B_row_cache[col_id] = double(rand()) / RAND_MAX *
                    (init_B_high_ - init_B_low_) + init_B_low_;
            }
            for (int col_id = 0; col_id < m; ++col_id) {
                B_update.Update(col_id, B_row_cache[col_id]);
            }
            B_table.BatchInc(row_id, B_update);
        }
    }

    // Stochastic Gradient Descent Optimization
    void SCEngine::Start() {
        // thread id on a client
        int thread_id = thread_counter_++;
        petuum::PSTableGroup::RegisterThread();
        LOG(INFO) << "client " << client_id_ << ", thread "
            << thread_id << " registers!";

        // Get dictionary table and loss table
        petuum::Table<float> B_table =
            petuum::PSTableGroup::GetTableOrDie<float>(0);
        petuum::Table<float> loss_table =
            petuum::PSTableGroup::GetTableOrDie<float>(1);

        // size of matrices
        int m = X_matrix_loader_.GetM();
        int client_n = X_matrix_loader_.GetClientN();

        // Cache dictionary table
        Eigen::MatrixXf petuum_table_cache(m, dictionary_size_);
        // Accumulate update of dictionary table in minibatch
        Eigen::MatrixXf petuum_update_cache(m, dictionary_size_);
        // Cache a column of coefficients S
        Eigen::VectorXf Sj(dictionary_size_);
        // Cache a column of update of S_j
        Eigen::VectorXf Sj_inc(dictionary_size_);
        // Cache a column of data X
        Eigen::VectorXf Xj(m);
        Eigen::VectorXf Xj_inc(m);
        // Cache a row of dictionary table
        std::vector<float> petuum_row_cache(m);

        // Register rows
        if (thread_id == 0) {
            for (int row_id = 0; row_id < dictionary_size_; ++row_id) {
                B_table.GetAsyncForced(row_id);
            }
            int loss_table_size = num_eval_per_client_ * num_clients_ * 2;
            for (int row_id = 0; row_id < loss_table_size; ++row_id) {
                loss_table.GetAsyncForced(row_id);
            }
        }
        petuum::PSTableGroup::GlobalBarrier();

        // initialize B
        STATS_APP_INIT_BEGIN();
        if (client_id_ == 0 && thread_id == 0) {
            LOG(INFO) << "starting to initialize B";
        }
        if (load_cache_) {// load B and S from cache
            LoadCache(thread_id, B_table);
        } else { // randomly init B
            InitRand(thread_id, B_table);
        }
        if (thread_id == 0 && client_id_ == 0) {
            LOG(INFO) << "matrix B initialization finished!";
        }
        petuum::PSTableGroup::GlobalBarrier();
        STATS_APP_INIT_END();

        // Optimization Loop
        // Timer
        boost::posix_time::ptime beginT =
            boost::posix_time::microsec_clock::local_time();
        // Step size for optimization
        float step_size_B = init_step_size_B_, step_size_S = init_step_size_S_;

        int num_minibatch = 0;
        for (int iter = 0; iter < num_epochs_; ++iter) {
            // how many minibatches per epoch
            int minibatch_per_epoch = (client_n / num_worker_threads_ > 0)?
                client_n / num_worker_threads_: 1;
            for (int iter_per_epoch = 0; iter_per_epoch * minibatch_size_
                    < minibatch_per_epoch; ++iter_per_epoch) {
                boost::posix_time::time_duration runTime =
                    boost::posix_time::microsec_clock::local_time() - initT_;
                // Terminate and save states to disk if running time exceeds
                // limit
                if (maximum_running_time_ > 0.0 &&
                        (float) runTime.total_milliseconds() >
                        maximum_running_time_*3600*1000) {
                    LOG(INFO) << "Maximum runtime limit activates, "
                        "terminating now!";
                    petuum::PSTableGroup::GlobalBarrier();
                    SaveResults(thread_id, B_table, loss_table);
                    petuum::PSTableGroup::DeregisterThread();
                    return;
                }
                // Update petuum table cache
                petuum::RowAccessor row_acc;
                for (int row_id = 0; row_id < dictionary_size_; ++row_id) {
                    const auto & row =
                        B_table.Get<petuum::DenseRow<float> >(row_id, &row_acc);
                    row.CopyToVector(&petuum_row_cache);
                    for (int col_id = 0; col_id < m; ++col_id) {
                        petuum_table_cache(col_id, row_id) =
                            petuum_row_cache[col_id];
                    }
                }
                // evaluate obj
                if (num_minibatch % num_eval_minibatch_ == 0) {
                    boost::posix_time::time_duration elapTime =
                        boost::posix_time::microsec_clock::local_time() - beginT;
                    // evaluate partial obj
                    double obj = 0.0;
                    double obj1 = 0.0;
                    double obj2 = 0.0;
                    int num_samples = num_eval_samples_;
                    for (int i = 0; i < dictionary_size_; ++i) {
                        float regularizer = petuum_table_cache.col(i).norm();
                        regularizer =
                            (regularizer > sqrt(C_))? sqrt(C_) / regularizer: 1.0;
                        petuum_table_cache.col(i) *= regularizer;
                    }
                    for (int i = 0; i < num_samples; ++i) {
                        int col_id_client = 0;
                        if (S_matrix_loader_.GetRandCol(col_id_client, Sj)
                                && X_matrix_loader_.GetCol(col_id_client, Xj)) {
                            Xj_inc = Xj - petuum_table_cache * Sj;
                            obj1 += Xj_inc.squaredNorm() / num_samples;
                            obj2 += lambda_ * Sj.lpNorm<1>() / num_samples;
                        }
                    }
                    obj = obj1 + obj2;
                    LOG(INFO) << "iter: " << num_minibatch << ", client "
                        << client_id_ << ", thread " << thread_id <<
                        " mean loss: " << obj
                        << ", mean reconstruction error: " << obj1
                        << ", mean regularization: " << obj2;
                    // update loss table
                    loss_table.Inc(client_id_ * num_eval_per_client_ +
                            num_minibatch / num_eval_minibatch_,  0,
                            obj/num_worker_threads_);
                    loss_table.Inc((num_clients_+client_id_) * num_eval_per_client_
                            + num_minibatch / num_eval_minibatch_, 0,
                            ((float) elapTime.total_milliseconds()) / 1000
                            / num_worker_threads_);
                    beginT = boost::posix_time::microsec_clock::local_time();
                }
                step_size_B = init_step_size_B_ *
                    pow(step_size_offset_B_ + num_minibatch,
                            -1.0*step_size_pow_B_);
                step_size_S = init_step_size_S_ *
                    pow(step_size_offset_S_ + num_minibatch,
                            -1.0*step_size_pow_S_);
                num_minibatch++;
                // clear update table
                petuum_update_cache.fill(0.0);
                // minibatch
                for (int k = 0; k < minibatch_size_; ++k) {
                    int col_id_client = 0;
                    if (S_matrix_loader_.GetRandCol(col_id_client, Sj)
                            && X_matrix_loader_.GetCol(col_id_client, Xj)) {
                        // update S_j
                        for (int iter_S = 0;
                                iter_S < num_iter_S_per_minibatch_; ++iter_S) {
                            // compute gradient of Sj
                            Sj_inc = step_size_S * 2.0 *
                                ( petuum_table_cache.transpose()
                                 * (Xj - petuum_table_cache * Sj) );
                            Sj.noalias() += Sj_inc;
                            // proximal gradient descent to deal with l1 subgradient
                            Sj = ((Sj.array() > step_size_S*lambda_).cast<float>()
                                * (Sj.array() - step_size_S*lambda_)
                                + (Sj.array() < -1.0*step_size_S*lambda_).cast<float>()
                                * (Sj.array() + step_size_S*lambda_)).matrix();
                            S_matrix_loader_.SetCol(col_id_client, Sj);

                            // get updated S_j
                            S_matrix_loader_.GetCol(col_id_client, Sj);
                        }
                        // update B
                        Xj_inc = Xj - petuum_table_cache * Sj;
                        petuum_update_cache.noalias() +=
                            step_size_B * 2.0 * Xj_inc * Sj.transpose();
                    }
                }
                // Update B_table
                for (int row_id = 0; row_id < dictionary_size_; ++row_id) {
                    petuum::UpdateBatch<float> B_update;
                    for (int col_id = 0; col_id < m; ++col_id) {
                        B_update.Update(col_id,
                                petuum_update_cache(col_id, row_id)
                                / minibatch_size_);
                    }
                    B_table.BatchInc(row_id, B_update);
                }
                petuum::PSTableGroup::Clock();
                // Update B_table to normalize l2-norm to C_
                std::vector<float> B_row_cache(m);
                for (int row_id = 0; row_id < dictionary_size_; ++row_id) {
                    const auto & row =
                        B_table.Get<petuum::DenseRow<float> >(row_id, &row_acc);
                    row.CopyToVector(&petuum_row_cache);
                    RegVec(petuum_row_cache, C_, B_row_cache);
                    petuum::UpdateBatch<float> B_update;
                    for (int col_id = 0; col_id < m; ++col_id) {
                        B_update.Update(col_id,
                                (-1.0 * petuum_row_cache[col_id] +
                                B_row_cache[col_id]) / num_clients_ /
                                num_worker_threads_);
                    }
                    B_table.BatchInc(row_id, B_update);
                }
                petuum::PSTableGroup::Clock();
            }
        }
        // Save results to disk
        petuum::PSTableGroup::GlobalBarrier();
        SaveResults(thread_id, B_table, loss_table);
        petuum::PSTableGroup::DeregisterThread();
    }

    SCEngine::~SCEngine() {
    }
} // namespace sparsecoding