#include "TH1.h"
#include "TH2.h"
#include "TCanvas.h"
#include "TProfile.h"
#include "TMath.h"
#include "TGraph.h"
#include "TRandom.h"
#include <iostream>
#include <vector>

#define mc_scale 1.0

std::vector<Double_t> trialsig_mc;
TH1F *mean_distro=0,*min_distro=0,*min_distro_over_mean=0,*distro=0;
TH1F *k_at_maxsig_mc_distro=0,*k_at_maxsig_mc_distro_over_mean=0;
TH2F *h_sig=0;
TCanvas *acanvas=0;
//     SRS  SWS  SBRS   SBWS  S       a (%)    B
// kpi 112  36   31     39    77+-12  9.42    35.24
// k3p 95   49   131    155   48+-12  9.73    46.67
// sat 271  180  217    289  114+-22  32.47   156.6
// dpl 860  514  1352   1349 345+-41  32.10   515

#define NPOINTS  11

TCanvas *mycanvas;
TGraph kpi(NPOINTS);
TGraph k3pi(NPOINTS);
TGraph sat(NPOINTS);
TGraph dpl(NPOINTS);
 
Double_t kpi_scale[NPOINTS];
Double_t k3pi_scale[NPOINTS];
Double_t sat_scale[NPOINTS];
Double_t dpl_scale[NPOINTS];
Double_t correlation[NPOINTS]={0.0,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,0.95};


void poisson(
	     Double_t srs_mean=112,
	     Double_t sws_mean=36,
	     Double_t sbrs_mean=31,
	     Double_t sbws_mean=39,
	     Double_t a=0.0942,
	     char name[]="kpi",
	     Double_t correlation=0.65,
	     Int_t ntrials=1000,
	     Int_t nexperiments=1000,
	     Bool_t quiet=false)
{
  double x_mc;
  double x_srs,x_sws;
  double x_sbrs,x_sbws;
  double x_b;
  double k,meank,mink;
  double sig_mc;
  double maxsig_mc=0.;
  double k_at_maxsig_mc;
  double s_mean,b_mean;

  b_mean=sws_mean+a*(sbrs_mean-sbws_mean);
  s_mean=srs_mean-b_mean;
  
  double mean=b_mean;

  char hname[1000];
  char htitle[1000];

  sprintf(hname,"%s_sig",name); sprintf(htitle,"%s significance",name);
  if (h_sig==0) h_sig=new TH2F(hname,htitle,ntrials,0,ntrials,50,0.0,2.0*s_mean/sqrt(s_mean+b_mean));
  else h_sig->Reset();

  // this histogram contains the difference between the value of b extracted (BG) and the central value
  // of b (BG~). It is filled for each cut in each toy experiment
  sprintf(hname,"%s_distro",name); sprintf(htitle,"%s distro",name);
  if (distro==0) distro=new TH1F(hname,htitle,50,-3.0*sqrt(mean),3.0*sqrt(mean));
  else distro->Reset();

  // this histogram contains the same quantity as above but averaged over the N(=ntrials) cuts in each toy experiment
  sprintf(hname,"%s_mean_distro",name); sprintf(htitle,"%s mean distro",name);
  if (mean_distro==0) mean_distro=new TH1F(hname,htitle,50,-3.0*sqrt(mean),3.0*sqrt(mean));
  else mean_distro->Reset();

  // this histogram is filled with x_b again, but this time we pick the minimum value of x_b among the N(=ntrials) cuts
  sprintf(hname,"%s_min_distro",name); sprintf(htitle,"%s min distro",name);
  if (min_distro==0) min_distro=new TH1F(hname,htitle,50,-mean,0);
  else min_distro->Reset();

  // same as above but divided by the expected number of BG events
  sprintf(hname,"%s_min_distro_over_mean",name); sprintf(htitle,"%s min distro/BG",name);
  if (min_distro_over_mean==0) min_distro_over_mean=new TH1F(hname,htitle,50,-1.0,0.0);
  else min_distro_over_mean->Reset();

  // Now we pick the value of x_b=BG-BG~ ath the maximum value of the estimator used in the optimization
  sprintf(hname,"%s_k_at_maxsig_mc_distro",name); sprintf(htitle,"%s dk at max significance (MC)",name);
  if (k_at_maxsig_mc_distro==0) k_at_maxsig_mc_distro=new TH1F(hname,htitle,50,-mean,0);
  else k_at_maxsig_mc_distro->Reset();

  // Same as above, but divided by BG
  sprintf(hname,"%s_k_at_maxsig_mc_distro_over_mean",name); sprintf(htitle,"%s k/BG at max significance (MC)",name);
  if (k_at_maxsig_mc_distro_over_mean==0) k_at_maxsig_mc_distro_over_mean=new TH1F(hname,htitle,50,-2.0,0.0);
  else k_at_maxsig_mc_distro_over_mean->Reset();

  for (Int_t j=0;j<nexperiments;j++)
    {
      meank=0; mink=0; maxsig_mc=0;

      x_mc=gRandom->Poisson(s_mean*mc_scale)/mc_scale;

      Double_t x_srs_base =gRandom->Poisson(srs_mean*correlation );
      Double_t x_sws_base =gRandom->Poisson(sws_mean*correlation );
      Double_t x_sbrs_base=gRandom->Poisson(sbrs_mean*correlation);
      Double_t x_sbws_base=gRandom->Poisson(sbws_mean*correlation);

      for (Int_t i=0;i<ntrials;i++)
	{
	  x_srs =x_srs_base +gRandom->Poisson(srs_mean*(1.0-correlation));
	  x_sbrs=x_sbrs_base+gRandom->Poisson(sbrs_mean*(1.0-correlation));
	  x_sws =x_sws_base +gRandom->Poisson(sws_mean*(1.0-correlation));
	  x_sbws=x_sbws_base+gRandom->Poisson(sbws_mean*(1.0-correlation));
	  
	  x_b=x_sws+a*(x_sbrs-x_sbws);

	  k=x_b-b_mean; 
	  sig_mc=x_mc/sqrt(x_mc+2.0*x_sws+a*a*(x_sbrs+x_sbws)+a*(x_sbrs-x_sbws));

	  meank+=k;
	  mink=(mink>k)?k:mink;

	  if (maxsig_mc<sig_mc)
	    {
	      maxsig_mc=sig_mc;
	      k_at_maxsig_mc=k;
	    }

	  distro->Fill(k);
	  trialsig_mc.push_back(sig_mc);
	}

      std::sort(trialsig_mc.begin(),trialsig_mc.end());
      for (Int_t i=0;i<ntrials;i++)
	h_sig->Fill(ntrials-i,trialsig_mc[i]);
      trialsig_mc.clear();

      meank/=ntrials;
      mean_distro->Fill(meank);
      min_distro->Fill(mink);
      min_distro_over_mean->Fill(mink/mean);

      k_at_maxsig_mc_distro->Fill(k_at_maxsig_mc);
      k_at_maxsig_mc_distro_over_mean->Fill(k_at_maxsig_mc/mean);
    }
  if (!quiet)
    {
      char cname[1000];
      sprintf(cname,"%s_canvas",name);
      if (acanvas==0) acanvas=new TCanvas(cname);

      acanvas->Divide(4,3);
      acanvas->cd(1); distro->Draw();
      acanvas->cd(5); mean_distro->Draw();
      
      acanvas->cd(2); min_distro->Draw();
      acanvas->cd(6); min_distro_over_mean->Draw();
      
      acanvas->cd(4); k_at_maxsig_mc_distro->Draw();
      acanvas->cd(8); k_at_maxsig_mc_distro_over_mean->Draw();
      
      acanvas->cd(11); h_sig->Draw();
      acanvas->cd(12); h_sig->ProfileX()->Draw();
      
      std::cout << "S: " << s_mean << "    B: " << b_mean << 
	" BG scale (MC): " <<  1.0/(1.0+k_at_maxsig_mc_distro_over_mean->GetMean()) << 
	" pessimistic: " << 1.0/(1.0+min_distro_over_mean->GetMean()) << std::endl;
    }
}

