#include "ROOT/RDataFrame.hxx"
#include "TCanvas.h"
#include "TLegend.h"
#include "TH1D.h"
#include "TProfile.h"

#include <iostream>

R__LOAD_LIBRARY(libeicd.so)
R__LOAD_LIBRARY(libDD4pod.so)
#include "dd4pod/Geant4ParticleCollection.h"
#include "eicd/TrackParametersCollection.h"
#include "eicd/ClusterCollection.h"
#include "eicd/ClusterData.h"
#include "eicd/TrackerHitCollection.h"

using ROOT::RDataFrame;
using namespace ROOT::VecOps;

auto p_track = [](std::vector<eic::TrackParametersData> const& in) {
  std::vector<double> result;
  for (size_t i = 0; i < in.size(); ++i) {
    result.push_back(std::abs(1.0/(in[i].qOverP)));
  }
  return result;
};


std::vector<float> pt (std::vector<dd4pod::Geant4ParticleData> const& in){
  std::vector<float> result;
  for (size_t i = 0; i < in.size(); ++i) {
    result.push_back(std::sqrt(in[i].ps.x * in[i].ps.x + in[i].ps.y * in[i].ps.y));
  }
  return result;
}

auto momentum = [](std::vector<ROOT::Math::PxPyPzMVector> const& in) {
  std::vector<double> result;
  for (size_t i = 0; i < in.size(); ++i) {
   result.push_back(in[i].P());
  }
  return result;
};
auto theta = [](std::vector<ROOT::Math::PxPyPzMVector> const& in) {
  std::vector<double> result;
  for (size_t i = 0; i < in.size(); ++i) {
   result.push_back(in[i].Theta()*180/M_PI);
  }
  return result;
};
auto fourvec = [](ROOT::VecOps::RVec<dd4pod::Geant4ParticleData> const& in) {
  std::vector<ROOT::Math::PxPyPzMVector> result;
  ROOT::Math::PxPyPzMVector lv;
  for (size_t i = 0; i < in.size(); ++i) {
    lv.SetCoordinates(in[i].ps.x, in[i].ps.y, in[i].ps.z, in[i].mass);
    result.push_back(lv);
  }
  return result;
};

auto delta_p = [](const std::vector<double>& tracks, const std::vector<double>& thrown) {
  std::vector<double> res;
  for (const auto& p1 : thrown) {
    for (const auto& p2 : tracks) {
      res.push_back(p1 - p2);
    }
  }
  return res;
};

auto delta_p_over_p = [](const std::vector<double>& tracks, const std::vector<double>& thrown) {
  std::vector<double> res;
  for (const auto& p1 : thrown) {
    for (const auto& p2 : tracks) {
      res.push_back((p1 - p2)/p1);
    }
  }
  return res;
};

