using namespace RooStats;
using namespace RooFit;
void poissonRuns(){
 

  int nruns;
  
 

  int nON_tot=0;
  int nOFF_tot=0;
  double tON_tot=0;
  double tOFF_tot=0;
  double Q_tot=0;

  ifstream fin("data.txt");
  fin>>nruns;

  int *runsID=new int[nruns];
  int *nON=new int[nruns];
  int *nOFF=new int[nruns];
  double *tON=new double[nruns];
  double *tOFF=new double[nruns];
  double *Q=new double[nruns];
  double *Qf=new double[nruns];
  double *tR=new double[nruns]; //Toff-over-Ton ratio
  
  
  int m_nON,m_nOFF,m_run;
  double m_tON,m_tOFF,m_Q;
  int nruns_good=0;
  for (int ii=0;ii<nruns;ii++){
    fin>>m_run>>m_nON>>m_nOFF>>m_tON>>m_tOFF>>m_Q;
    if (m_Q!=0){
      runsID[nruns_good]=m_run;
      nON[nruns_good]=(int)(m_nOFF*m_tON/m_tOFF); //A.C.
      nOFF[nruns_good]=m_nOFF;
      tON[nruns_good]=m_tON;
      tOFF[nruns_good]=m_tOFF;
      Q[nruns_good]=m_Q;
      
      tR[nruns_good]=tOFF[nruns_good]/tON[nruns_good];
      
      //totals
      nON_tot+=nON[nruns_good];
      nOFF_tot+=nOFF[nruns_good];
      tON_tot+=tON[nruns_good];
      tOFF_tot+=tOFF[nruns_good];
      Q_tot+=Q[nruns_good];
      nruns_good++;
    }
  }
  
  for (int ii=0;ii<nruns_good;ii++){
    Qf[ii]=Q[ii]/Q_tot;
    cout<<ii<<" "<<Qf[ii]<<" "<<Qf[ii]*nON_tot<<endl;
  }
  
  RooWorkspace w("w",true);
  
  //for the signal: s_i = s*Qi/Qtot, so that sum(s_i)=s
  RooRealVar s("s","signal",nON_tot,0,100*nON_tot);
  w.import(s);
  
  //create the b_i-nuisance parameters and the nb_i observables
  RooRealVar *Vb[nruns_good];
  RooProduct *Vbprod[nruns_good];
  RooRealVar *Vnb[nruns_good];
  for (int ii=0;ii<nruns_good;ii++){
    Vb[ii]=new RooRealVar(Form("b_%i",ii),Form("b_%i",ii),nOFF[ii],0,100*nOFF[ii]);
    Vbprod[ii]=new RooProduct(Form("bprod_%i",ii),Form("bprod_%i",ii),RooArgList(*Vb[ii],RooConst(tR[ii])));
    Vnb[ii]=new RooRealVar(Form("nb_%i",ii),Form("nb_%i",ii),nOFF[ii],0,100*nOFF[ii]);
  }

  //create the nsb_i observables
  RooRealVar *Vnsb[nruns_good];
  RooProduct *Vsprod[nruns_good];
  RooAddition *Vsb[nruns_good];
  
  for (int ii=0;ii<nruns_good;ii++){
    //RooConstVar f(Form("qf_%i",ii),Form("qf_%i",ii),Qf[ii]);
    Vsprod[ii]=new RooProduct(Form("sprod_%i",ii),Form("sprod_%i",ii),RooArgList(s,RooConst(Qf[ii])));
  }
  for (int ii=0;ii<nruns_good;ii++){
    Vsb[ii]=new RooAddition(Form("sb_%i",ii),Form("sb_%i",ii),RooArgList(*Vsprod[ii],*Vb[ii]));
  }
  for (int ii=0;ii<nruns_good;ii++){
    Vnsb[ii]=new RooRealVar(Form("nsb_%i",ii),Form("nsb_%i",ii),nON[ii],0,100*nON[ii]);
  }

  //create the two poisson PDFs, and their product
  RooPoisson *Vpsb[nruns_good];
  RooPoisson *Vpb[nruns_good];
  RooProdPdf *Vprod[nruns_good];

  for (int ii=0;ii<nruns_good;ii++){
    Vpsb[ii]=new RooPoisson(Form("psb_%i",ii),Form("psb_%i",ii),*Vnsb[ii],*Vsb[ii]);
  }
  for (int ii=0;ii<nruns_good;ii++){
    Vpb[ii]=new RooPoisson(Form("pb_%i",ii),Form("pb_%i",ii),*Vnb[ii],*Vbprod[ii]); //need Vbprod, due to the Toff/Ton ratio
  }
  for (int ii=0;ii<nruns_good;ii++){
    Vprod[ii]=new RooProdPdf(Form("prod_%i",ii),Form("prod_%i",ii),RooArgList(*Vpsb[ii],*Vpb[ii]));
  }

  //Now, the final product
  RooArgList factors("factors");
  for (int ii=0;ii<nruns_good;ii++){
    factors.add(*Vprod[ii]);
  }
   
  
  RooProdPdf model("model","model",factors);
  w.import(model);
  
  
  //Set the observed nOFF observables
  for (int ii=0;ii<nruns_good;ii++){
    w.var(Form("nb_%i",ii))->setVal(nOFF[ii]);
    w.var(Form("nb_%i",ii))->setConstant(true);//needed for being treated as global observables
  }
  
  //Set the observed nON observables
  for (int ii=0;ii<nruns_good;ii++){
    w.var(Form("nsb_%i",ii))->setVal(nON[ii]);
  }

  //make data set with the namber of observed nsb events
  RooArgList nsb_obs("nsb_obs");
  RooArgList nb_obs("nb_obs");
  RooArgList b_obs("b_obs");
  for (int ii=0;ii<nruns_good;ii++){
    nsb_obs.add(*Vnsb[ii]);
  }
  for (int ii=0;ii<nruns_good;ii++){
    nb_obs.add(*Vnb[ii]);
  }
  for (int ii=0;ii<nruns_good;ii++){
    b_obs.add(*Vb[ii]);
  }
		   
  RooDataSet data("data","",RooArgSet(nsb_obs,nb_obs));
  data.add(RooArgSet(nsb_obs,nb_obs));
  w.import(data);


   w.Print();
  
   ModelConfig mc("Smodel",&w);
   mc.SetPdf(*w.pdf("model"));
   mc.SetParametersOfInterest(*w.var("s"));
   mc.SetObservables(RooArgSet(nsb_obs,nb_obs));
  
   // these are needed for the hypothesis tests
   mc.SetSnapshot(*w.var("s"));
   mc.SetNuisanceParameters(b_obs);


   //construct the B hypo (s=0)
   RooStats::ModelConfig* bModel = (RooStats::ModelConfig*) mc.Clone("Bmodel") ;
   RooRealVar* poi = (RooRealVar*) bModel->GetParametersOfInterest()->first();
   poi->setVal(0);
   bModel->SetSnapshot(*poi);
   
   mc.Print();
   cout<<"---------"<<endl;
   bModel->Print();
   // import model in the workspace
   cout<<"Import MC"<<endl;
   w.import(mc);
   cout<<"Import Bmodel"<<endl;
   w.import(*bModel);
   cout<<"OK"<<endl;

  
  


   //this is, instead, the part for the TOTAL
   RooWorkspace w2("w2",true);
   RooRealVar nb_tot("nb_tot","nb_tot",nOFF_tot,0,10*nOFF_tot);
   RooRealVar nsb_tott("nsb_tott","nsb_tott",nON_tot,0,10*nON_tot); //yes, two t
   RooRealVar b_tot("b_tot","b_tot",nOFF_tot,0,10*nOFF_tot); 
   RooRealVar s_tot("s_tot","s_tot",nON_tot,0,10*nON_tot);
   w2.import(nb_tot);
   w2.import(nsb_tott);
   w2.import(b_tot);
   w2.import(s_tot);
   
   RooProduct bprod_tot("bprod_tot","bprod_tot",RooArgList(b_tot,RooConst(tOFF_tot/tON_tot)));
   RooAddition sb_tot("sb_tot","sb_tot",RooArgList(s_tot,b_tot));

   RooPoisson pb_tot("pb_tot","pb_tot",nb_tot,bprod_tot);
   RooPoisson psb_tot("psb_tot","psb_tot",nsb_tott,sb_tot);
   RooProdPdf model_tot("model_tot","model_tot",RooArgList(pb_tot,psb_tot));
   w2.import(model_tot);

   //set value of observed events
  RooRealVar *nsb_tot = w2.var("nsb_tott");
  nsb_tot->setVal(nON_tot);

 

  // make data set with the namber of observed events
  RooDataSet data_tot("data_tot","",RooArgSet(*nsb_tot,nb_tot));
  data_tot.add(RooArgSet(*nsb_tot,nb_tot));
  w2.import(data_tot);

  w2.Print();


   ModelConfig mc_tot("Smodel_tot",&w2);
   mc_tot.SetPdf(*w2.pdf("model_tot"));
   mc_tot.SetParametersOfInterest(*w2.var("s_tot"));
   mc_tot.SetObservables(RooArgList(*w2.var("nsb_tott"),*w2.var("nb_tot")));
  
   // these are needed for the hypothesis tests
   mc_tot.SetSnapshot(*w2.var("s_tot"));
   mc_tot.SetNuisanceParameters(*w2.var("b_tot"));


   //construct the B hypo (s=0)
   RooStats::ModelConfig* bModel_tot = (RooStats::ModelConfig*) mc_tot.Clone("Bmodel_tot") ;
   RooRealVar* poi_tot = (RooRealVar*) bModel_tot->GetParametersOfInterest()->first();
   poi_tot->setVal(0);
   bModel_tot->SetSnapshot(*poi_tot);
   
   mc_tot.Print();
   cout<<"---------"<<endl;
   bModel_tot->Print();
   // import model in the workspace 
   w2.import(mc_tot);
   w2.import(*bModel_tot);

 
    TString fileName = "poissonRuns.out_separate.root"; 
   // write workspace in the file (recreate file if already existing)
   w.writeToFile(fileName, true);

   TString fileName_tot = "poissonRuns.out_tot.root"; 
   // write workspace in the file (recreate file if already existing)
   w2.writeToFile(fileName_tot, true);
   


   
}