#include <TMath.h>
#include <TSystem.h>
#include <TFile.h>
#include <TTree.h>
#include <TChain.h>
#include <TString.h>

#include "TROOT.h"
#include <FairMCEventHeader.h>
#include <MpdEvent.h>
#include <TClonesArray.h>
#include <MpdTrack.h>
#include <FairMCTrack.h>
#include <TH1.h>
#include <TF1.h>

#include <iostream>

//#define WEIGHT_B

using namespace std;

void get_multiplicity(TString inFileName, TString outFileName, TString dcaFileName)
{

	const float pt_bins[] = {0., 0.2, 0.4, 0.6, 0.8, 1.0, 1.2, 1.4, 1.6, 1.8, 2.0, 2.5, 3.};
	const int n_pt_bin = 12;

	const float eta_bins[] = {-1.5, -1.2, -1., -0.8, -0.6, -0.4, -0.2, 0., 0.2, 0.4, 0.6, 0.8, 1., 1.2, 1.5};
	const int n_eta_bin = 14;

	const Int_t n_proj = 3;
	const Int_t n_hit_cut = 32;

	TH1F *h_multiplicity = new TH1F("h_multiplicity", "multiplicity in TPC after cuts;N_{tracks} in TPC;N_{events};", 1600, 0., 1600.);
	TH1F *h_mult_compare = new TH1F("h_mult_compare", "Normalized multiplicity in TPC after cuts;N_{tracks} in TPC;N_{events};", 1600, 0., 1600.);
	TH1F *h_mult_old = new TH1F("h_mult_old", "multiplicity in TPC before cuts;N_{tracks} in TPC;N_{events};", 1600, 0., 1600.);

        TChain *dstTree = new TChain("cbmsim");
        dstTree->Add(inFileName.Data());

	TFile *outFile = new TFile(outFileName.Data(), "RECREATE");

	//TFile  *dcaFile = new TFile("/lustre/nyx/hades/user/parfenov/dca_out_file_TDR.root","READ");
	TFile *dcaFile = new TFile(dcaFileName.Data(), "READ");

	MpdEvent *MPDEvent = nullptr;
	dstTree->SetBranchAddress("MPDEvent.", &MPDEvent);
	TClonesArray *MpdGlobalTracks = 0;
	TClonesArray *MCTracks = 0;
	dstTree->SetBranchAddress("MCTrack", &MCTracks);
	#ifdef WEIGHT_B
	FairMCEventHeader *MCHeader = 0;
	dstTree->SetBranchAddress("MCEventHeader.", &MCHeader);
	#endif

	//TF1* f_dca[n_proj][n_pt_bin][n_eta_bin];
	TF1 *f_pt_fit[n_proj][n_eta_bin];

	for (Int_t i_proj = 0; i_proj < n_proj; i_proj++)
	{
		//for (Int_t i_pt=0;i_pt<n_pt_bin;i_pt++){
		for (Int_t i_eta = 0; i_eta < n_eta_bin; i_eta++)
		{
			//f_dca[i_proj][i_pt][i_eta] = (TF1*) dcaFile->Get(Form("dca_fit[%i][%i][%i]",i_proj,i_pt,i_eta));
			f_pt_fit[i_proj][i_eta] = (TF1 *)dcaFile->Get(Form("sigma_pt_fit%i%i", i_proj, i_eta));
		}
		//}
	}

	Int_t n_events = dstTree->GetEntries();
	//Int_t n_events = 50;
	Int_t pt_bin;
	Int_t eta_bin;
	Double_t wB,B,dB;
	for (int event = 0; event < n_events; ++event)
	{
		cout << "EVENT N " << event << endl;
		dstTree->GetEntry(event);
		MpdGlobalTracks = MPDEvent->GetGlobalTracks();

		#ifdef WEIGHT_B
		B  = MCHeader->GetB();
		dB = 0.25;
		wB = 2*TMath::Pi()*B*dB;
		#else
		B  = 0.; dB = 0.; wB = 1.;
		#endif

		h_mult_old->Fill(MpdGlobalTracks->GetEntriesFast(),wB);

		Long_t n_tracks_mpd = 0;
                Long_t n_tracks_for_FlowQA = 0;
		for (Int_t track = 0; track < MpdGlobalTracks->GetEntriesFast(); track++)
		{
			MpdTrack *mpdtrack = (MpdTrack *)MpdGlobalTracks->UncheckedAt(track);

			pt_bin = -1;
			eta_bin = -1;
			if (mpdtrack->GetNofHits() < n_hit_cut)
				continue;

			for (Int_t i_pt = 0; i_pt < n_pt_bin; i_pt++)
				if (TMath::Abs(mpdtrack->GetPt()) >= pt_bins[i_pt] && TMath::Abs(mpdtrack->GetPt()) <= pt_bins[i_pt + 1])
					pt_bin = i_pt;
			for (Int_t i_eta = 0; i_eta < n_eta_bin; i_eta++)
				if (mpdtrack->GetEta() > eta_bins[i_eta] && mpdtrack->GetEta() <= eta_bins[i_eta + 1])
					eta_bin = i_eta;

			if (pt_bin == -1)
				continue;
			if (eta_bin == -1)
				continue;
			Double_t Pt = TMath::Abs(mpdtrack->GetPt());
			TF1 sigma_fit_X = *f_pt_fit[0][eta_bin];
			TF1 sigma_fit_Y = *f_pt_fit[1][eta_bin];
			TF1 sigma_fit_Z = *f_pt_fit[2][eta_bin];

                        if (mpdtrack->GetDCAX()*mpdtrack->GetDCAX()+
                            mpdtrack->GetDCAY()*mpdtrack->GetDCAY()+
                            mpdtrack->GetDCAZ()*mpdtrack->GetDCAZ() < 0.04 && 
                            TMath::Abs(mpdtrack->GetPt()) >= 0.2 && 
                            TMath::Abs(mpdtrack->GetPt()) <= 3 &&
                            TMath::Abs(mpdtrack->GetEta()) <= 1.5)
                          n_tracks_for_FlowQA++;

			if (TMath::Abs(mpdtrack->GetDCAX()) >= sigma_fit_X(Pt) * 2)
				continue;
			if (TMath::Abs(mpdtrack->GetDCAY()) >= sigma_fit_Y(Pt) * 2)
				continue;
			if (TMath::Abs(mpdtrack->GetDCAZ()) >= sigma_fit_Z(Pt) * 2)
				continue;

			if (TMath::Abs(mpdtrack->GetEta()) > 1.5)
				continue;
			if (TMath::Abs(mpdtrack->GetPt()) < 0 && TMath::Abs(mpdtrack->GetPt()) > 3)
				continue;

			n_tracks_mpd++;
		}

		if (n_tracks_mpd > 0)
		{
			h_multiplicity->Fill(n_tracks_mpd,wB);
			h_mult_compare->Fill((Float_t)(n_tracks_mpd * 1.5),wB);
                        h_mult_old    ->Fill(n_tracks_for_FlowQA,wB);
		}
	}

	outFile->cd();
	h_multiplicity->Write();
	//h_mult_compare->Write();
	h_mult_old->Write();
	outFile->Close();
}

