#include <vector>
#include <algorithm>
#include <iostream>
#include "var_alea.hpp"
#include <gsl/gsl_cdf.h>
#include <functional>

using namespace std;

vector<double> quantile_norm(int n, double sigma);


struct gaussian_truncated : var_alea<double>
{
    gaussian_truncated(double a, double b, double mu = 0, double sigma = 1)
        :a(a), b(b), V(gsl_cdf_ugaussian_P(a), gsl_cdf_ugaussian_P(b)), mu(mu), sigma(sigma) {};

    double operator()() {
        double v = V();
        return mu + gsl_cdf_gaussian_Pinv(v,sigma);
    }

    private:
        double a, b;
        uniform V;
        double mu;
        double sigma;
    };

struct multi_gaussian_truncated : public var_alea<std::vector<double> >
{
    multi_gaussian_truncated(double a, double b, const std::vector<double> u)
        :a(a), b(b), V(gsl_cdf_ugaussian_P(a), gsl_cdf_ugaussian_P(b)), G(0,1), u(u), d(u.size()) {};

    std::vector<double> operator()() {
        double v = V();
        double Z = gsl_cdf_gaussian_Pinv(v,1);
        std::vector<double> Y(d);
        for(int i=0; i<d; i++){
            Y[i] = G();
        }
        double scal = 0;
        for(int j=0; j<d; j++){
            scal += Y[j]*u[j];
        }
        std::vector<double> X(d);
        for(int i=0; i<d; i++){
            X[i] = u[i]*Z + Y[i] - u[i]*scal;
        }
        return X;
    }

    private:
        double a, b;
        uniform V;
        gaussian G;
        std::vector<double> u;
        int d;
    };

template <typename L>
struct stratified_sampling {
    stratified_sampling(vector<double> p, vector<L> X)
        :p(p), X(X), mean(p.size(), 0), sigma2(p.size(), 0), I(p.size()){};
    void draw(int N);
    vector<double> get_mean() const;
    vector<double> get_var() const;
    void print_mean() const;
    void print_sigma() const;
    pair<double,double> estimator() const;
private:
    void update(int N);
    vector<double> p;
    vector<L> X;
    vector<int> M;
    vector<int> cumM;
    vector<double> mean;
    vector<double> sigma2;
    const int I;
};

//actualisation du nombre de tirages à faire par strates
template <typename L>
void stratified_sampling<L>::update(int Nk) {
    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-1; i++){
        M[i] += floor(current+m[i]) - floor(current);
        current += m[i];
    }
    M[I-1]+=Nk-I-floor(current);
}

template <typename L>
void stratified_sampling<L>::draw(int N) {
		update(N);
    double m, s, oldmean;
    for(int i=0;i<I;i++){
        m=0;
        s=0;
        for(int j=0;j<M[i];j++){
            double temp = X[i]();
            m=m+temp;
            s=s+temp*temp;
        }
        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 L>
vector<double> stratified_sampling<L>::get_mean() const {
    return mean;
};

template <typename L>
vector<double> stratified_sampling<L>::get_var() const {
    return sigma2;
};

template <typename L>
void stratified_sampling<L>::print_mean() const {
	cout<<"les espérances :"<<endl;
	    for(int i=0;i<I;i++){
        cout<<mean[i]<<"\t";
			}
			cout<<endl;
};

template <typename L>
void stratified_sampling<L>::print_sigma() const {
	cout<<"les écarts types :"<<endl;
	    for(int i=0;i<I;i++){
        cout<<sqrt(sigma2[i])<<"\t";
			}
			cout<<endl;
};

template <typename L>
pair<double,double> stratified_sampling<L>::estimator() const {
	double est_mean = 0;
	double est_std = 0;
	for (int i=0; i<I; i++) {
		est_mean += mean[i]*p[i];
		est_std += sqrt(sigma2[i])*p[i];
	}
	return {est_mean, est_std};
}

template <typename Fct>
struct exponential_tilt : public std::function<double(std::vector<double>)>
{
    exponential_tilt(std::vector<double> mu, Fct f) : mu(mu), f(f){
        norm_mu = 0;
        for(unsigned int i=0; i<mu.size(); i++) {
            norm_mu += mu[i]*mu[i];
        }
    };

    double operator()(std::vector<double> X) {
        std::vector<double> Y(X.size());
        double scal = 0;
        for (unsigned int i=0; i<X.size(); i++){
            Y[i] = X[i] + mu[i];
            scal+=X[i]*mu[i];
        }
        return f(Y) * exp(-scal-0.5*norm_mu);
    };

    private:
        std::vector<double> mu;
        Fct f;
        double norm_mu;
    };