#ifndef FASTKDESMOOTHING_H
#define FASTKDESMOOTHING_H

#include <vector>
#include <memory>
#include <algorithm>

#include "BackwardFilterTransition.h"
#include "BackwardFilter.h"

#include "../Particle.h"

#include "../../floorplan/v2/FloorplanHelper.h";
#include "../../grid/Grid.h";

#include "../filtering/resampling/ParticleFilterResampling.h"
#include "../filtering/estimation/ParticleFilterEstimation.h"
#include "../filtering/ParticleFilterEvaluation.h"
#include "../filtering/ParticleFilterInitializer.h"
#include "../filtering/ParticleFilterTransition.h"

#include "../sampling/ParticleTrajectorieSampler.h"

#include "../../math/boxkde/benchmark.h"
#include "../../math/boxkde/DataStructures.h"
#include "../../math/boxkde/Image2D.h"
#include "../../math/boxkde/BoxGaus.h"
#include "../../math/boxkde/Grid2D.h"

#include "../../Assertions.h"
namespace SMC {

    template <typename State, typename Control, typename Observation>
    class FastKDESmoothing : public BackwardFilter<State, Control, Observation>{

    private:

        /** all smoothed particles T -> 1*/
        std::vector<std::vector<Particle<State>>> backwardParticles;

        /** all estimations calculated */
        std::vector<State> estimatedStates;

        /** the estimation function to use */
        std::unique_ptr<ParticleFilterEstimation<State>> estimation;

        /** the transition function to use */
        std::unique_ptr<ParticleFilterTransition<State, Control>> transition;

        /** the resampler to use */
        std::unique_ptr<ParticleFilterResampling<State>> resampler;

        /** the sampler for drawing trajectories */
        std::unique_ptr<ParticleTrajectorieSampler<State>> sampler;

        /** the percentage-of-efficient-particles-threshold for resampling */
        double nEffThresholdPercent = 0.25;

        /** number of realizations to be calculated */
        int numParticles;

        /** is update called the first time? */
        bool firstFunctionCall;

        /** boundingBox for the boxKDE */
        BoundingBox<float> bb;

        /** histogram/grid holding the particles*/
        Grid2D<float> grid;

        /** bandwith for KDE */
        Point2 bandwith;

    public:

        /** ctor */
        FastKDESmoothing(int numParticles, const Floorplan::IndoorMap* map, const int gridsize_cm, const Point2 bandwith) {
            this->numParticles = numParticles;
            //backwardParticles.reserve(numParticles);
            this->firstFunctionCall = true;

            const Point3 maxBB = FloorplanHelper::getBBox(map).getMax() * 100.0;
            const Point3 minBB = FloorplanHelper::getBBox(map).getMin() * 100.0;
            this->bb = BoundingBox<float>(minBB.x, maxBB.x, minBB.y, maxBB.y);

            // Create histogram
            size_t nBinsX = static_cast<size_t>((maxBB.x - minBB.x) / gridsize_cm);
            size_t nBinsY = static_cast<size_t>((maxBB.y - minBB.y) / gridsize_cm);
            this->grid = Grid2D<float>(bb, nBinsX, nBinsY);

            this->bandwith = bandwith;
        }

        /** dtor */
        ~FastKDESmoothing() {
            backwardParticles.clear();
            estimatedStates.clear();
        }

        /** access to all backward / smoothed particles */
        const std::vector<std::vector<Particle<State>>>& getbackwardParticles() {
            return backwardParticles;
        }

        /** set the estimation method to use */
        void setEstimation(std::unique_ptr<ParticleFilterEstimation<State>> estimation) {
            Assert::isNotNull(estimation, "setEstimation() MUST not be called with a nullptr!");
            this->estimation = std::move(estimation);
        }

        /** set the transition method to use */
        void setTransition(std::unique_ptr<ParticleFilterTransition<State, Control>> transition) {
            Assert::isNotNull(transition, "setTransition() MUST not be called with a nullptr!");
            this->transition = std::move(transition);
        }

        /** set the transition method to use */
        void setTransition(std::unique_ptr<BackwardFilterTransition<State, Control>> transition) {
            Assert::doThrow("Do not use a backward transition for fast smoothing! Forward Transition");
            //TODO: two times setTransition is not the best solution
        }
        /** set the resampling method to use */
        void setResampling(std::unique_ptr<ParticleFilterResampling<State>> resampler) {
            Assert::isNotNull(resampler, "setResampling() MUST not be called with a nullptr!");
            this->resampler = std::move(resampler);
        }

        /** set the sampler method to use */
        void setSampler(std::unique_ptr<ParticleTrajectorieSampler<State>> sampler){
            Assert::isNotNull(sampler, "setSampler() MUST not be called with a nullptr!");
            this->sampler = std::move(sampler);
        }


        /** set the resampling threshold as the percentage of efficient particles */
        void setNEffThreshold(const double thresholdPercent) {
            this->nEffThresholdPercent = thresholdPercent;
        }

        /** get the used transition method */
        BackwardFilterTransition<State, Control>* getTransition() {
            return nullptr;
        }