int main(int argc, char **argv)
{
	TString iFileName, oFileName, dcaFileName;

	if (argc < 3)
	{
		std::cerr << "./get_multiplicity -i inputfile -o outputfile -dca dcafile" << std::endl;
		return 1;
	}

	for (int i = 1; i < argc; i++)
	{
		if (std::string(argv[i]) != "-i" &&
				std::string(argv[i]) != "-o" &&
				std::string(argv[i]) != "-dca")
		{
			std::cerr << "\n[ERROR]: Unknown parameter " << i << ": " << argv[i] << std::endl;
			return 1;
		}
		else
		{
			if (std::string(argv[i]) == "-i" && i != argc - 1)
			{
				iFileName = argv[++i];
				continue;
			}
			if (std::string(argv[i]) == "-i" && i == argc - 1)
			{
				std::cerr << "\n[ERROR]: Input file name was not specified " << std::endl;
				return 1;
			}
			if (std::string(argv[i]) == "-o" && i != argc - 1)
			{
				oFileName = argv[++i];
				continue;
			}
			if (std::string(argv[i]) == "-o" && i == argc - 1)
			{
				std::cerr << "\n[ERROR]: Output file name was not specified " << std::endl;
				return 1;
			}
			if (std::string(argv[i]) == "-dca" && i != argc - 1)
			{
				dcaFileName = argv[++i];
				continue;
			}
			if (std::string(argv[i]) == "-dca" && i == argc - 1)
			{
				std::cerr << "\n[ERROR]: DCA file name was not specified " << std::endl;
				return 1;
			}
		}
	}

	if (iFileName == "" || oFileName == "" || dcaFileName == "")
	{
		std::cerr << "\n[ERROR]: One of the need file is absent " << std::endl;
		std::cerr << "./get_multiplicity -i inputfile -o outputfile -dca dcafile" << std::endl;
		return 1;
	}
	get_multiplicity(iFileName, oFileName, dcaFileName);
	return 0;
}