#include <TCanvas.h>
#include <TF1.h>
#include <TMath.h>
#include <TLegend.h>
#include <TAxis.h>
#include <TLatex.h>

#include <Math/WrappedTF1.h>
#include <Math/GSLIntegrator.h>

// This analysis uses the "h_n" likelihoods from my thesis.

double P1(const double * const x, const double * const pars)
{
  const double & b = x[0];
  const double & rho = pars[0];
  const double & x1 = pars[1];

  if(b > x1) std::cerr << "Error, b (first argument, 0th element) "
    "cannot be greater than x1 (second argument, 1th element)!\n";

  const double y = TMath::Sqrt(x1*x1-b*b);
  const double z = rho*x1;

  return 2.0*TMath::Exp(-rho*y)/(TMath::Pi()*y*(TMath::BesselI0(z)-TMath::StruveL0(z)));
}

double N2_integrand(const double * const x,const double * const pars)
{
  const double & b = x[0];
  const double & rho = pars[0];
  const double & x1 = pars[1];
  const double & x2 = pars[2];

  const double y1 = TMath::Sqrt(x1*x1-b*b);
  const double y2 = TMath::Sqrt(x2*x2-b*b);

  return TMath::Exp(-rho*y2)/(y1*y2);
}

double N2(double rho, double x1, double x2)
{
  TF1 n2_integrand("n2_integrand",&N2_integrand,0,x1,3);
  n2_integrand.SetParameters(rho,x1,x2);
  ROOT::Math::WrappedTF1 wn2(n2_integrand);

  ROOT::Math::GSLIntegrator ig(ROOT::Math::IntegrationOneDim::kADAPTIVE);

  ig.SetFunction(wn2);
  ig.SetRelTolerance(0.0001);

  return ig.Integral(0,x1);
}

double P2(const double * const x, const double * const pars)
{
  const double & b = x[0];
  const double & rho = pars[0];
  const double & x1 = pars[1];
  const double & x2 = pars[2];

  const double y1 = TMath::Sqrt(x1*x1-b*b);
  const double y2 = TMath::Sqrt(x2*x2-b*b);

  return TMath::Exp(-rho*y2)/(y1*y2*N2(rho,x1,x2));
}

double N3_integrand(const double * const x,const double * const pars)
{
  const double & b = x[0];
  const double & rho = pars[0];
  const double & x1 = pars[1];
  const double & x2 = pars[2];
  const double & x3 = pars[3];

  const double y1 = TMath::Sqrt(x1*x1-b*b);
  const double y2 = TMath::Sqrt(x2*x2-b*b);
  const double y3 = TMath::Sqrt(x3*x3-b*b);

  return TMath::Exp(-rho*y3)/(y3*y2*y1);
}

double N3(double rho, double x1, double x2, double x3)
{
  TF1 n3_integrand("n3_integrand",&N3_integrand,0,x1,4);
  n3_integrand.SetParameters(rho,x1,x2,x3);
  ROOT::Math::WrappedTF1 wn3(n3_integrand);

  ROOT::Math::GSLIntegrator ig(ROOT::Math::IntegrationOneDim::kADAPTIVE);

  ig.SetFunction(wn3);
  ig.SetRelTolerance(0.0001);

  return ig.Integral(0,x1);
}

double P3(const double * const x, const double * const pars)
{
  const double & b = x[0];
  const double & rho = pars[0];
  const double & x1 = pars[1];
  const double & x2 = pars[2];
  const double & x3 = pars[3];

  const double y1 = TMath::Sqrt(x1*x1-b*b);
  const double y2 = TMath::Sqrt(x2*x2-b*b);
  const double y3 = TMath::Sqrt(x3*x3-b*b);

  return TMath::Exp(-rho*y3)/(y3*y2*y1*N3(rho,x1,x2,x3));
}



//////////////////////////////////////////////////
// GRAPHS 
//////////////////////////////////////////////////



