/*
 * CondensationBackwardFilter.h
 *
 *  Created on: Jun 23, 2015
 *      Author: Toni Fetzer
 */

#ifndef BACKWARDSIMULATION_H_
#define BACKWARDSIMULATION_H_


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

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

#include "../Particle.h"

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

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

#include "../../Assertions.h"

namespace SMC {

    /**
     * the main-class for the Backward Simulation Filter
     * running "backwards" in time, generates multiple backwards trajectories
     * (Realizations) by repeating the backward simulation M time.
     * it can be started at a random time T of any forward particle filter
     * [Monte Carlo smoothing for non-linear time series Godsill et al. '03]
     * @param State the (user-defined) state for each particle
     * @param numRealizations is the number of backward trajectories starting
     */
    template <typename State, typename Control, typename Observation>
    class BackwardSimulation : 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<BackwardFilterTransition<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 numRealizations;

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


    public:

        /** ctor */
        BackwardSimulation(int numRealizations) {
            this->numRealizations = numRealizations;
            backwardParticles.reserve(numRealizations);
            firstFunctionCall = true;
        }

        /** dtor */
        ~BackwardSimulation() {
            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<BackwardFilterTransition<State, Control>> transition) {
            Assert::isNotNull(transition, "setTransition() MUST not be called with a nullptr!");
            this->transition = std::move(transition);
        }

        /** 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 this->transition.get();
        }

        /**
         * @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(numRealizations);
            firstFunctionCall = true;

            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>> forwardParticles = 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(forwardParticles, numRealizations);

                    firstFunctionCall = false;
                    backwardParticles.push_back(smoothedParticles);

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

                // compute weights using the transition model
                // transitionWeigths[numRealizations][numParticles]
                std::vector<std::vector<double>> transitionWeights = transition->transition(forwardParticles, backwardParticles.back());

                //get the next trajectorie for a realisation
                for(int j = 0; j < numRealizations; ++j){

                    //vector for the current smoothedWeights at time t
                    std::vector<Particle<State>> smoothedWeights;
                    smoothedWeights.resize(forwardParticles.size());
                    smoothedWeights = forwardParticles;

                    //check if all transitionWeights are zero
                    double weightSumTransition = std::accumulate(transitionWeights[j].begin(), transitionWeights[j].end(), 0.0);
                    Assert::isNot0(weightSumTransition, "all transition weights for smoothing are zero");

                    int k = 0;
                    for (auto& w : transitionWeights.at(j)) {

                        // multiply the weight of the particles at time t and normalize
                        smoothedWeights.at(k).weight = (smoothedWeights.at(k).weight * w);
                        if(smoothedWeights.at(k).weight != smoothedWeights.at(k).weight) {throw "detected NaN";}

                        // iter
                        ++k;
                    }

                    //get the sum of all weights
                    auto lambda = [](double current, const Particle<State>& a){return current + a.weight; };
                    double weightSumSmoothed = std::accumulate(smoothedWeights.begin(), smoothedWeights.end(), 0.0, lambda);

                    //normalize the weights
                    if(weightSumSmoothed != 0.0){
                        for (int l = 0; l < smoothedWeights.size(); ++l){
                            smoothedWeights.at(l).weight /= weightSumSmoothed;
                        }

                        //check if normalization worked
                        double normWeightSum = std::accumulate(smoothedWeights.begin(), smoothedWeights.end(), 0.0, lambda);
                        Assert::isNear(normWeightSum, 1.0, 0.001, "Smoothed weights do not sum to 1");
                    }


                    //draw the next trajectorie at time t for a realization and save them
                    smoothedParticles.push_back(sampler->drawSingleParticle(smoothedWeights));

                    //throw if weight of smoothedParticle is zero
                    //in practice this is possible, if a particle is completely separated from the rest and is therefore
                    //weighted zero or very very low.
                    Assert::isNot0(smoothedParticles.back().weight, "smoothed particle has zero weight");
                }


                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 lag
                    State est = estimation->estimate(smoothedParticles);
                    estimatedStates.push_back(est);
                }
                else if (i == lag) { //fixed interval
                    State est = estimation->estimate(smoothedParticles);
                    estimatedStates.push_back(est);
                }

            }

            //return the calculated estimations
            // TODO: Wir interessieren uns beim fixed-lag smoothing immer nur für die letzte estimation und den letzten satz gesmoothet particle da wir ja weiter vorwärts in der Zeit gehen
            // und pro zeitschritt ein neues particle set hinzu kommt. also wenn lag = 3, dann smoothen wir t - 3 und sollten auch nur die estimation von t-3 und das particle set von t-3 abspeichern
            //
            // beim interval smoothing dagegen interessieren uns alle, da ja keine neuen future informationen kommen und wir einfach sequentiell zurück in die zeit wandern.
            //
            // lösungsvorschlag.
            // - observer-pattern? immer wenn eine neue estimation und ein neues particle set kommt, sende. (wird nix bringen, da keine unterschiedlichen threads?)
            // - wie vorher machen, also pro update eine estimation aber jedem update einfach alle observations und alle controls mitgeben?!?!?
            // - ich gebe zusätzlich den lag mit an. dann kann die forwardfilterhistory auch ständig alles halten. dann gibt es halt keine ständigen updates, sondern man muss die eine
            // berechnung abwarten. eventl. eine art ladebalken hinzufügen. (30 von 120 timestamps done) (ich glaube das ist die beste blackboxigste version) man kann den lag natürlich auch beim
            // init des backwardsimulation objects mit übergeben. ABER: damit könntem an kein dynamic-lag smoothing mehr machen. also lieber variable lassen :). durch den lag wissen wir einfach was wir
            // genau in estimatedStates und backwardParticles speichern müssen ohne über das problem oben zu stoßen. haben halt keine ständigen updates. observer-pattern hier nur bei mehrere threads,
            // das wäre jetzt aber overkill und deshalb einfach ladebalken :):):):)
            // - oder man macht einfach zwei update funktionen mit den beiden möglichkeiten. halte ich aber für nen dummen kompromiss. )

            // würde es sinn machen, die estimations auch mit zu speichern?

            return estimatedStates.back();

        }

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


    };


}


#endif /* BACKWARDSIMULATION_H_ */