void plotcorr()
{
  Int_t ntr=1000;
  Int_t nexp=1000;
  for (Int_t i=0;i<NPOINTS;i++)
    {
      poisson(112,36,31,39,0.0942,"kpi"    ,correlation[i],ntr,nexp,true);
      kpi_scale[i] =1.0/(1.0+k_at_maxsig_mc_distro_over_mean->GetMean());

      poisson(95,49,131,155,0.0973,"k3pi"  ,correlation[i],ntr,nexp,true);
      k3pi_scale[i]=1.0/(1.0+k_at_maxsig_mc_distro_over_mean->GetMean());

      poisson(271,180,217,289,0.3247,"sat" ,correlation[i],ntr,nexp,true);
      sat_scale[i] =1.0/(1.0+k_at_maxsig_mc_distro_over_mean->GetMean());

      poisson(860,514,1352,1349,0.321,"dpl",correlation[i],ntr,nexp,true);
      dpl_scale[i] =1.0/(1.0+k_at_maxsig_mc_distro_over_mean->GetMean());
                                                                                                                   
      kpi.SetPoint(i,correlation[i],kpi_scale[i]);
      k3pi.SetPoint(i,correlation[i],k3pi_scale[i]);
      sat.SetPoint(i,correlation[i],sat_scale[i]);
      dpl.SetPoint(i,correlation[i],dpl_scale[i]);
                                                                                                                   
      std::cout << "Point " << i << " Corr: " << correlation[i]
                << " KPI: " << kpi_scale[i]
                << " K3PI: " << k3pi_scale[i]
                << " SAT: " << sat_scale[i]
                << " DPL: " << dpl_scale[i]
                << std::endl;
                                                                                                                   
    }
                                                                                                                   
  kpi.SetTitle("K#pi event scale vs correlation");
  k3pi.SetTitle("K3#pi event scale vs correlation");
  sat.SetTitle("Satellite event scale vs correlation");
  dpl.SetTitle("D^{+} event scale vs correlation");
  mycanvas=new TCanvas("a");
  mycanvas->Divide(2,2);
  mycanvas->cd(1);  kpi.Draw("ALS");
  mycanvas->cd(2); k3pi.Draw("ALS");
  mycanvas->cd(3);  sat.Draw("ALS");
  mycanvas->cd(4);  dpl.Draw("ALS");
                                                                                                                   
}