#include <TLorentzVector.h>

void reconstruct_D0()
  
{
  TFile *f = new TFile("FileName.root");
  TTreeReader reader("events", f);
  
  TTreeReaderArray<Float_t> pxArray(reader, "ReconstructedParticles.p.x");
  TTreeReaderArray<Float_t> pyArray(reader, "ReconstructedParticles.p.y");
  TTreeReaderArray<Float_t> pzArray(reader, "ReconstructedParticles.p.z");
  TTreeReaderArray<Float_t> energyArray(reader, "ReconstructedParticles.energy");
  TTreeReaderArray<Int_t> pdgArray(reader, "mcparticles2.pdgID");

  reader.Next();                                                                                                                              
  cout << "Length of pxArray is: " <<pxArray.GetSize()<<"\n";
  cout << "Length of pyArray is: " <<pyArray.GetSize()<<"\n";
  cout << "Length of pzArray is: " <<pzArray.GetSize()<<"\n";
  cout << "Length of energyArray is: " <<energyArray.GetSize()<<"\n";
  cout << "Length of pdgArray is: " <<pdgArray.GetSize()<<"\n";
  
  TH1F *h1 = new TH1F("K_px","",100, -20, 20);
  h1->GetXaxis()->SetTitle("K_px");
  h1->GetXaxis()->SetTitleSize(0.05);
  h1->GetXaxis()->CenterTitle();
  h1->SetLineWidth(2);

  TH1F *h2 = new TH1F("K_py","",100, -20, 20);
  h2->GetXaxis()->SetTitle("K_py");
  h2->GetXaxis()->SetTitleSize(0.05);
  h2->GetXaxis()->CenterTitle();
  h2->SetLineWidth(2);

  TH1F *h3 = new TH1F("K_pz","",100, -20, 20);
  h3->GetXaxis()->SetTitle("K_pz");
  h3->GetXaxis()->SetTitleSize(0.05);
  h3->GetXaxis()->CenterTitle();
  h3->SetLineWidth(2);

  TH1F *h4 = new TH1F("K_energy","",100, 0, 80);
  h4->GetXaxis()->SetTitle("K_energy");
  h4->GetXaxis()->SetTitleSize(0.05);
  h4->GetXaxis()->CenterTitle();
  h4->SetLineWidth(2);
  /////////  

  TH1F *h5 = new TH1F("pi_px","",100, -20, 20);
  h5->GetXaxis()->SetTitle("pi_px");
  h5->GetXaxis()->SetTitleSize(0.05);
  h5->GetXaxis()->CenterTitle();
  h5->SetLineWidth(2);

  TH1F *h6 = new TH1F("pi_py","",100, -20, 20);
  h6->GetXaxis()->SetTitle("pi_py");
  h6->GetXaxis()->SetTitleSize(0.05);
  h6->GetXaxis()->CenterTitle();
  h6->SetLineWidth(2);

  TH1F *h7 = new TH1F("pi_pz","",100, -20, 20);
  h7->GetXaxis()->SetTitle("pi_pz");
  h7->GetXaxis()->SetTitleSize(0.05);
  h7->GetXaxis()->CenterTitle();
  h7->SetLineWidth(2);

  TH1F *h8 = new TH1F("pi_energy","",100, 0, 80);
  h8->GetXaxis()->SetTitle("pi_energy");
  h8->GetXaxis()->SetTitleSize(0.05);
  h8->GetXaxis()->CenterTitle();
  h8->SetLineWidth(2);

  
  TH1F *h9 = new TH1F("K_M","",100, 0, 1);
  h9->GetXaxis()->SetTitle("K_mass");
  h9->GetXaxis()->SetTitleSize(0.05);
  h9->GetXaxis()->CenterTitle();
  h9->SetLineWidth(2);

  
  
 
    
  Int_t n = 8;
  Double_t K_px[n];
  Double_t K_py[n];
  Double_t K_pz[n];
  Double_t K_energy[n];
  
 

           
  while (reader.Next())
    {

      for (auto pdg: pdgArray)
      {
	  if(pdg==-321)
	    {	  
	      for (auto K_px_new: pxArray)
		h1->Fill(K_px_new);
		  
	      for (auto K_py_new: pyArray)
		h2->Fill(K_py_new);
		  
	      for (auto K_pz_new: pzArray)
		h3->Fill(K_pz_new);
		  
	      for (auto K_energy_new: energyArray)
		h4->Fill(K_energy_new);
		  
	      /* Double_t K_M;
	      for (int i = 0; i < pxArray.GetSize(); ++i)
	      TLorentzVector K_V(K_px[i], K_py[i], K_pz[i], K_energy[i]);
	      K_M = K_V.Mag();
	      h9->Fill(K_M);
	      */

	       std::vector<TLorentzVector> lorentzvectors;
	      for (int i = 0; i < pxArray.GetSize(); ++i)
		lorentzvectors.emplace_back(K_px[i], K_py[i], K_pz[i], K_energy[i]);
	      K_M = lorentzvectors.Mag();
	      h9->Fill(K_M);
	    }
     
	  

      if(pdg==211)
	{
	  for (auto pi_px: pxArray)
	    h5->Fill(pi_px);
	  for (auto pi_py: pyArray)
	    h6->Fill(pi_py);
	  for (auto pi_pz: pzArray)
                h7->Fill(pi_pz);
	  for (auto pi_energy: energyArray)
	    h8->Fill(pi_energy); 
	}
      
      }
  
  
    }


  TFile * f1 = new TFile("output.root","RECREATE");
  h1->Write();
  h2->Write();
  h3->Write();
  h4->Write();
  h5->Write();
  h6->Write();
  h7->Write();
  h8->Write();
  h9->Write();

  
  f1->Close(); 
}