#include "thdm.h"
#include <cmath>
#include <TCanvas.h>
#include <TPaveText.h>
#include <TString.h>
#include <TStyle.h>
#include <TGraph2D.h>
#include <TH2.h>

using namespace std;


void xsvsm(double cba, int type, int tb=10, bool WW = false, bool noVBF = true) {
  thdm t;

  TGraph2D* width = new TGraph2D(100);
  width->SetTitle("; m_{H} [ GeV ] ;tan#beta;#sigma * BR [ pb ] ");
  width->SetDirectory(0);

  int ipt(0);
  double epsilon(0.00001);
  while (t.next()) {

    //std::cout << t.mH << std::endl;
    if (t.type == type  &&                                                                        
	// mH == mA == mCh 
	fabs(t.mA-t.mH) < epsilon &&                                                             
	fabs(t.mH-t.mCh) < epsilon &&
	// mass
	t.mH < 1500 + epsilon  &&                                                                   
	// cba/tanb
	t.cosba() > cba-epsilon && t.cosba() < cba+epsilon     &&     
	t.tanb < tb+epsilon && t.tanb > 0.5-epsilon                                               
	// theory constraints status
	// && t.status == 0
	){           

//       double tbtest = 2.;
//       if (t.cosba() > cba-epsilon && t.cosba() < cba+epsilon && t.tanb > tbtest-epsilon && t.tanb < tbtest+epsilon && (int(t.mH) % 100) == 0)
// 	std::cout << (WW ? "WW" : "ZZ") << " " 
// 		  << t.type << " " << t.mH << " " 
// 		  << t.cosba() << " " << t.tanb << " " 
// 		  << (t.xsec_H_gg + (noVBF ? 0 : t.xsec_H_VBF)) * (WW ? t.br_H_WW : t.br_H_ZZ) << std::endl;
                  
                       
      //std::cout << t.mA << " " << t.cosba() << " " << t.tanb << " " << t.br_H_ZZ << std::endl;
      width->SetPoint(ipt++, t.mH,t.tanb,(t.xsec_H_gg + (noVBF ? 0 : t.xsec_H_VBF))  * (WW ? t.br_H_WW : t.br_H_ZZ));
    }
  }

  std::cout << "Plotting ..." << std::endl;
  if (WW)
    t.label->AddText(Form("#sigma * BR(WW)  #bf{#it{ATLAS}} Internal, SusHi+2HDMC, Type %d, cos(#beta-#alpha)=%.1f",type,cba));
  else
    t.label->AddText(Form("#sigma * BR(ZZ)     #bf{#it{ATLAS}} Internal, SusHi+2HDMC, Type %d, cos(#beta-#alpha)=%.1f",type,cba));
  TString fn = Form("xsectbr_ZZ_mH_vs_tanb_type%d_cba_%.1f_tanb_%d",type, cba, tb);
  if (WW) fn = Form("xsectbr_WW_mH_vs_tanb_type%d_cba_%.1f_tanb_%d",type, cba, tb);
  TCanvas* canvas = new TCanvas("","",800,600);
  gPad->SetRightMargin(0.15);
  gStyle->SetPaintTextFormat(".1f");      
  //gStyle->SetPalette(4,0);
  width->SetMarkerSize(0.7);
  width->Draw("COLZ");  
  width->Draw("TEXT SAME");  
  //gPad->SetLogy(true);
  gPad->RedrawAxis();
  t.label->Draw();
  canvas->SaveAs("plots/"+fn+".eps");
  canvas->Print("plots/"+fn+".png");
  delete canvas;
}

void xs(int m , int type, int tb=10, bool WW = false, bool noVBF = true) {
  thdm t;

  TGraph2D* width = new TGraph2D(100);
  width->SetTitle("; cos(#beta-#alpha);tan#beta;#sigma * BR [ pb ]");
  width->SetDirectory(0);

  int ipt(0);
  double epsilon(0.00001);
  while (t.next()) {

    //std::cout << t.mH << std::endl;
    if (t.type == type  &&                                                                        
	// mH == mA == mCh 
	fabs(t.mA-t.mH) < epsilon &&                                                             
	fabs(t.mH-t.mCh) < epsilon &&
	// mass
	fabs(t.mH-m) < epsilon  &&                                                                   
	// cba/tanb
	fabs(t.cosba()) < 0.99+epsilon     &&                                                        
	t.tanb < tb+epsilon && t.tanb > 0.5-epsilon                                               
	// theory constraints status
	// && t.status == 0
	){                             
     
      width->SetPoint(ipt++, t.cosba(),t.tanb,(t.xsec_H_gg + (noVBF ? 0 : t.xsec_H_VBF)) * (WW ? t.br_H_WW : t.br_H_ZZ));
    }
  }

  std::cout << "Plotting ..." << std::endl;
  if (WW)
    t.label->AddText(Form("#sigma_{ggH+VBFH} * BR(WW)     #bf{#it{ATLAS}} Internal, SusHi+2HDMC, Type %d, m_{H}=%d",type,m));
  else
    t.label->AddText(Form("#sigma_{ggH+VBFH} * BR(ZZ)     #bf{#it{ATLAS}} Internal, SusHi+2HDMC, Type %d, m_{H}=%d",type,m));
  TString fn = Form("xsectbr_ZZ_cba_vs_tanb_type%d_cba_%d_tanb_%d",type, m, tb);
  if (WW) fn = Form("xsectbr_WW_cba_vs_tanb_type%d_cba_%d_tanb_%d",type, m, tb);
  TCanvas* canvas = new TCanvas("","",800,600);
  gPad->SetRightMargin(0.15);
  gStyle->SetPaintTextFormat(".2f");      
  //gStyle->SetPalette(4,0);
  width->SetMarkerSize(0.7);
  width->Draw("COLZ");  
  width->Draw("TEXT SAME");  
  //gPad->SetLogy(true);
  gPad->RedrawAxis();
  t.label->Draw();
  canvas->SaveAs("plots/"+fn+".eps");
  canvas->Print("plots/"+fn+".png");
  delete canvas;
}

int main() {
//   //xsvsm(0.1, 1, 10);
//   xsvsm(-0.1, 1, 10);
//   //xsvsm(0.1, 2, 10);
//   xsvsm(-0.1, 2, 10);
// 
//   //xsvsm(0.1, 1, 10, true);
//   xsvsm(-0.1, 1, 10, true);
//   //xsvsm(0.1, 2, 10, true);
//   xsvsm(-0.1, 2, 10, true);


//   xsvsm(0.1, 3, 10);
//   xsvsm(-0.1, 3, 10);
//   xsvsm(0.1, 4, 10);
//   xsvsm(-0.1, 4, 10);

//    xs(200, 1, 10);
//    xs(200, 2, 10);
//    xs(200, 1, 10, true);
//    xs(200, 2, 10, true);

//   xs(200, 3, 10);
//   xs(200, 4, 10);

    xs(500, 1, 10, true);
    xs(500, 2, 10, true);
    xs(500, 1, 10);
    xs(500, 2, 10);

}