using namespace RooFit ;

void ex4_makesample(Int_t code=1)
{
  RooAbsData::setDefaultStorageType(RooAbsData::Tree) ;

  // Create workspace
  RooWorkspace w("w",kTRUE) ;

  switch(code) {
  case 0: {
    // -----------------------------------------------------------------------
    // Sample 1 -- sig = Gaussian in x at -3, 
    //             bkg = Gaussian in x at +3 
    // -----------------------------------------------------------------------
    w.factory("PROD::sig(Gaussian(x[-10,10],-3,3),Uniform(y[-10,10]))") ;
    w.factory("PROD::bkg(Gaussian(x,3,3),Uniform(y))") ;
    w.defineSet("obs","x,y") ;
    break ;
  }

  case 1: {
    // -----------------------------------------------------------------------
    // Sample 1 -- sig = Gaussian in (x,y,z)
    //             bkg = Uniform(x,y,z)
    // -----------------------------------------------------------------------
    w.factory("PROD::sig(Gaussian(x[-10,10],0,3),Gaussian(y[-10,10],0,3),Gaussian(z[-10,10],0,3))") ;
    w.factory("PROD::bkg(Polynomial(x),Polynomial(y),Polynomial(z))") ;
    w.defineSet("obs","x,y,z") ;
    break ;
  }

  case 2: {
    // -----------------------------------------------------------------------
    // Sample 2 -- sig is Gaussian in (x,y,z) with strong correlations, 
    //             bkg is Gaussian in (x,y,z) with strong anticorrelations
    // -----------------------------------------------------------------------
    TMatrixDSym Vsig(3) ;
    Vsig(0,0) =  9   ; Vsig(1,0) =  7.2 ; Vsig(2,0) = 0 ;
    Vsig(0,1) =  7.2 ; Vsig(1,1) =  9   ; Vsig(2,1) = 4.5 ;
    Vsig(0,2) =  0   ; Vsig(1,2) =  4.5 ; Vsig(2,2) = 9 ;
    
    TMatrixDSym Vbkg(3) ;
    Vbkg(0,0) =  9   ; Vbkg(1,0) = -7.2 ; Vbkg(2,0) =   0 ;
    Vbkg(0,1) = -7.2 ; Vbkg(1,1) =    9 ; Vbkg(2,1) = -4.5 ;
    Vbkg(0,2) =   0  ; Vbkg(1,2) = -4.5 ; Vbkg(2,2) =  9 ;

    w.factory("{x[-10,10],y[-10,10],z[-10,10]}") ;
    TVectorD muSig(3) ; muSig(0) = 0 ; muSig(1) = 0 ; muSig(2) = 0 ; 
    TVectorD muBkg(3) ; muBkg(0) = 0 ; muBkg(1) = 0 ; muBkg(2) = 0 ; 
    RooMultiVarGaussian sig("sig","sig",RooArgList(w::x,w::y,w::z),muSig,Vsig) ;
    RooMultiVarGaussian bkg("bkg","bkg",RooArgList(w::x,w::y,w::z),muBkg,Vbkg) ;
    w.import(sig) ;
    w.import(bkg) ;
    w.defineSet("obs","x,y,z") ;
    break ;

  }

  case 3: {
    // -----------------------------------------------------------------------
    // Sample 3 -- sig = Gaussian in (x,y,z) at -3, with width 3
    //             bkg = Gaussian in (x,y,z) at +3, with width 5
    // -----------------------------------------------------------------------
    w.factory("PROD::sig(Gaussian(x[-10,10],-3,3),Gaussian(y[-10,10],-3,3),Gaussian(z[-10,10],-3,3))") ;
    w.factory("PROD::bkg(Gaussian(x,3,5),Gaussian(y,3,5),Gaussian(z,3,5))") ;
    w.defineSet("obs","x,y,z") ;
    break ;
  }

  case 4: {
    // -----------------------------------------------------------------------
    // Sample 4 -- sig = Circle in (x,y) center = (-2,-2), radius 5
    //             bkg = Circle in (x,y) center = ( 2, 2), radius 5
    // -----------------------------------------------------------------------
    w.factory("Gaussian::sig(expr('sqrt(((x+2)*(x+2))+((y+2)*(y+2)))',x[-10,10],y[-10,10]),5,1)") ;
    w.factory("Gaussian::bkg(expr('sqrt(((x-2)*(x-2))+((y-2)*(y-2)))',x,y),5,1)") ;
    w.defineSet("obs","x,y") ;
    break ;
  }



  default: {
    cout << "ERROR: sample code " << code << " is not defined" << endl ;
    return ;
  }
  }

  // Generate events
  RooDataSet* data_sig = w::sig.generate(*w.set("obs"),10000) ;
  RooDataSet* data_bkg = w::bkg.generate(*w.set("obs"),10000) ;    


  data_sig->Print("v") ;
  data_bkg->Print("v") ;
  
  // And write to file
  TFile f(Form("sample%d.root",code),"RECREATE") ;
  ((RooTreeDataStore*)data_sig->store())->tree().Clone()->Write("TreeS") ;
  ((RooTreeDataStore*)data_bkg->store())->tree().Clone()->Write("TreeB") ;
  f.Close() ;
  cout << "finished" << endl ; 
}