// This macro demonstrates generation of random points from 3D Gaussian distribution

//ROOT headers
#include "TH3D.h"
#include "TCanvas.h"

//roofit headers
#include "RooGlobalFunc.h"
#include "RooMultiVarGaussian.h"
#include "RooArgList.h"
#include "RooRealVar.h"
#include "RooAbsReal.h"
#include "RooFitResult.h"
#include "RooDataSet.h"

void Generate(Int_t nPart, TH3D *Output, Double_t XOffset, Double_t YOffset, Double_t ZOffset, Double_t SigmaX, Double_t SigmaY, Double_t SigmaZ, Double_t RhoXY=0, Double_t RhoXZ=0, Double_t RhoYZ=0)
{
  RooArgList xVec; // to store the random variable
  RooArgList muVec; // vector of mean values


  Int_t dim = 3; // 3D case
  RooRealVar* x; // auxilary variables
  RooRealVar* mu_x;
  for (Int_t i = 0; i < dim; i++) {	//Cycle to setup x variable and mu's 
    
    char* name = Form("x%d", i);
    if(i==0)
    x = new  RooRealVar(name, name, 0, -5*SigmaX, 5*SigmaX);// random variable in x-direction
    if(i==1)
    x = new  RooRealVar(name, name, 0, -5*SigmaY, 5*SigmaY);// random variable in y-direction
    if(i==2)
    x = new  RooRealVar(name, name, 0, -5*SigmaZ, 5*SigmaZ);// random variable in z-direction

    xVec.add(*x);// adding variable into a vector
   
    char* mu_name = Form("mu_x%d",i); // mean values
    if(i==0)
    mu_x = new RooRealVar(mu_name, mu_name, XOffset, -5*SigmaX+XOffset, 5*SigmaX+XOffset);// mu_x with initial value XOffset
    if(i==1)
    mu_x = new RooRealVar(mu_name, mu_name, YOffset, -5*SigmaY+YOffset, 5*SigmaY+YOffset);// mu_y with initial value YOffset
    if(i==2)
    mu_x = new RooRealVar(mu_name, mu_name, ZOffset, -5*SigmaZ+ZOffset, 5*SigmaZ+ZOffset);// mu_z with initial value ZOffset

    muVec.add(*mu_x);// adding into a vector
  	}

   // make a covariance matrix of 3D gaussian
   TMatrixDSym cov(dim);
   cov(1,0) = cov(0,1) = RhoXY*SigmaX*SigmaY;
   cov(0,2) = cov(2,0) = RhoXZ*SigmaX*SigmaZ;
   cov(1,2) = cov(2,1) = RhoYZ*SigmaY*SigmaZ;
   cov(0,0) = SigmaX*SigmaX;
   cov(1,1) = SigmaY*SigmaY;
   cov(2,2) = SigmaZ*SigmaZ;
   
   
   RooMultiVarGaussian mvg3D("mvg3D", "mvg3D",xVec, muVec, cov);// now make the multivariate Gaussian in 3D by giving it vector of random variables, mu's and covariance matrix

   RooDataSet* data = mvg3D.generate(xVec,nPart);// generate nPart (number) random points

   // Create Variables to store points from the distribution

   RooRealVar *a;
   RooRealVar *b;
   RooRealVar *c;
   RooArgList *l;

   for(Int_t i=0;i<nPart;i++) //for cycle to go through all generated data
	{
     	l = (RooArgList*) data->get(i);// Get point from data and split them into 3 dimentions.
     	a = (RooRealVar *) l->at(0);
     	b = (RooRealVar *) l->at(1);
	c = (RooRealVar *) l->at(2);          
	Output->Fill(a->getVal(),b->getVal(),c->getVal()); //Fill into 3D histogram
	}

}




void Sample()
{

  TH3D *test = new TH3D("Test", "Histogram with 3D Gaussian points", 100,-1,1, 100,-1,1, 100,-1,1); //Create histogram to put into Generate function
  Generate(1000, test, 0,0,0, 0.1,0.1,0.2);
  //Generate(1000, test, 0,0,0, 0.1,0.1,0.2, 0.9,0.9,0); //To show corellation
  test->Draw();

}