int rec_multiple_tracks(const char* fname = "topside/rec_multiple_tracks.root")
{

  ROOT::EnableImplicitMT();
  ROOT::RDataFrame df("events", fname);

  auto df0 = df.Define("isThrown", "mcparticles.genStatus == 1")
                 .Define("thrownParticles", "mcparticles[isThrown]")
                 .Define("thrownP", fourvec, {"thrownParticles"})
                 .Define("p_thrown", momentum, {"thrownP"})
                 .Define("theta_thrown", theta, {"thrownP"})
                 .Define("theta0", "theta_thrown[0]")
                 .Define("nTracks", "outputTrackParameters.size()")
                 .Define("p_track", p_track, {"outputTrackParameters"})
                 .Define("delta_p0",delta_p, {"p_track", "p_thrown"})
                 .Define("delta_p_over_p0",delta_p_over_p, {"p_track", "p_thrown"})
                 .Define("N_Hits",       [](std::vector<eic::TrackerHitData> hits) { return hits.size();}, {"trackingHits"})
                 .Define("N_BarrelHits", [](std::vector<eic::TrackerHitData> hits) { return hits.size();}, {"TrackerBarrelRecHits"})
                 .Define("N_EndcapHits", [](std::vector<eic::TrackerHitData> hits) { return hits.size();}, {"TrackerEndcapRecHits"})
                 ;

  auto h_nTracks_vs_theta = df0.Histo2D({"h_nTracks_vs_theta", "; #theta; N tracks ", 40,0,180,10, 0, 10}, "theta0","nTracks");

  auto h_nTracks = df0.Histo1D({"h_nTracks", "; N tracks ", 10, 0, 10}, "nTracks");
  auto h_pTracks = df0.Histo1D({"h_pTracks", "; GeV/c ", 100, 0, 10}, "p_track");

  auto h_delta_p0  = df0.Histo1D({"h_delta_p0", "Truth Track Init; GeV/c ",  100, -10, 10}, "delta_p0");

  auto h_delta_p0_over_p = df0.Histo1D({"h_delta_p0_over_p",  "Truth Track Init; delta p/p ",  100, -0.1, 0.1}, "delta_p_over_p0");

  auto hNhits_vs_theta = df0.Histo1D({"hNhits_vs_theta", "; #theta [deg.]",   40, 0, 180 }, "theta0", "N_Hits");
  auto hBarrel_N_vs_theta = df0.Histo1D({"hBarrel_N_vs_theta", "; #theta [deg.]",   40, 0, 180 }, "theta0", "N_BarrelHits");
  auto hEndcap_N_vs_theta = df0.Histo1D({"hEndcap_N_vs_theta", "; #theta [deg.]",   40, 0, 180 }, "theta0", "N_EndcapHits");

  auto hHits_Nhits  = df0.Histo1D({"hHits_Nhits", "; #theta [deg.]",       20, 0, 20 }, "N_Hits");
  auto hBarrel_Nhits  = df0.Histo1D({"hBarrel_Nhits", "; #theta [deg.]",   20, 0, 20 }, "N_BarrelHits");
  auto hEndcap_Nhits  = df0.Histo1D({"hEndcap_Nhits", "; #theta [deg.]",   20, 0, 20 }, "N_EndcapHits");

  auto hHits_Ntheta = df0.Histo1D({"hHits_Ntheta", "; #theta [deg.]",       40, 0, 180 }, "theta0");
  auto hBarrel_Ntheta = df0.Histo1D({"hBarrel_Ntheta", "; #theta [deg.]",   40, 0, 180 }, "theta0");
  auto hEndcap_Ntheta = df0.Histo1D({"hEndcap_Ntheta", "; #theta [deg.]",   40, 0, 180 }, "theta0");

  // -----------------------------------------------
  auto c = new TCanvas();

  h_nTracks->DrawCopy();
  c->SaveAs("results/tracking/rec_multiple_tracks_nTracks.png");
  c->SaveAs("results/tracking/rec_multiple_tracks_nTracks.pdf");

  // -----------------------------------------------
  h_pTracks->DrawCopy();
  c->SaveAs("results/tracking/rec_multiple_tracks_pTracks.png");
  c->SaveAs("results/tracking/rec_multiple_tracks_pTracks.pdf");

  // -----------------------------------------------
  c = new TCanvas();
  THStack * hs = new THStack("hs_delta_p","; GeV/c ");
  TH1D* h1 = (TH1D*) h_delta_p0->Clone();
  hs->Add(h1);
  hs->Draw("nostack");
  c->BuildLegend();
  c->SaveAs("results/tracking/rec_multiple_tracks_delta_p.png");
  c->SaveAs("results/tracking/rec_multiple_tracks_delta_p.pdf");

  // -----------------------------------------------
  c  = new TCanvas();
  hs = new THStack("hs_delta_p_over_p","; delta p/p ");
  h1 = (TH1D*) h_delta_p0_over_p->Clone();
  hs->Add(h1);
  hs->Draw("nostack");
  c->BuildLegend();
  c->SaveAs("results/tracking/rec_multiple_tracks_delta_p_over_p.png");
  c->SaveAs("results/tracking/rec_multiple_tracks_delta_p_over_p.pdf");

  // -----------------------------------------------
  c  = new TCanvas();
  hs = new THStack("n_hits","; #theta  ");

  h1 = (TH1D*) hBarrel_N_vs_theta->Clone();
  auto h2 = (TH1D*) hBarrel_Ntheta->Clone();
  h1->SetLineColor(4);
  h1->Divide(h2);
  hs->Add(h1);

  h1 = (TH1D*) hEndcap_N_vs_theta->Clone();
  h2 = (TH1D*) hEndcap_Ntheta->Clone();
  h1->Divide(h2);
  h1->SetLineColor(2);
  hs->Add(h1);

  h1 = (TH1D*) hNhits_vs_theta->Clone();
  h2 = (TH1D*) hHits_Ntheta->Clone();
  h1->Divide(h2);
  h1->SetLineColor(1);
  hs->Add(h1);

  hs->Draw("nostack, hist");
  c->BuildLegend();
  c->SaveAs("results/tracking/rec_multiple_tracks_n_hits_vs_theta.png");
  c->SaveAs("results/tracking/rec_multiple_tracks_n_hits_vs_theta.pdf");

  // -----------------------------------------------
  c  = new TCanvas();
  hs = new THStack("theta","; #theta  ");

  h2 = (TH1D*) hBarrel_Ntheta->Clone();
  h2->SetLineColor(4);
  hs->Add(h2);

  h2 = (TH1D*) hEndcap_Ntheta->Clone();
  h2->SetLineColor(2);
  hs->Add(h2);

  h2 = (TH1D*) hHits_Ntheta->Clone();
  h2->SetLineColor(1);
  hs->Add(h2);

  hs->Draw("nostack hist");
  c->BuildLegend();
  c->SaveAs("results/tracking/rec_multiple_tracks_theta.png");
  c->SaveAs("results/tracking/rec_multiple_tracks_theta.pdf");

  // -----------------------------------------------
  c  = new TCanvas();
  hs = new THStack("hits","; hits  ");
  h1 = (TH1D*) hBarrel_Nhits->Clone();
  h1->SetLineColor(4);
  hs->Add(h1);
  h1 = (TH1D*) hEndcap_Nhits->Clone();
  h1->SetLineColor(2);
  hs->Add(h1);
  h1 = (TH1D*) hHits_Nhits->Clone();
  h1->SetLineColor(2);
  hs->Add(h1);
  c->BuildLegend();
  c->SaveAs("results/tracking/rec_multiple_tracks_nhits.png");
  c->SaveAs("results/tracking/rec_multiple_tracks_nhits.pdf");

  // -----------------------------------------------
  c  = new TCanvas();
  h_nTracks_vs_theta->DrawCopy("colz");
  c->SaveAs("results/tracking/rec_multiple_tracks_nTracks_vs_theta.png");
  c->SaveAs("results/tracking/rec_multiple_tracks_nTracks_vs_theta.pdf");
  return 0;
}