#include <iostream>
#include <cmath>
#include <vector>

#include "TChain.h"
#include "TTreeReader.h"
#include "TTreeReaderArray.h"
#include "TCanvas.h"
#include "TH1D.h"
#include "TLegend.h"
#include "TLorentzVector.h"
#include "TStyle.h"


// ============================================================
// Electron-method DIS kinematics
//
// Return order:
//   [0] = x
//   [1] = y
//   [2] = Q2
// ============================================================

std::vector<double> GetElectronKinematics(const TLorentzVector& ei, const TLorentzVector& ef, const TLorentzVector& pni)
{
  TLorentzVector q_e = ei - ef;
  double Q2 = -1.0 * (q_e * q_e);
  double y = (pni * q_e) / (pni * ei);
  double x = Q2 / (2.0 * (pni * q_e));

  return {x, y, Q2};
}


// ============================================================
// Main
// ============================================================

void CheckKinematics()
{
  gStyle->SetOptStat(0);

  TChain chain("events");

  // Change to your input file
  chain.Add("pythia*100to1000*.root");


  // ==========================================================
  // Beam configuration
  // ==========================================================

  const double Ee = 9.0;
  const double Ep = 130.0;
  const double crossingAngle = 25e-3;


  // ==========================================================
  // TTreeReader
  // ==========================================================

  TTreeReader reader(&chain);


  // ----------------------------------------------------------
  // MCParticles
  // ----------------------------------------------------------

  TTreeReaderArray<int> MCParticles_PDG(reader, "MCParticles.PDG");
  TTreeReaderArray<int> MCParticles_generatorStatus(reader, "MCParticles.generatorStatus");

  TTreeReaderArray<double> MCParticles_momentum_x(reader, "MCParticles.momentum.x");
  TTreeReaderArray<double> MCParticles_momentum_y(reader, "MCParticles.momentum.y");
  TTreeReaderArray<double> MCParticles_momentum_z(reader, "MCParticles.momentum.z");

  TTreeReaderArray<double> MCParticles_mass(reader, "MCParticles.mass");


  // ----------------------------------------------------------
  // ReconstructedParticles
  // ----------------------------------------------------------

  TTreeReaderArray<float> ReconstructedParticles_momentum_x(reader, "ReconstructedParticles.momentum.x");
  TTreeReaderArray<float> ReconstructedParticles_momentum_y(reader, "ReconstructedParticles.momentum.y");
  TTreeReaderArray<float> ReconstructedParticles_momentum_z(reader, "ReconstructedParticles.momentum.z");
  TTreeReaderArray<float> ReconstructedParticles_energy(reader, "ReconstructedParticles.energy");


  // ----------------------------------------------------------
  // ScatteredElectronsTruth object index
  // ----------------------------------------------------------

  TTreeReaderArray<int> ScatteredElectronsTruth_objIdx_index(reader, "ScatteredElectronsTruth_objIdx.index");


  // ==========================================================
  // Histograms
  // ==========================================================

  TH1D* hXTrue = new TH1D("hXTrue", ";x;Normalised events", 1000, 1e-4, 1.0);
  TH1D* hXReco = new TH1D("hXReco", ";x;Normalised events", 1000, 1e-4, 1.0);

  TH1D* hYTrue = new TH1D("hYTrue", ";y;Normalised events", 1000, 1e-4, 1.0);
  TH1D* hYReco = new TH1D("hYReco", ";y;Normalised events", 1000, 1e-4, 1.0);

  TH1D* hQ2True = new TH1D("hQ2True", ";Q^{2} [GeV^{2}];Normalised events", 1000, 1.0, 1e3);
  TH1D* hQ2Reco = new TH1D("hQ2Reco", ";Q^{2} [GeV^{2}];Normalised events", 1000, 1.0, 1e3);

  TH1D* hXReso = new TH1D("hXReso", ";(x_{rec}-x_{true})/x_{true};Events", 100, -1.0, 1.0);
  TH1D* hYReso = new TH1D("hYReso", ";(y_{rec}-y_{true})/y_{true};Events", 100, -1.0, 1.0);
  TH1D* hQ2Reso = new TH1D("hQ2Reso", ";(Q^{2}_{rec}-Q^{2}_{true})/Q^{2}_{true};Events", 100, -1.0, 1.0);

  // const Int_t nbins = 40;
  // Double_t xmin = 1e-4;  // Must be strictly greater than 0
  // Double_t xmax = 1e0;
  // Double_t Q2min = 1e-1;
  // Double_t Q2max = 1e4;
  //  
  // Double_t logxmin = TMath::Log10(xmin);
  // Double_t logxmax = TMath::Log10(xmax);
  // Double_t xbinwidth = (logxmax - logxmin) / nbins;
  // 
  // Double_t logQ2min = TMath::Log10(Q2min);
  // Double_t logQ2max = TMath::Log10(Q2max);
  // Double_t Q2binwidth = (logQ2max - logQ2min) / nbins;
  //  
  // Double_t xbins[nbins + 1]; 
  // Double_t Q2bins[nbins + 1]; 
  //  
  // for (Int_t i = 0; i <= nbins; i++) {
  // xbins[i] = TMath::Power(10, logxmin + i * xbinwidth);
  // Q2bins[i] = TMath::Power(10, logQ2min + i * Q2binwidth);
  // }
  //  
  // hXTrue->SetBins(nbins, xbins);
  // hXReco->SetBins(nbins, xbins);
  // hYTrue->SetBins(nbins, xbins);
  // hYReco->SetBins(nbins, xbins);
  // hQ2True->SetBins(nbins, Q2bins);
  // hQ2Reco->SetBins(nbins, Q2bins);


  // ==========================================================
  // Event loop
  // ==========================================================

  Long64_t nEntries = chain.GetEntries();

  Long64_t nTruth = 0;
  Long64_t nReco = 0;

  std::cout << "Entries: " << nEntries << std::endl;


  while (reader.Next())
    {
      // ------------------------------------------------------
      // Find truth incoming electron, proton and scattered
      // electron
      // ------------------------------------------------------

      TLorentzVector ei_true;
      TLorentzVector ef_true;
      TLorentzVector pni_true;

      bool found_ei_true = false;
      bool found_ef_true = false;
      bool found_pni_true = false;


      for (size_t i = 0; i < MCParticles_PDG.GetSize(); ++i)
        {
          int pdg = MCParticles_PDG[i];
          int status = MCParticles_generatorStatus[i];

          double px = MCParticles_momentum_x[i];
          double py = MCParticles_momentum_y[i];
          double pz = MCParticles_momentum_z[i];
          double mass = MCParticles_mass[i];

          double E = std::sqrt(px*px + py*py + pz*pz + mass*mass);

          TLorentzVector particle(px, py, pz, E);


          // Incoming electron
          if (status == 4 && pdg == 11)
            {
              ei_true = particle;
              found_ei_true = true;
            }


          // Incoming proton
          if (status == 4 && pdg == 2212)
            {
              pni_true = particle;
              found_pni_true = true;
            }


          // First status-1 electron
          if (!found_ef_true && status == 1 && pdg == 11)
            {
              ef_true = particle;
              found_ef_true = true;
            }
        }


      if (!found_ei_true || !found_pni_true || !found_ef_true)
        continue;


      // ------------------------------------------------------
      // Truth kinematics
      // ------------------------------------------------------

      auto kin_true = GetElectronKinematics(ei_true, ef_true, pni_true);

      double x_true = kin_true[0];
      double y_true = kin_true[1];
      double Q2_true = kin_true[2];

      if (y_true < 0.01) continue;


      // ------------------------------------------------------
      // Nominal reconstructed beam four-vectors
      // ------------------------------------------------------

      TLorentzVector ei_reco(0.0, 0.0, -Ee, Ee);

      TLorentzVector pni_reco(-1 * Ep * std::sin(crossingAngle),
                              0.0,
                              Ep * std::cos(crossingAngle),
                              Ep);


      // ------------------------------------------------------
      // Get reconstructed scattered electron
      //
      // ScatteredElectronsTruth_objIdx.index[0]
      // gives the index into ReconstructedParticles.
      // ------------------------------------------------------

      if (ScatteredElectronsTruth_objIdx_index.GetSize() == 0)
        continue;


      int recoIndex = ScatteredElectronsTruth_objIdx_index[0];


      if (recoIndex < 0 || recoIndex >= (int)ReconstructedParticles_momentum_x.GetSize())
        continue;


      double px_reco = ReconstructedParticles_momentum_x[recoIndex];
      double py_reco = ReconstructedParticles_momentum_y[recoIndex];
      double pz_reco = ReconstructedParticles_momentum_z[recoIndex];
      double E_reco = ReconstructedParticles_energy[recoIndex];


      TLorentzVector ef_reco(px_reco, py_reco, pz_reco, E_reco);


      // ------------------------------------------------------
      // Reconstructed kinematics
      // ------------------------------------------------------

      auto kin_reco = GetElectronKinematics(ei_reco, ef_reco, pni_reco);

      double x_reco = kin_reco[0];
      double y_reco = kin_reco[1];
      double Q2_reco = kin_reco[2];


      // ------------------------------------------------------
      // Truth sanity cuts
      // ------------------------------------------------------

      if (x_true <= 0.0 || x_true > 1.0)
        continue;

      if (y_true <= 0.0 || y_true >= 1.0)
        continue;

      if (Q2_true <= 0.0)
        continue;


      hXTrue->Fill(x_true);
      hYTrue->Fill(y_true);
      hQ2True->Fill(Q2_true);

      ++nTruth;


      // ------------------------------------------------------
      // Reco sanity cuts
      // ------------------------------------------------------

      if (x_reco <= 0.0 || x_reco > 1.0)
        continue;

      if (y_reco <= 0.0 || y_reco >= 1.0)
        continue;

      if (Q2_reco <= 0.0)
        continue;


      hXReco->Fill(x_reco);
      hYReco->Fill(y_reco);
      hQ2Reco->Fill(Q2_reco);


      // ------------------------------------------------------
      // Resolutions
      // ------------------------------------------------------

      hXReso->Fill((x_reco - x_true) / x_true);
      hYReso->Fill((y_reco - y_true) / y_true);
      hQ2Reso->Fill((Q2_reco - Q2_true) / Q2_true);

      ++nReco;
    }


  // ==========================================================
  // Normalise distributions
  // ==========================================================

  if (hXTrue->Integral() > 0)
    hXTrue->Scale(1.0 / hXTrue->Integral());

  if (hXReco->Integral() > 0)
    hXReco->Scale(1.0 / hXReco->Integral());

  if (hYTrue->Integral() > 0)
    hYTrue->Scale(1.0 / hYTrue->Integral());

  if (hYReco->Integral() > 0)
    hYReco->Scale(1.0 / hYReco->Integral());

  if (hQ2True->Integral() > 0)
    hQ2True->Scale(1.0 / hQ2True->Integral());

  if (hQ2Reco->Integral() > 0)
    hQ2Reco->Scale(1.0 / hQ2Reco->Integral());


  // ==========================================================
  // Style
  // ==========================================================

  hXTrue->SetLineColor(kBlack);
  hXTrue->SetLineWidth(2);

  hXReco->SetLineColor(kRed);
  hXReco->SetLineWidth(2);

  hYTrue->SetLineColor(kBlack);
  hYTrue->SetLineWidth(2);

  hYReco->SetLineColor(kRed);
  hYReco->SetLineWidth(2);

  hQ2True->SetLineColor(kBlack);
  hQ2True->SetLineWidth(2);

  hQ2Reco->SetLineColor(kRed);
  hQ2Reco->SetLineWidth(2);


  // ==========================================================
  // Truth vs reconstructed
  // ==========================================================

  TCanvas* cKin = new TCanvas("cKin", "Electron kinematics", 1500, 500);
  cKin->Divide(3, 1);


  cKin->cd(1);
  gPad->SetLogx();

  hXReco->Draw("HIST");
  hXTrue->Draw("HIST SAME");

  TLegend* legX = new TLegend(0.60, 0.75, 0.88, 0.88);
  legX->SetBorderSize(0);
  legX->AddEntry(hXTrue, "Truth", "l");
  legX->AddEntry(hXReco, "Reconstructed", "l");
  legX->Draw();


  cKin->cd(2);
  gPad->SetLogx();

  hYTrue->Draw("HIST");
  hYReco->Draw("HIST SAME");

  TLegend* legY = new TLegend(0.60, 0.75, 0.88, 0.88);
  legY->SetBorderSize(0);
  legY->AddEntry(hYTrue, "Truth", "l");
  legY->AddEntry(hYReco, "Reconstructed", "l");
  legY->Draw();


  cKin->cd(3);
  // gPad->SetLogx();

  hQ2Reco->Draw("HIST");
  hQ2True->Draw("HIST SAME");

  TLegend* legQ2 = new TLegend(0.60, 0.75, 0.88, 0.88);
  legQ2->SetBorderSize(0);
  legQ2->AddEntry(hQ2True, "Truth", "l");
  legQ2->AddEntry(hQ2Reco, "Reconstructed", "l");
  legQ2->Draw();


  // ==========================================================
  // Resolutions
  // ==========================================================

  TCanvas* cReso = new TCanvas("cReso", "Kinematic resolutions", 1500, 500);
  cReso->Divide(3, 1);


  cReso->cd(1);
  hXReso->Draw("HIST");


  cReso->cd(2);
  hYReso->Draw("HIST");


  cReso->cd(3);
  hQ2Reso->Draw("HIST");


  // ==========================================================
  // Print summary
  // ==========================================================

  std::cout << std::endl;
  std::cout << "Entries in tree:        " << nEntries << std::endl;
  std::cout << "Events with truth kine: " << nTruth << std::endl;
  std::cout << "Events with reco kine:  " << nReco << std::endl;
  std::cout << std::endl;

  std::cout << "x resolution mean  = "
            << hXReso->GetMean()
            << " +/- "
            << hXReso->GetMeanError()
            << std::endl;

  std::cout << "y resolution mean  = "
            << hYReso->GetMean()
            << " +/- "
            << hYReso->GetMeanError()
            << std::endl;

  std::cout << "Q2 resolution mean = "
            << hQ2Reso->GetMean()
            << " +/- "
            << hQ2Reso->GetMeanError()
            << std::endl;
}