#include <vector>
#include <algorithm>
#include <iostream>
#include <gsl/gsl_rng.h>
#include "rtnorm.hpp"

using namespace std;

template <typename T>
struct var_alea {
    typedef T result_type;
    var_alea() : value(0) {};
    var_alea(T value) : value(value) {};
    virtual ~var_alea() {};
    virtual T operator()() = 0;
    T current() const { return value; };
protected:
    T value;
};

typedef var_alea<double> var_alea_real;

struct gaussian_truncated : public var_alea_real
{
    gaussian_truncated(double a, double b, double mean=0, double sigma2=1, int seed=0)
        :a(a), b(b), mean(mean), sigma2(sigma2), seed(seed) {
        const gsl_rng_type* type = gsl_rng_default;
        gen = gsl_rng_alloc(type);
        gsl_rng_set(gen, seed);
    };
    gaussian_truncated(gaussian_truncated const &other)
        :a(other.a),b(other.b), mean(other.mean),sigma2(other.sigma2){
        gen = gsl_rng_clone(other.gen);
    };

    double operator()() {
        pair<double, double> p = rtnorm(gen, a, b, mean, sigma2);
        return value = p.first;
    };

    ~gaussian_truncated() { gsl_rng_free(gen); }
private:
    int seed;
    double mean, sigma2, a, b;
    gsl_rng *gen;
};

template <typename Gen>
struct stratified_sampling {
    stratified_sampling(vector<double> p, vector<Gen> gen)
        :p(p), gen(gen), mean(p.size(), 0), sigma2(p.size(), 0){};
    void update(int N);
    void draw();
    vector<double> get_mean();
    //double estimator();
private:
    vector<double> p;
    vector<int> M;
    vector<int> cumM;
    vector<double> mean;
    vector<double> sigma2;
    vector<Gen> gen;

};

//actualisation du nombre de tirages à faire par strates
template <typename Gen>
void stratified_sampling<Gen>::update(int Nk) {
    int I = p.size();
    bool first_step = M.empty();
    //reinitialistation du vecteur M du nombre de tirages par strates
    if (first_step) {
        M.resize(I,1);
        cumM.resize(I,0);
    }
    else {
        for(int i=0; i<I; i++){
            cumM[i]=cumM[i]+M[i];
            M[i]=1;
        }
    }

    std::vector<double> m(I, 0); //le vecteur des m_i idéals

    if (first_step) {
        for (int i=0; i<I; i++) {
            m[i] = (Nk-I)*p[i];
        }
    }
    else {
				//On remplit un vecter des écarts types à parti de notre vecteur de variance
				std::vector<double> sigma(p.size(),0);
				for (int i=0; i < I; i++) {
					sigma[i]=sqrt(sigma2[i]);
				}
        double scal = std::inner_product(p.begin(), p.end(), sigma.begin(), (double) 0);
        for (int i=0; i < I; i++) {
            m[i] = (Nk-I)*p[i]*sigma[i]/scal;
            std::cout<<m[i]<<std::endl;
        }
    }
    M[0]+=floor(m[0]);
    double current = m[0];
    for (int i=1; i<I; i++){
        M[i] += floor(current+m[i]) - floor(current);
        current += m[i];
        cout<<M[i]<<"\t";
    }
    cout<<endl;
}

template <typename Gen>
void stratified_sampling<Gen>::draw() {
    int I = p.size();
    double m, s, oldmean;
    for(int i=0;i<I;i++){
        m=0;
        s=0;
        for(int j=0;j<M[i];j++){
            m=m+gen[i]();
            s=s+gen[i].current()*gen[i].current();
        }
        oldmean=mean[i];
        mean[i]=(mean[i]*cumM[i]+m)/(cumM[i]+M[i]);
        sigma2[i]=((sigma2[i]+oldmean*oldmean)*cumM[i] + s)/(cumM[i]+M[i]) - mean[i]*mean[i];
    }
};

template <typename Gen>
vector<double> stratified_sampling<Gen>::get_mean() {
    return mean;
};