#pragma once

#include "../include/openpgl/config.h"

#include "../openpgl_common.h"

#include "field/ISurfaceVolumeField.h"
#include "field/SurfaceVolumeField.h"

#include "directional/vmm/ParallaxAwareVonMisesFisherMixture.h"
#include "directional/vmm/AdaptiveSplitandMergeFactory.h"
#include "directional/vmm/VMMSurfaceSamplingDistribution.h"
#include "directional/vmm/VMMVolumeSamplingDistribution.h"
#include "directional/vmm/VMMPhaseFunctions.h"
#include "directional/dqt/DQT.h"
#include "directional/dqt/DQTFactory.h"
#include "directional/dqt/DQTSurfaceSamplingDistribution.h"
#include "directional/dqt/DQTVolumeSamplingDistribution.h"
#include "directional/dqt/SphereToSquare.h"

#include "spatial/kdtree/KDTreeBuilder.h"

#include "tbb/tbb.h"

#define OPENPGL_TASK_CONTROL

namespace openpgl {

#ifdef OPENPGL_TASK_CONTROL
#if TBB_INTERFACE_VERSION >= 11005
  static tbb::global_control* g_opgl_tbb_thread_control = nullptr;
#else
  static tbb::task_scheduler_init g_opgl_tbb_threads(tbb::task_scheduler_init::deferred);
#endif
#endif

struct IDevice {
    virtual ~IDevice() {};
    virtual ISurfaceVolumeField* newField(PGLFieldArguments args) const = 0;
    virtual ISurfaceVolumeField* newFieldFromFile(const std::string fieldFileName) const = 0;
};

template<int VecSize>
struct Device: public IDevice {

private:
#ifdef OPENPGL_TASK_CONTROL
    size_t m_numThreads {0};
#endif
public:
    Device(size_t numThreads = 0)
    {
#ifdef OPENPGL_TASK_CONTROL
        if (numThreads == 0) 
        {
#if TBB_INTERFACE_VERSION >= 9100
            m_numThreads = tbb::this_task_arena::max_concurrency();
#else
            m_numThreads = tbb::task_scheduler_init::default_num_threads();
#endif
        } else {
#if TBB_INTERFACE_VERSION >= 9100
            m_numThreads = std::min(numThreads, (size_t)tbb::this_task_arena::max_concurrency());
#else
            m_numThreads = std::min(numThreads, (size_t)tbb::task_scheduler_init::default_num_threads());
#endif
            
        }
#if TBB_INTERFACE_VERSION >= 11005
        g_opgl_tbb_thread_control = new tbb::global_control(tbb::global_control::max_allowed_parallelism, m_numThreads);
#else
        g_opgl_tbb_threads.initialize(int(m_numThreads));
#endif
#endif
        VMMSingleLobeHenyeyGreensteinOracle::init();
    }

    ~Device() override
    {
#ifdef OPENPGL_TASK_CONTROL
#if TBB_INTERFACE_VERSION >= 11005
        delete g_opgl_tbb_thread_control;
        g_opgl_tbb_thread_control = nullptr;
#else
        g_opgl_tbb_threads.terminate();
#endif
#endif
    }

