#ifndef MIXINGSAMPLERDIVERGENCY_H
#define MIXINGSAMPLERDIVERGENCY_H

#include "MixingSampler.h"

#include <algorithm>
#include <random>
#include <eigen3/Eigen/Dense>

namespace SMC {

    /**
     * Using the direct sampling approach of Driessen and Boers in
     * "Efficient particle filter for jump Markov nonlinear systems".
     * as transition probability matrix for the markov chain we use
     * a divergence based on Jensen–Shannon divergence
     */
    template <typename State, typename Control, typename Observation> class MixingSamplerDivergency
        : public MixingSampler<State, Control, Observation> {

    /** random number generator */
    std::minstd_rand genNorm;
    std::minstd_rand genPart;

    /** copy of the modes with cumulative probabilities for easy drawing*/
    std::vector<ParticleFilterMixing<State, Control, Observation>> copyModes;

    public:

         void mixAndSample(std::vector<ParticleFilterMixing<State, Control, Observation>>& modes, Eigen::MatrixXd transitionProbabilityMatrix) override{

            genNorm.seed(std::chrono::system_clock::now().time_since_epoch().count());
            genPart.seed(std::chrono::system_clock::now().time_since_epoch().count() + 233);

            // set copyModes for later drawing
            copyModes = modes;

            // create cumulative particlesets for the copy
            // Note: in most cases, the particles are already resampled within the filtering stage
            // but in some cases they are not and therefore we need to draw cumulatively and not equally
            for(ParticleFilterMixing<State, Control, Observation>& copyFilter : copyModes){
                double cumWeight = 0;
                std::vector<Particle<State>> copyParticles;
                copyParticles = copyFilter.getParticles();

                for(int i = 0; i < copyParticles.size(); ++i){
                    cumWeight += copyParticles[i].weight;
                    copyParticles[i].weight = cumWeight;
                }

                copyFilter.setParticles(copyParticles);
            }

            // calculate the new predicted mode prob P(m_t|Z_t-1) and P(m_t-1 | m_t, Z_t-1)
            int m = 0; //this is m_t
            for(ParticleFilterMixing<State, Control, Observation>& focusedFilter : modes){

                // P(m_t|Z_t-1) = sum(P(m_t | m_t-1) P(m_t-1 | Z_t-1))
                int i = 0; //this are all possible m_t-1
                double predictedModeProbability = 0;
                for(ParticleFilterMixing<State, Control, Observation>& filter : modes){
                    predictedModeProbability += transitionProbabilityMatrix(m,i) * filter.getModePosteriorProbability();
                    ++i;
                }

                //Assert::isNotNull(predictedModeProbability, "predictedModeProbability is zero.. thats not possible!");
                focusedFilter.setPredictedModeProbability(predictedModeProbability);

                // cumulative sum of TransitionModeProbabilities for drawing modes from the perspective of ONE filter!
                // this means, the transition mode probabilities are calculated for each filter NEW!
                // calculate P(m_t-1 | m_t, Z_t-1) = P(m_t | m_t-1) * p(m_t-1 | Z_t-1) / P(m_t|Z_t-1)
                double cumTransitionModeProbability = 0;
                i = 0;
                for(ParticleFilterMixing<State, Control, Observation>& filter : modes){

                    double prob = (transitionProbabilityMatrix(m,i) * filter.getModePosteriorProbability()) / focusedFilter.getPredictedModeProbability();
                    filter.setTransitionModeProbability(prob);

                    cumTransitionModeProbability += prob;
                    copyModes[i].setTransitionModeProbability(cumTransitionModeProbability);

                    //std::cout << "Draw Mode Probability from mode " << m << i << " : " << prob << std::endl;
                    ++i;
                }

                // draw new modes and particles
                // Note: in most cases, the particles are already resampled within the filtering stage
                // but in some cases they are not and therefore we need to draw cumulatively and not equally

                // todo: make the particle size dynamic depending on the kld or something else
                // number of particles for this timestep
                int numParticles = focusedFilter.getParticles().size();

                std::vector<Particle<State>> newParticles;
                newParticles.resize(numParticles);
                double equalWeight = 1.0 / numParticles;

                // draw modes cumulative
                for(int k = 0; k < numParticles; ++k){

                    auto mode = drawMode(cumTransitionModeProbability);
                    newParticles[k] = drawParticle(mode);
                    newParticles[k].weight = equalWeight; //todo: why equal weight? i guess if the different filters have different representations of probability?
                }

                focusedFilter.setParticles(newParticles);

                //iter
                ++m;

            }
        }

    private:

         /** draw a mode depending upon the transition mode probabilities */
         ParticleFilterMixing<State, Control, Observation>& drawMode(const double cumTransitionModeProbabilities){

             // generate random values between [0:cumWeight]
             std::uniform_real_distribution<float> dist(0, cumTransitionModeProbabilities);

             // draw a random value between [0:cumWeight]
             const float rand = dist(genNorm);

             // search comparator (cumWeight is ordered -> use binary search)
             auto comp = [] (ParticleFilterMixing<State, Control, Observation>& filter, const float d) {return filter.getTransitionModeProbability() < d;};
             auto it = std::lower_bound(copyModes.begin(), copyModes.end(), rand, comp);
             return *it;

         }


        /** draw one particle according to its weight from the copy vector of a given mode */
        const Particle<State>& drawParticle(ParticleFilterMixing<State, Control, Observation>& filter) {

            //double weights = filter.getParticles().back().weight;

            // generate random values between [0:cumWeight]
            std::uniform_real_distribution<float> dist(0, filter.getParticles().back().weight);

            // draw a random value between [0:cumWeight]
            const float rand = dist(genPart);

            // search comparator (cumWeight is ordered -> use binary search)
            auto comp = [] (const Particle<State>& s, const float d) {return s.weight < d;};
            auto it = std::lower_bound(filter.getParticles().begin(), filter.getParticles().end(), rand, comp);
            return *it;

        }

    };


}


#endif // MIXINGSAMPLERDIVERGENCY_H