TCanvas * graph_p1()
{
  double rho = 1;
  double x1_0 = 0.8;
  double x1_1 = 1;
  double x1_2 = 1.2;

  double lowlim = 0.7, highlim = 1;

  TF1 * fP1_0 = new TF1("fP1_0",&P1,0,x1_0,2);
  TF1 * fP1_1 = new TF1("fP1_0",&P1,0,x1_1,2);
  TF1 * fP1_2 = new TF1("fP1_0",&P1,0,x1_2,2);

  fP1_0->SetParameters(rho,x1_0);
  fP1_1->SetParameters(rho,x1_1);
  fP1_2->SetParameters(rho,x1_2);

  fP1_0->SetLineWidth(1);
  fP1_1->SetLineWidth(1);
  fP1_2->SetLineWidth(1);
  
  fP1_1->SetLineColor(kBlue);
  fP1_2->SetLineColor(kRed);

  fP1_0->SetNpx(1000);
  fP1_1->SetNpx(1000);
  fP1_2->SetNpx(1000);

  fP1_0->SetMaximum(20);

  TCanvas * c1 = new TCanvas();
  c1->cd();
  c1->DrawFrame(0.5,0.25,1.25,30);

  TLatex * tlt = new TLatex;
  tlt->DrawLatexNDC(0.6,0.02,"Impact parameter b");
  tlt->SetTextAngle(90);
  tlt->DrawLatexNDC(0.04,0.5,"Probability density");
  
  //fP1_2->GetXaxis()->SetLimits(lowlim,1.25);
  fP1_0->Draw("same");
  fP1_1->Draw("same");
  fP1_2->Draw("same");

  gPad->SetLogy();

  fP1_2->SetTitle("");
  fP1_2->GetXaxis()->SetTitle("Impact parameter b");
  fP1_2->GetYaxis()->SetTitle("Probability Density");


  TLegend * tl = new TLegend(0.12,0.6,0.35,0.88);
  tl->SetFillStyle(0);
  tl->SetTextSize(0.05);
  tl->AddEntry(fP1_0,TString::Format("x_{1} = %g",x1_0).Data());
  tl->AddEntry(fP1_1,TString::Format("x_{1} = %g",x1_1).Data());
  tl->AddEntry(fP1_2,TString::Format("x_{1} = %g",x1_2).Data());
  tl->Draw();

  c1->SaveAs("theory_p1.pdf");
  return c1;
}


TCanvas * graph_p2()
{
  double lowlim = 0.7, highlim = 1;

  TF1 * fP1 = new TF1("fP1",&P1,lowlim,highlim,2);
  TF1 * fP2_0 = new TF1("fP2_0",&P2,lowlim,highlim,3);
  TF1 * fP2_1 = new TF1("fP2_1",&P2,lowlim,highlim,3);

  double rho = 1;
  double x1 = 1;
  double x2_0 = 1.1;
  double x2_1 = 2;

  fP1->SetParameters(rho,x1);
  fP2_0->SetParameters(rho,x1,x2_0);
  fP2_1->SetParameters(rho,x1,x2_1);

  fP1->SetLineWidth(1);
  fP2_0->SetLineWidth(1);
  fP2_0->SetLineColor(kBlue);
  fP2_1->SetLineWidth(1);
  fP2_1->SetLineColor(kRed);

  fP1->SetNpx(1000);
  fP2_0->SetNpx(1000);
  fP2_1->SetNpx(1000);

  fP1->SetMaximum(20);

  TCanvas * c1 = new TCanvas();
  c1->cd();
  c1->DrawFrame(lowlim,0.5,1.01,30);

  TLatex * tlt = new TLatex;
  tlt->DrawLatexNDC(0.6,0.02,"Impact parameter b");
  tlt->SetTextAngle(90);
  tlt->DrawLatexNDC(0.04,0.5,"Probability density");

  fP1->Draw("same");
  fP2_0->Draw("same");
  fP2_1->Draw("same");

  gPad->SetLogy();

  fP1->SetTitle("");
  fP1->GetXaxis()->SetTitle("Impact parameter b");
  fP1->GetYaxis()->SetTitle("Probability Density Function");

  TLegend * tl = new TLegend(0.12,0.5,0.6,0.88);
  tl->SetFillStyle(0);
  tl->SetTextSize(0.04);
  tl->AddEntry(fP1,TString::Format("First Cluster Only x_{1} = %g",x1));
  tl->AddEntry(fP2_0,TString::Format("2nd Cluster x_{2} = %g",x2_0));
  tl->AddEntry(fP2_1,TString::Format("2nd Cluster x_{2} = %.42g",x2_1));
  tl->Draw();

  c1->SaveAs("theory_h_p2.pdf");
  return c1;
}