    ISurfaceVolumeField* newField(PGLFieldArguments args) const override {
        ISurfaceVolumeField* gField;

        if (args.spatialStructureType == PGL_SPATIAL_STRUCTURE_KDTREE &&
            args.directionalDistributionType ==  PGL_DIRECTIONAL_DISTRIBUTION_PARALLAX_AWARE_VMM )
        {
            using DirectionalDistributionFactory = AdaptiveSplitAndMergeFactory<ParallaxAwareVonMisesFisherMixture<VecSize, 32, true>>;
            using GuidingField = SurfaceVolumeField<VecSize, DirectionalDistributionFactory, KDTreePartitionBuilder, VMMSurfaceSamplingDistribution<typename DirectionalDistributionFactory::Distribution, true>, VMMVolumeSamplingDistribution<typename DirectionalDistributionFactory::Distribution, true>>;

            typename GuidingField::Settings gFieldSettings;
            gFieldSettings.settings.decayOnSpatialSplit   = 0.25f;
            gFieldSettings.settings.deterministic         = args.deterministic;
            gFieldSettings.debugSettings.fitRegions       = args.debugArguments.fitRegions;

            PGLKDTreeArguments *spatialSturctureArguments = (PGLKDTreeArguments*)args.spatialSturctureArguments;
            gFieldSettings.settings.useStochasticNNLookUp = spatialSturctureArguments->knnLookup;
            gFieldSettings.settings.spatialSubdivBuilderSettings.minSamples = spatialSturctureArguments->minSamples;
            gFieldSettings.settings.spatialSubdivBuilderSettings.maxSamples = spatialSturctureArguments->maxSamples;
            gFieldSettings.settings.spatialSubdivBuilderSettings.maxDepth   = spatialSturctureArguments->maxDepth;
            delete spatialSturctureArguments;

            PGLVMMFactoryArguments *directionalDistributionArguments = (PGLVMMFactoryArguments*)args.directionalDistributionArguments;

            gFieldSettings.distributionFactorySettings.weightedEMCfg.initK = directionalDistributionArguments->initK;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.initKappa = directionalDistributionArguments->initKappa;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.maxK = directionalDistributionArguments->maxK;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.maxEMIterrations = directionalDistributionArguments->maxEMIterrations;

            gFieldSettings.distributionFactorySettings.weightedEMCfg.maxKappa = directionalDistributionArguments->maxKappa;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.maxMeanCosine = openpgl::KappaToMeanCosine<float>(gFieldSettings.distributionFactorySettings.weightedEMCfg.maxKappa);
            gFieldSettings.distributionFactorySettings.weightedEMCfg.convergenceThreshold = directionalDistributionArguments->convergenceThreshold;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.weightPrior = directionalDistributionArguments->weightPrior;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.meanCosinePriorStrength = directionalDistributionArguments->meanCosinePriorStrength;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.meanCosinePrior = directionalDistributionArguments->meanCosinePrior;

            gFieldSettings.distributionFactorySettings.splittingThreshold = directionalDistributionArguments->splittingThreshold;
            gFieldSettings.distributionFactorySettings.mergingThreshold = directionalDistributionArguments->mergingThreshold;


            gFieldSettings.distributionFactorySettings.partialReFit = directionalDistributionArguments->partialReFit;
            gFieldSettings.distributionFactorySettings.maxSplitItr = directionalDistributionArguments->maxSplitItr;

            gFieldSettings.distributionFactorySettings.useSplitAndMerge = directionalDistributionArguments->useSplitAndMerge;
            gFieldSettings.distributionFactorySettings.minSamplesForSplitting = directionalDistributionArguments->minSamplesForSplitting;
            gFieldSettings.distributionFactorySettings.minSamplesForPartialRefitting = directionalDistributionArguments->minSamplesForPartialRefitting;
            gFieldSettings.distributionFactorySettings.minSamplesForMerging = directionalDistributionArguments->minSamplesForMerging;
            delete directionalDistributionArguments;

            gField = new GuidingField(gFieldSettings);
        } else if (args.spatialStructureType == PGL_SPATIAL_STRUCTURE_KDTREE &&
            args.directionalDistributionType ==  PGL_DIRECTIONAL_DISTRIBUTION_VMM )
        {
            using DirectionalDistributionFactory = AdaptiveSplitAndMergeFactory<ParallaxAwareVonMisesFisherMixture<VecSize, 32, false>>;
            using GuidingField = SurfaceVolumeField<VecSize, DirectionalDistributionFactory, KDTreePartitionBuilder, VMMSurfaceSamplingDistribution<typename DirectionalDistributionFactory::Distribution, false>, VMMVolumeSamplingDistribution<typename DirectionalDistributionFactory::Distribution, false>>;

            typename GuidingField::Settings gFieldSettings;
            gFieldSettings.settings.decayOnSpatialSplit   = 0.25f;
            gFieldSettings.settings.deterministic         = args.deterministic;
            gFieldSettings.debugSettings.fitRegions       = args.debugArguments.fitRegions;

            PGLKDTreeArguments *spatialSturctureArguments = (PGLKDTreeArguments*)args.spatialSturctureArguments;
            gFieldSettings.settings.useStochasticNNLookUp = spatialSturctureArguments->knnLookup;
            gFieldSettings.settings.spatialSubdivBuilderSettings.minSamples = spatialSturctureArguments->minSamples;
            gFieldSettings.settings.spatialSubdivBuilderSettings.maxSamples = spatialSturctureArguments->maxSamples;
            gFieldSettings.settings.spatialSubdivBuilderSettings.maxDepth   = spatialSturctureArguments->maxDepth;
            delete spatialSturctureArguments;

            PGLVMMFactoryArguments *directionalDistributionArguments = (PGLVMMFactoryArguments*)args.directionalDistributionArguments;

            gFieldSettings.distributionFactorySettings.weightedEMCfg.initK = directionalDistributionArguments->initK;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.initKappa = directionalDistributionArguments->initKappa;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.maxK = directionalDistributionArguments->maxK;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.maxEMIterrations = directionalDistributionArguments->maxEMIterrations;

            gFieldSettings.distributionFactorySettings.weightedEMCfg.maxKappa = directionalDistributionArguments->maxKappa;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.maxMeanCosine = openpgl::KappaToMeanCosine<float>(gFieldSettings.distributionFactorySettings.weightedEMCfg.maxKappa);
            gFieldSettings.distributionFactorySettings.weightedEMCfg.convergenceThreshold = directionalDistributionArguments->convergenceThreshold;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.weightPrior = directionalDistributionArguments->weightPrior;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.meanCosinePriorStrength = directionalDistributionArguments->meanCosinePriorStrength;
            gFieldSettings.distributionFactorySettings.weightedEMCfg.meanCosinePrior = directionalDistributionArguments->meanCosinePrior;

            gFieldSettings.distributionFactorySettings.splittingThreshold = directionalDistributionArguments->splittingThreshold;
            gFieldSettings.distributionFactorySettings.mergingThreshold = directionalDistributionArguments->mergingThreshold;


            gFieldSettings.distributionFactorySettings.partialReFit = directionalDistributionArguments->partialReFit;
            gFieldSettings.distributionFactorySettings.maxSplitItr = directionalDistributionArguments->maxSplitItr;

            gFieldSettings.distributionFactorySettings.useSplitAndMerge = directionalDistributionArguments->useSplitAndMerge;
            gFieldSettings.distributionFactorySettings.minSamplesForSplitting = directionalDistributionArguments->minSamplesForSplitting;
            gFieldSettings.distributionFactorySettings.minSamplesForPartialRefitting = directionalDistributionArguments->minSamplesForPartialRefitting;
            gFieldSettings.distributionFactorySettings.minSamplesForMerging = directionalDistributionArguments->minSamplesForMerging;
            delete directionalDistributionArguments;

            gField = new GuidingField(gFieldSettings);
        } else if (args.spatialStructureType == PGL_SPATIAL_STRUCTURE_KDTREE &&
                   args.directionalDistributionType ==  PGL_DIRECTIONAL_DISTRIBUTION_QUADTREE )
        {
            using DirectionalDistributionFactory = DirectionalQuadtreeFactory<DirectionalQuadtree<SphereToSquareCylindrical>>;
            using GuidingField = SurfaceVolumeField<VecSize, DirectionalDistributionFactory, KDTreePartitionBuilder, DQTSurfaceSamplingDistribution<DirectionalDistributionFactory::Distribution>, DQTVolumeSamplingDistribution<DirectionalDistributionFactory::Distribution>>;

            typename GuidingField::Settings gFieldSettings;
            gFieldSettings.settings.decayOnSpatialSplit   = 0.25f;
            gFieldSettings.settings.deterministic         = args.deterministic;
            gFieldSettings.debugSettings.fitRegions       = args.debugArguments.fitRegions;

            PGLKDTreeArguments *spatialSturctureArguments = (PGLKDTreeArguments*)args.spatialSturctureArguments;
            gFieldSettings.settings.useStochasticNNLookUp = spatialSturctureArguments->knnLookup;
            gFieldSettings.settings.spatialSubdivBuilderSettings.minSamples = spatialSturctureArguments->minSamples;
            gFieldSettings.settings.spatialSubdivBuilderSettings.maxSamples = spatialSturctureArguments->maxSamples;
            gFieldSettings.settings.spatialSubdivBuilderSettings.maxDepth   = spatialSturctureArguments->maxDepth;

            PGLDQTFactoryArguments *directionalDistributionArguments = (PGLDQTFactoryArguments*)args.directionalDistributionArguments;
            gFieldSettings.distributionFactorySettings.leafEstimator = (LeafEstimator)directionalDistributionArguments->leafEstimator;
            gFieldSettings.distributionFactorySettings.splitMetric = (SplitMetric)directionalDistributionArguments->splitMetric;
            gFieldSettings.distributionFactorySettings.splitThreshold = directionalDistributionArguments->splitThreshold;
            gFieldSettings.distributionFactorySettings.footprintFactor = directionalDistributionArguments->footprintFactor;
            gFieldSettings.distributionFactorySettings.maxLevels = directionalDistributionArguments->maxLevels;

            gField = new GuidingField(gFieldSettings);
        }else {
            throw std::runtime_error("error: unrecognized field type");
        }

        return gField;
    }

