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

RooDataSet * GetExponentialDataset(RooRealVar& x, double slope, int nEvents) {
    RooRealVar genSlope("genSlope", "genSlope", slope);
    RooExponential genExpo("genExpo", "genExpo", x, genSlope);
    RooArgSet genVariables(x);
    auto * dataset = genExpo.generate(genVariables, nEvents);
    return dataset;
}

RooDataSet * GetGaussianDataset(RooRealVar& x, double mean, double sigma, int nEvents) {
    RooRealVar genMean("genMean", "genMean", mean);
    RooRealVar genSigma("genSigma", "genSigma", sigma);
    RooGaussian genGauss("genGauss", "genGauss", x, genMean, genSigma);
    RooArgSet genVariables(x);
    auto * dataset = genGauss.generate(genVariables, nEvents);
    return dataset;
}

void PrintToFile(RooFitResult * _fitResults, TString _outFileName) {
    double _edm = _fitResults->edm();
    double _minNLL = _fitResults->minNll();
    int _covQual = _fitResults->covQual();
    int _fitStatus = _fitResults->status();

    ofstream _outFile(_outFileName);
    _outFile << "EDM\t: " << _edm << std::endl;
    _outFile << "likehood\t: " << _minNLL << std::endl;
    _outFile << "fitStatus\t: " << _covQual << std::endl;
    _outFile << "covQual\t: " << _fitStatus << std::endl;

    _fitResults->printMultiline(_outFile, 1111, kTRUE, "\t");
    _outFile.close();
}

void testMultiVarGaussianSmallCovariance() {

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

    // Create the datasets
    RooDataSet fullDataset("fullDataset", "fullDataset", RooArgSet(x));
    auto datasetGauss1 = GetGaussianDataset(x, 1e-4, 1e-4, 10000);
    auto datasetGauss2 = GetGaussianDataset(x, -2e-4, 1.5e-4, 15000);
    auto datasetBackground = GetExponentialDataset(x, -5e-6, 20000);
    fullDataset.append(*datasetGauss1);
    fullDataset.append(*datasetGauss2);
    fullDataset.append(*datasetBackground);

    // 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 constrainedPDFset(constraintPDF);

    // Fit to data with a covariance matrix
    auto * _fitResult = model.fitTo(fullDataset, RooFit::Extended(1), RooFit::Strategy(2), RooFit::Offset(kTRUE), RooFit::Optimize(kFALSE), RooFit::Save(1), RooFit::ExternalConstraints(constraintPDF));
    PrintToFile(_fitResult, "FitSmallCovariance.log");

    auto frame = x.frame();
    fullDataset.plotOn(frame);
    model.plotOn(frame, RooFit::LineStyle(kSolid), RooFit::LineColor(kRed));

    TCanvas c("c", "c", 800, 600);
    frame->Draw();
    c.SaveAs("FitSmallCovariance.pdf");
    delete frame;
    delete _fitResult;
}