TCanvas * graph_p3(double x3_0 = 1.6, double x3_1 = 2.5)
{
  double lowlim = 0.7, highlim = 1;

  TF1 * fP1 = new TF1("fP1",&P1,lowlim,highlim,2);
  TF1 * fP2 = new TF1("fP2",&P2,lowlim,highlim,3);
  TF1 * fP3_0 = new TF1("fP3_0",&P3,lowlim,highlim,4);
  TF1 * fP3_1 = new TF1("fP3_1",&P3,lowlim,highlim,4);

  double rho = 1;
  double x1 = 1;
  double x2 = 1.5;
  //double x3_0 = 1.6;
  //double x3_1 = 2.5;

  fP1->SetParameters(rho,x1);
  fP2->SetParameters(rho,x1,x2);
  fP3_0->SetParameters(rho,x1,x2,x3_0);
  fP3_1->SetParameters(rho,x1,x2,x3_1);

  fP1->SetLineWidth(1);
  fP2->SetLineWidth(1);
  fP2->SetLineColor(kBlue);
  fP3_0->SetLineWidth(1);
  fP3_1->SetLineWidth(1);
  fP3_0->SetLineColor(kRed);
  fP3_1->SetLineColor(kMagenta);

  fP1->SetNpx(1000);
  fP2->SetNpx(1000);
  fP3_0->SetNpx(1000);
  fP3_1->SetNpx(1000);

  fP1->SetMaximum(20);

  TCanvas * c1 = new TCanvas();
  c1->cd();
  c1->DrawFrame(lowlim,0.5,1.01,30);

  TLatex * tlt = new TLatex;
  tlt->DrawLatexNDC(0.6,0.02,"Impact parameter b");
  tlt->SetTextAngle(90);
  tlt->DrawLatexNDC(0.04,0.5,"Probability density");

  fP1->Draw("same");
  fP2->Draw("same");
  fP3_0->Draw("same");
  fP3_1->Draw("same");

  gPad->SetLogy();

  //TString ftitle = TString::Format("The effect of the 3rd cluster (rho = %g)",rho);
  //fP1->SetTitle(ftitle);
  fP1->SetTitle("");
  fP1->GetXaxis()->SetTitle("Impact parameter b");
  fP1->GetYaxis()->SetTitle("Probability Density Function");

  TLegend * tl = new TLegend(0.12,0.5,0.55,0.88);
  tl->SetFillStyle(0);
  tl->SetTextSize(0.03);
  tl->AddEntry(fP1,TString::Format("First Cluster Only x_{1} = %g",x1));
  tl->AddEntry(fP2,TString::Format("1st & 2nd Clusters only x_{2} = %g",x2));
  tl->AddEntry(fP3_0,TString::Format("3rd Cluster x_{3} = %g",x3_0));
  tl->AddEntry(fP3_1,TString::Format("3rd Cluster x_{3} = %g",x3_1));
  tl->Draw();

  c1->SaveAs("theory_p3.pdf");
  return c1;
}

void theory()
{
  graph_p1();
  graph_p2();
  graph_p3();
}