    ISurfaceVolumeField* newFieldFromFile(const std::string fieldFileName) const override {
        std::filebuf fb;
        fb.open (fieldFileName, std::ios::in | std::ios::binary);
        if (!fb.is_open()) throw std::runtime_error("error: couldn't open file");
        std::istream is(&fb);

        auto size = strlen(FIELD_FILE_HEADER_STRING) + 1;
        OPENPGL_ASSERT(size <= 256);
        char buf[256];
        is.read(&buf[0], size);
        if (!is) throw std::runtime_error("error: invalid file header");
#ifdef OPENPGL_STRICT_IO_VERSION_CHECKING        
        for (auto i = 0; i < size; i++)
        {
            if (buf[i] != FIELD_FILE_HEADER_STRING[i])
                throw std::runtime_error("error: invalid file header");
        }
#endif        
        PGL_SPATIAL_STRUCTURE_TYPE spatialStructureType;
        is.read(reinterpret_cast<char*>(&spatialStructureType), sizeof(spatialStructureType));
        PGL_DIRECTIONAL_DISTRIBUTION_TYPE directionalDistributionType;
        is.read(reinterpret_cast<char*>(&directionalDistributionType), sizeof(directionalDistributionType));

        ISurfaceVolumeField* gField;

        if (spatialStructureType == PGL_SPATIAL_STRUCTURE_KDTREE && directionalDistributionType == PGL_DIRECTIONAL_DISTRIBUTION_PARALLAX_AWARE_VMM)
        {
            using DirectionalDistributionFactory = AdaptiveSplitAndMergeFactory<ParallaxAwareVonMisesFisherMixture<VecSize, 32, true>>;
            using GuidingField = SurfaceVolumeField<VecSize, DirectionalDistributionFactory, KDTreePartitionBuilder, VMMSurfaceSamplingDistribution<typename DirectionalDistributionFactory::Distribution, true>, VMMVolumeSamplingDistribution<typename DirectionalDistributionFactory::Distribution, true>>;

            gField = (ISurfaceVolumeField *)new GuidingField();
        } else if (spatialStructureType == PGL_SPATIAL_STRUCTURE_KDTREE && directionalDistributionType == PGL_DIRECTIONAL_DISTRIBUTION_VMM)
        {
            using DirectionalDistributionFactory = AdaptiveSplitAndMergeFactory<ParallaxAwareVonMisesFisherMixture<VecSize, 32, false>>;
            using GuidingField = SurfaceVolumeField<VecSize, DirectionalDistributionFactory, KDTreePartitionBuilder, VMMSurfaceSamplingDistribution<typename DirectionalDistributionFactory::Distribution, false>, VMMVolumeSamplingDistribution<typename DirectionalDistributionFactory::Distribution, false>>;

            gField = (ISurfaceVolumeField *)new GuidingField();
        } else if (spatialStructureType == PGL_SPATIAL_STRUCTURE_KDTREE &&
                   directionalDistributionType == PGL_DIRECTIONAL_DISTRIBUTION_QUADTREE)
        {
            using DirectionalDistributionFactory = DirectionalQuadtreeFactory<DirectionalQuadtree<SphereToSquareCylindrical>>;
            using GuidingField = SurfaceVolumeField<VecSize, DirectionalDistributionFactory, KDTreePartitionBuilder, DQTSurfaceSamplingDistribution<DirectionalDistributionFactory::Distribution>, DQTVolumeSamplingDistribution<DirectionalDistributionFactory::Distribution>>;

            gField = (ISurfaceVolumeField *)new GuidingField();
        } else {
            fb.close();
            throw std::runtime_error("error: unrecognized field type");
        }

        gField->deserialize(is);

        fb.close();

        return gField;
    }
};

#ifdef OPENPGL_DEVICE_TYPE_CPU_4
IDevice* newDeviceCPU4(size_t numThreads = 0);
#endif
#ifdef OPENPGL_DEVICE_TYPE_CPU_8
IDevice* newDeviceCPU8(size_t numThreads = 0);
#endif
#ifdef OPENPGL_DEVICE_TYPE_CPU_16
IDevice* newDeviceCPU16(size_t numThreads = 0);
#endif

}

#undef FILE_HEADER_STRING
