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

using namespace std;

void xs(int type, float tb, TString bm = "mCHmin", TString decay = "BOTH") {
  thdm t(TString("/hepstore/gwilliam/thdm_grid_v165_AZHHZA-") + bm + ".root");

  gStyle->SetOptStat("111111");

  //TGraph2D* width = new TGraph2D(100);
  TH2D* width = new TH2D("name", "tilte", 9, 150, 1050, 9, 150, 1050);
  width->SetTitle("; m_{H} [ GeV ] ;m_{A} [ GeV ];#sigma * BR [ pb ] ");
  width->SetDirectory(0);

  bool azh = (decay == "AZH" || decay == "BOTH");
  bool hza = (decay == "HZA" || decay == "BOTH");

  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 > tb-epsilon                                               
	// theory constraints status
	// && t.status == 0
	){           

	//std::cout << ipt << " " << t.mH << " " << t.mA << " " << t.tanb << endl;

      if (azh && t.mA > t.mH) {
	width->SetBinContent(width->FindBin(t.mH), width->FindBin(t.mA), (t.xsec_A_gg + t.xsec_A_bb) * t.br_A_ZH * t.br_H_bb * 0.066);
	//width->SetPoint(ipt++, t.mH,t.mA,(t.xsec_A_gg + t.xsec_A_bb) * t.br_A_ZH);
      } else if (hza && t.mA < t.mH) {
	width->SetBinContent(width->FindBin(t.mH), width->FindBin(t.mA), (t.xsec_H_gg + t.xsec_H_bb) * t.br_H_ZA * t.br_A_bb * 0.066);
 	//width->SetPoint(ipt++, t.mH,t.mA,(t.xsec_H_gg + t.xsec_H_bb) * t.br_H_ZA);
       }
    }
  }

  std::cout << "Plotting ..." << std::endl;
  t.label->AddText(Form("#sigma * BR (%s)  #bf{#it{ATLAS}} Internal, SusHi+2HDMC, Type %d, tan(#beta)=%.1f", bm.Data(), type,tb));
  TString fn = Form("xsectbr_%s_%s_mH_vs_mA_type%d_tanb_%.1f", decay.Data(), bm.Data(), type, 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;
}


int main() {

  //xs(int type, float tb, TString bm = "mCHmin")
  float tanb[] = {0.5, 1.0, 5.0, 10.0, 20.0};
  float type[] = {1,2};
  TString decay[] = {"BOTH", "AZH", "HZA"};

  for (int itb(0); itb < 5; itb++) {
    for (int it(0); it < 2; it++) {
      for (int id(0); id < 1; id++) {
	xs(type[it], tanb[itb], "mCHmin", decay[id]);
	xs(type[it], tanb[itb], "mCHmax", decay[id]);
      }
    }
  }
}