        /**
         * @brief update
         * @param forwardHistory
         * @return
         */
        State update(ForwardFilterHistory<State, Control, Observation>& forwardHistory, int lag = 666) {

            // sanity checks (if enabled)
            Assert::isNotNull(transition, "transition MUST not be null! call setTransition() first!");
            Assert::isNotNull(estimation, "estimation MUST not be null! call setEstimation() first!");

            //init for backward filtering
            std::vector<Particle<State>> smoothedParticles;
            smoothedParticles.reserve(numParticles);
            firstFunctionCall = true;

            // if no lag is given, we have a fixed interval smoothing
            if(lag == 666){
                lag = forwardHistory.size() - 1;
            }

            //check if we have enough data for lag
            if(forwardHistory.size() <= lag){
                lag = forwardHistory.size() - 1;
            }

            //iterate through all forward filtering steps
            for(int i = 0; i <= lag; ++i){
                std::vector<Particle<State>> forwardParticlesForTransition_t1 = forwardHistory.getParticleSet(i);

                //storage for single trajectories / smoothed particles
                smoothedParticles.clear();

                // Choose \tilde x_T = x^(i)_T with probability w^(i)_T
                // Therefore sample independently from the categorical distribution of weights.
                if(firstFunctionCall){

                    smoothedParticles = sampler->drawTrajectorie(forwardParticlesForTransition_t1, numParticles);

                    firstFunctionCall = false;
                    backwardParticles.push_back(smoothedParticles);

                    State est = estimation->estimate(smoothedParticles);
                    estimatedStates.push_back(est);
                    continue;
                }

                // transition p(q_t+1* | q_t): so we are performing again a forward transition step.
                // we are doing this to track single particles between two timesteps! normaly, resampling would destroy
                // any identifier given to particles.
                // Node: at this point we can integrate future observation and control information for better smoothing
                Control controls = forwardHistory.getControl(i-1);
                Observation obs = forwardHistory.getObservation(i-1);
                transition->transition(forwardParticlesForTransition_t1, &controls);

                // KDE auf q_t+1 Samples = p(q_t+1 | o_1:T) - Smoothed samples from the future
                grid.clear();
                for (Particle<State> p : backwardParticles.back())
                {
                    grid.add(p.state.position.x_cm, p.state.position.y_cm, p.weight);
                }

                int nFilt = 3;
                float sigmaX = bandwith.x / grid.binSizeX;
                float sigmaY = bandwith.y / grid.binSizeY;
                BoxGaus<float> boxGaus;
                boxGaus.approxGaus(grid.image(), sigmaX, sigmaY, nFilt);

                /**
                  * Das hier wird nicht funktionieren!
                  * Wir machen eine Transition von t zu t+1 und nehmen dann diese position um von der kde aus t+1 ein
                  * Gewicht zu erhalten... aber das bringt ja nichts, diese Gewicht haben wir doch schon auf den partikeln in t+1
                  * die transition macht ja nicht 2x was anderes, sondern wieder genau das gleiche.
                  * was das smoothing gut macht, sind die Faltungen mit allen Partikeln (jeder mit jedem vergleichen).
                **/


                // Apply Position from Samples from q_t+1* into KDE of p(q_t+1 | o_1:T) to get p(q_t+1* | o_1:T)
                // Calculate new weight w(q_(t|T)) = w(q_t) * p(q_t+1* | o_1:T) * p(q_t+1* | q_t) * normalisation
                smoothedParticles = forwardHistory.getParticleSet(i);
                for(Particle<State> p : smoothedParticles){
                    p.weight = p.weight * grid.fetch(p.state.position.x_cm, p.state.position.y_cm);
                    //Assert::isNot0(p.weight, "smoothed particle has zero weight");
                }

                //normalization
                auto lambda = [](double current, const Particle<State>& a){return current + a.weight; };
                double weightSumSmoothed = std::accumulate(smoothedParticles.begin(), smoothedParticles.end(), 0.0, lambda);

                if(weightSumSmoothed != 0.0){

                    for (Particle<State> p : smoothedParticles){
                        p.weight /= weightSumSmoothed;
                    }

                    //check if normalization worked
                    double normWeightSum = std::accumulate(smoothedParticles.begin(), smoothedParticles.end(), 0.0, lambda);
                    Assert::isNear(normWeightSum, 1.0, 0.001, "Smoothed weights do not sum to 1");
                } else {
                    Assert::doThrow("Weight Sum of smoothed particles is zero!");
                }


                if(resampler)
                {
                    //TODO - does this even make sense?
                    Assert::doThrow("Warning - Resampling is not yet implemented!");
                }

                // push_back the smoothedParticles
                backwardParticles.push_back(smoothedParticles);

                // estimate the current state
                if(lag == (forwardHistory.size() - 1) ){ //fixed interval
                    State est = estimation->estimate(smoothedParticles);
                    estimatedStates.push_back(est);
                }
                else if (i == lag) { //fixed lag
                    State est = estimation->estimate(smoothedParticles);
                    estimatedStates.push_back(est);
                }

            }
            return estimatedStates.back();

        }

        std::vector<State> getEstimations(){
            return estimatedStates;
        }


    };

}
#endif // FASTKDESMOOTHING_H