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

#include <vector>

using namespace std;



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

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

  vector<float> mass;
  vector<float> tanb;
  vector<float> value;

  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 < 400 + 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
	){                             

      mass.push_back(t.mH);
      tanb.push_back(t.tanb);
      value.push_back((t.xsec_H_gg + (noVBF ? 0 : t.xsec_H_VBF))); // * t.br_H_ZZ);
    }
  }

  int type2 = (type == 1 ? 3 : 4);

  t.reset();
  while (t.next()) {

    //std::cout << t.mH << std::endl;
    if (t.type == type2  &&                                                                        
	// mH == mA == mCh 
	fabs(t.mA-t.mH) < epsilon &&                                                             
	fabs(t.mH-t.mCh) < epsilon &&
	// mass
	t.mH < 400 + 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
	){                             

      value[ipt] = ((t.xsec_H_gg + (noVBF ? 0 : t.xsec_H_VBF)) /** t.br_H_ZZ*/) / value[ipt];
      ipt++;
    }
  }

  int imax(0);
  float vmax(0);

  for (int i(0); i < mass.size(); i++) {
    width->SetPoint(i, mass[i] ,tanb[i], value[i]);
    if (fabs(value[i]) > vmax) {imax = i; vmax = fabs(value[i]);}
  }

  std::cout << "Max (" << mass[imax] << ", " << tanb[imax] << ") = " << value[imax] << std::endl;

  std::cout << "Plotting ..." << std::endl;
  t.label->AddText(Form("#sigma     #bf{#it{ATLAS}} Internal, SusHi+2HDMC, Type %d vs %d, cos(#beta-#alpha)=%.1f",type,type2,cba));
  TString fn = Form("xsectdiff_mH_vs_tanb_type%dvs%d_cba_%.1f_tanb_%d",type, type2, cba, tb);
  TCanvas* canvas = new TCanvas("","",800,600);
  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 pwvsm(double cba, int type, int tb=10, bool total = false) {
  thdm t;

  TGraph2D* width = new TGraph2D(100);
  if (total) 
    width->SetTitle("; m_{H};tan#beta;#Gamma");
  else
    width->SetTitle("; m_{H};tan#beta;#Gamma * BR");
  width->SetDirectory(0);

  vector<float> mass;
  vector<float> tanb;
  vector<float> value;

  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 < 400 + 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
	){                             

      mass.push_back(t.mH);
      tanb.push_back(t.tanb);
      value.push_back(t.width_H * (total ? 1 : t.br_H_ZZ));
    }
  }

  int type2 = (type == 1 ? 3 : 4);

  t.reset();
  while (t.next()) {

    //std::cout << t.mH << std::endl;
    if (t.type == type2  &&                                                                        
	// mH == mA == mCh 
	fabs(t.mA-t.mH) < epsilon &&                                                             
	fabs(t.mH-t.mCh) < epsilon &&
	// mass
	t.mH < 400 + 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
	){                             

      value[ipt] = (t.width_H * (total ? 1 : t.br_H_ZZ)) / value[ipt];
      ipt++;
    }
  }

  int imax(0);
  float vmax(0);

  for (int i(0); i < mass.size(); i++) {
    width->SetPoint(i, mass[i] ,tanb[i], value[i]);
    if (fabs(value[i]) > vmax) {imax = i; vmax = fabs(value[i]);}
  }

  std::cout << "Max (" << mass[imax] << ", " << tanb[imax] << ") = " << value[imax] << std::endl;

  std::cout << "Plotting ..." << std::endl;
  
  TString fn;
  
  if (total) {
    t.label->AddText(Form("#Gamma    #bf{#it{ATLAS}} Internal, SusHi+2HDMC, Type %d vs %d, cos(#beta-#alpha)=%.1f",type,type2,cba));
    fn = Form("twidthdiff_mH_vs_tanb_type%dvs%d_cba_%.1f_tanb_%d",type, type2, cba, tb);
  } else {
    t.label->AddText(Form("#Gamma * BR    #bf{#it{ATLAS}} Internal, SusHi+2HDMC, Type %d vs %d, cos(#beta-#alpha)=%.1f",type,type2,cba));
    fn = Form("pwidthdiff_mH_vs_tanb_type%dvs%d_cba_%.1f_tanb_%d",type, type2, cba, tb);
  }

  TCanvas* canvas = new TCanvas("","",800,600);
  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 brvsm(double cba, int type, int tb=10, bool alt=false) {
  thdm t;

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

  vector<float> mass;
  vector<float> tanb;
  vector<float> value;

  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 < 400 + 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
	){                             

      mass.push_back(t.mH);
      tanb.push_back(t.tanb);
      value.push_back(t.br_H_ZZ);
    }
  }

  int type2 = (type == 1 ? 3 : 4);

  if (alt) {
    type2 = (type == 1 ? 4 : 3);
  }

  t.reset();
  while (t.next()) {

    //std::cout << t.mH << std::endl;
    if (t.type == type2  &&                                                                        
	// mH == mA == mCh 
	fabs(t.mA-t.mH) < epsilon &&                                                             
	fabs(t.mH-t.mCh) < epsilon &&
	// mass
	t.mH < 400 + 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
	){                             

      value[ipt] = t.br_H_ZZ / value[ipt];
      ipt++;
    }
  }

  int imax(0);
  float vmax(0);

  for (int i(0); i < mass.size(); i++) {
    width->SetPoint(i, mass[i] ,tanb[i], value[i]);
    if (fabs(value[i]) > vmax) {imax = i; vmax = fabs(value[i]);}
  }

  std::cout << "Max (" << mass[imax] << ", " << tanb[imax] << ") = " << value[imax] << std::endl;

  std::cout << "Plotting ..." << std::endl;
  
  TString fn;
  t.label->AddText(Form("BR    #bf{#it{ATLAS}} Internal, SusHi+2HDMC, Type %d vs %d, cos(#beta-#alpha)=%.1f",type,type2,cba));
  fn = Form("brdiff_mH_vs_tanb_type%dvs%d_cba_%.1f_tanb_%d",type, type2, cba, tb);

  TCanvas* canvas = new TCanvas("","",800,600);
  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() {
  xsvsm(0.1, 1, 2);
  xsvsm(-0.1, 1, 4);
  xsvsm(0.1, 2, 2);
  xsvsm(-0.1, 2, 3);

  pwvsm(0.1, 1, 2);
  pwvsm(-0.1, 1, 4);
  pwvsm(0.1, 2, 2);
  pwvsm(-0.1, 2, 3);

  pwvsm(0.1, 1, 2, true);
  pwvsm(-0.1, 1, 4, true);
  pwvsm(0.1, 2, 2, true);
  pwvsm(-0.1, 2, 3, true);

  brvsm(0.1, 1, 2);
  brvsm(-0.1, 1, 4);
  brvsm(0.1, 2, 2);
  brvsm(-0.1, 2, 3);

//   brvsm(0.1, 1, 2, true);
//   brvsm(-0.1, 1, 4, true);
//   brvsm(0.1, 2, 2, true);
//   brvsm(-0.1, 2, 3, true);

}