#include "RooGaussian.h"
#include "RooExponential.h"
#include "RooRealVar.h"
#include "TCanvas.h"
#include "RooDataSet.h"
#include "RooPlot.h"
#include <map>

void testNormalisation() {

    RooRealVar x("x", "x", -1e-3, 1e-3);

    // Define the model
    RooRealVar mean1("mean1", "mean1", 1e-4, -5e-4, 5e-4);
    RooRealVar sigma1("sigma1", "sigma1", 1e-4, 1e-5, 2e-4);
    RooGaussian gauss1("gauss1", "gauss1", x, mean1, sigma1);

    RooRealVar mean2("mean2", "mean2", -2e-4, -5e-4, 5e-4);
    RooRealVar sigma2("sigma2", "sigma2", 1.5e-4, 1e-5, 2e-4);
    RooGaussian gauss2("gauss2", "gauss2", x, mean2, sigma2);

    RooRealVar slope("slope", "slope", -4e-6, -1e-4, 0.);
    RooExponential expo("expo", "expo", x, slope);

    RooRealVar yieldGauss1("yieldGauss1", "yieldGauss1", 11000, 5000, 20000);
    RooRealVar yieldGauss2("yieldGauss2", "yieldGauss2", 14000, 5000, 20000);
    RooRealVar yieldExpo("yieldExpo", "yieldExpo", 22000, 10000, 50000);

    RooAddPdf model("model", "model", RooArgList(gauss1, gauss2, expo), RooArgList(yieldGauss1, yieldGauss2, yieldExpo));

    // Define the multi-variate constraint
    RooConstVar constrMean1("constrMean1", "constrMean1", 1e-4);
    RooConstVar constrMean2("constrMean2", "constrMean2", -2e-4);
    RooConstVar constrSigma1("constrSigma1", "constrSigma1", 1e-4);
    RooConstVar constrSigma2("constrSigma2", "constrSigma2", 1.5e-4);

    vector <double> errors(4);

    TRandom3 randomNumberGenerator;
    errors[0] = 2e-5 + 1e-5 * randomNumberGenerator.Uniform(-1, 1);
    errors[1] = 2e-5 + 1e-5 * randomNumberGenerator.Uniform(-1, 1);
    errors[2] = 2e-5 + 1e-5 * randomNumberGenerator.Uniform(-1, 1);
    errors[3] = 2e-5 + 1e-5 * randomNumberGenerator.Uniform(-1, 1);

    TMatrixDSym covarianceMatrix(4);
    for (int i = 1; i < 4; i++) {
        for (int j = i; j < 4; j++) {
//            double covariance = errors[i] * errors[j] * 0.16 * randomNumberGenerator.Uniform(-1, 1);
            covarianceMatrix(i,j) = 0;
            covarianceMatrix(j,i) = 0;
        }
    }

    for (int i = 0; i < 4; i++) {
        covarianceMatrix(i, i) = errors[i] * errors[i];
    }

    RooArgSet constrainedVariablesSet(mean1, mean2, sigma1, sigma2);
    RooArgList constrainedVariables(mean1, mean2, sigma1, sigma2);
    RooArgList constrainedMeans(constrMean1, constrMean2, constrSigma1, constrSigma2);

    RooMultiVarGaussian constraintPDF("constraintPDF", "constraintPDF", constrainedVariables, constrainedMeans, covarianceMatrix);
    RooArgSet multiVarGaussianObservables(mean1, mean2, sigma1, sigma2);
    std::cout << "With normalisation \t :" << constraintPDF.getVal(multiVarGaussianObservables) << std::endl;
    std::cout << "Without normalisation \t :" << constraintPDF.getVal() << std::endl;

    std::cout << "Test log with normalisation \t :" << constraintPDF.getLogVal(&multiVarGaussianObservables) << std::endl;
    std::cout << "Test log without normalisation \t :" << constraintPDF.getLogVal() << std::endl;
}