#ifndef  ___STRUCTS_HELPERS___
#define  ___STRUCTS_HELPERS___

#include "TVector3.h"
#include "TLorentzVector.h"
#include "TDatabasePDG.h"
#include "TParticlePDG.h"
#include "TDecayChannel.h"
#include "TObjArray.h"
#include "TString.h"
#include "TFile.h"
#include "TTree.h"
#include "TClonesArray.h"
#include <iostream>
#include <iomanip>
#include <vector>
#include <map>
#include <algorithm>
#include <fstream>
#include "pid_converters.h"

using namespace std;

//######################## INPUT PARAMETERS ##############################
struct parameters {

    Int_t   verbose_flag;
    TString infile;
    TString outfile;
    Int_t   nevents;
    Int_t   help;
    Int_t   doIncludeSpectators;
    Int_t   doRecoRho;
    Int_t   doRecoOmega;
    Int_t   doRecoPhi;
    Int_t   doSecondDecay;
    Int_t   doPluto;
    Int_t   doTree;
    Int_t   doBoost;
    TString boostConfig;
    UInt_t  seed;
    TString databasefile;

    void print()
    {
	cerr<<"current parameters :"<<endl;
	cerr<<"infile              : "<<infile<<endl;
	cerr<<"outfile             : "<<outfile<<endl;
	cerr<<"nevents             : "<<nevents<<endl;
	cerr<<"verbose             : "<<verbose_flag<<endl;
	cerr<<"doIncludeSpectators : "<<doIncludeSpectators<<endl;
	cerr<<"doRecoRho           : "<<doRecoRho<<endl;
	cerr<<"doRecoOmega         : "<<doRecoOmega<<endl;
	cerr<<"doRecoPhi           : "<<doRecoPhi<<endl;
	cerr<<"doSecondDecay       : "<<doSecondDecay<<endl;
	cerr<<"doPluto             : "<<doPluto<<endl;
	cerr<<"doTree              : "<<doTree<<endl;
	cerr<<"doBoost             : "<<doBoost<<endl;
	cerr<<"seed                : "<<seed<<endl;
	cerr<<"databasefile        : "<<databasefile<<endl;
    }


};

static parameters params;
//########################################################################

//####################### EVENT STRUCTURE ###################################

struct smash_evt {

    Int_t    evtNr;
    Int_t    nParticle;
    Double_t impact;
    Double_t eBeam;
    Double_t gamma;
    Double_t beta;
    Int_t    flag;

    //------------------------------------------------------
    // Collision information for BOOST CM -> LAB
    // (provived by the -b option string)

    Double_t calcGamma(Double_t beta) { return sqrt(1 - beta*beta); }
    Double_t calcBeta (Double_t Ekin=1.23) // Ekin/u [GeV],
    {
        const Double_t mP = 0.938272, mN = 0.939565; // GeV
        Double_t m = (mP+mN)/2.;
	return sqrt(pow(Ekin + m, 2.) - pow(m, 2.)) / (Ekin + 2*m);
    }

    //------------------------------------------------------

    void  reset(){
	evtNr = nParticle =0;
        impact=eBeam=gamma=beta=flag=0;
    }

    void print(){
	cout<<"-------------------------------------------------------------------"<<endl;
        cout<<"evtNr "<<setw(6)<<evtNr<<" NPart "<<setw(5)<<nParticle<<" impact "<<setw(12)<< impact<<" eBeam "<<setw(12)<<eBeam<<" flag "<<flag<<endl;
    }
};



struct smash_p {

    Int_t    t;
    Double_t turqmd;
    Double_t x;
    Double_t y;
    Double_t z;
    Double_t mass;
    Double_t p0;
    Double_t px;
    Double_t py;
    Double_t pz;
    Int_t    pdg;
    Int_t    ID;
    Int_t    charge;
    Int_t    iso3;

    // extended
    Int_t    ncoll;
    Double_t form_time;
    Double_t form_time_org;
    Int_t    xsecfac;
    Int_t    proc_id_origin;
    Int_t    proc_type_origin;
    Double_t time_last_coll;
    Int_t    pdg_mother1;
    Int_t    pdg_mother2;
    Int_t    baryon_number;

    // helpers
    Int_t    geantID;
    Double_t E;
    Double_t weight;
    Bool_t   isUsed;
    Bool_t   isNucleon;
    Int_t    genInfo1;
    Int_t    genInfo2;
    Int_t    genInfo3;
    Int_t    state;  // 1:  smash,  -1: smash suppressed, 2: RECO, -2: RECO suppressed; 3: PLUTO  -3 : PLUTO suppressed
    TLorentzVector v;

    void reset(){

	t = 0;
	turqmd = 0;
	x = y = z = mass = p0 =  px =  py = pz = 0.;
	pdg = ID = charge = iso3 = 0;

	// extended
	ncoll = form_time = xsecfac = proc_id_origin =  proc_type_origin = time_last_coll = pdg_mother1 = pdg_mother2 = baryon_number = 0;

	// helpers
	geantID = 0;
	E = weight = 0;
	isUsed = isNucleon = kFALSE;
	v.SetXYZM(0,0,0,0);

        genInfo1 = genInfo2 = genInfo3 = 0;
        state = 1;

    }

    void print(Int_t i) {

	cout<<setw(4)
	    <<i
	    <<" state "  <<setw(2)<<state
	    //<<" t "<<t
	    <<" : x "    <<setw(12)<<x
	    <<" y "      <<setw(12)<<y
	    <<" z "      <<setw(12)<<z
	    <<" mass "   <<setw(12)<<mass
	    <<" p0 "     <<setw(12)<<p0
	    <<" px "     <<setw(12)<<px
	    <<" py "     <<setw(12)<<py
	    <<" pz "     <<setw(12)<<pz
	    <<" pdg "    <<setw(12)<<pdg
	    <<" geantId "<<setw(6) <<geantID
	    <<" ID "     <<setw(6) <<ID
	    <<" chrg "   <<setw(2) <<charge
	    <<" ncoll "  <<setw(2) <<ncoll
	    <<" time "   <<setw(12)<<form_time
	    //<<setw(2)<<xsecfac<<" "
	    //<<setw(4)<<proc_id_origin<<" "
	    //<<setw(2)<<proc_type_origin<<" "
	    //<<setw(12)<<time_last_coll<<" "
	    <<" moth1 "  <<setw(8) <<pdg_mother1
	    <<" moth2 "  <<setw(8) <<pdg_mother2
	    <<endl;
            //v.Print();
    }
};

Bool_t sortFormTime(smash_p& a, smash_p& b) {
    if(a.state < 0 && b.state > 0 ) return kFALSE;
    if(b.state < 0 && a.state > 0 ) return kTRUE;
    return (a.form_time) < (b.form_time);
}
Bool_t sortState   (smash_p& a, smash_p& b) { return (a.state) > (b.state); }


//########################################################################


struct mapDecay {

    map<Int_t,vector<vector<Int_t> > > mPDGDecayList;
    map<Int_t,vector<Int_t> >          mPDGDecayListStat;
    map<Int_t,vector<Int_t> >          mPDGDecayListID;
    map<Int_t,vector<vector<Int_t> > >::iterator it;
    map<Int_t,vector<Int_t> >::iterator it2;
    vector<Int_t> vIds;

    Bool_t compDaughters(const vector<Int_t>& d1,const vector<Int_t>& d2)
    {
	if(d1.size() != d2.size()) return kFALSE;
	for(UInt_t i = 0 ; i < d1.size(); i++){
	    if(find(d2.begin(),d2.end(),d1[i]) == d2.end()) return kFALSE;
	}
        return kTRUE;
    }


    Bool_t findDecay(Int_t pdg,UInt_t& ich,Int_t& srcID, const vector<Int_t>& vdaughters,Bool_t count=kTRUE)
    {
	ich   =  99;
        srcID = -1;
	it = mPDGDecayList.find(pdg);
	if(it == mPDGDecayList.end()) return kFALSE;

	vector<vector<Int_t> >& vdecays = it->second;
	for(UInt_t i = 0; i < vdecays.size(); i++){
	    vector<Int_t>& d = vdecays[i];
	    if(compDaughters(vdaughters,d)){
                if(count)mPDGDecayListStat[pdg][i]++;
                srcID = mPDGDecayListID[pdg][i];
		ich = i;
		return kTRUE;
	    }
	}
        return kFALSE;
    }

    const vector<Int_t>& getDecayDaughters(Int_t pdgID,UInt_t ich){
         return mPDGDecayList[pdgID][ich];
    }

    Int_t getDecayStat(Int_t pdgID,UInt_t ich){
         return mPDGDecayListStat[pdgID][ich];
    }

    Int_t getDecaySourceID(Int_t pdgID,UInt_t ich){
         return mPDGDecayListID[pdgID][ich];
    }

    void printDecay(TDatabasePDG& PDG,Int_t id,Int_t ich){

	it = mPDGDecayList.find(id);
	if(it == mPDGDecayList.end()) return;

	it2 = mPDGDecayListID.find(id);

	vector<Int_t>& d = mPDGDecayList[id][ich];
	cout<<"Mother PDG "<<setw(6)<<id<<" ("<<setw(4)<<PDG.GetParticle(id)->GetName()<<") "<<endl;
	cout<<"channel  "<<setw(2)<<ich<<" : "<<flush;

	TString msg1;
	TString msg2;
	for(UInt_t j = 0; j < d.size(); j++){
	    msg1+= Form("%4d " ,d[j]);
	    msg2+= Form("%6s ",PDG.GetParticle(d[j])->GetName());
	}
	cout<<msg1<<"  ("<<msg2<<")  ,  source ID = "<<it2->second[ich]<<endl;
    }

    void printDecayMap(TDatabasePDG& pdg)
    {
	cout<<"###########################################################################"<<endl;
	cout<<"DECAY MAP SMASH:"<<endl;
	for(it = mPDGDecayList.begin(); it != mPDGDecayList.end(); ++ it){
	    cout<<"---------------------------------------------------------------------------"<<endl;
	    cout<<"Mother PDG "<<setw(6)<<it->first<<" ("<<setw(4)<<pdg.GetParticle(it->first)->GetName()<<") "<<endl;
	    vector<vector<Int_t> >& vdecays = it->second;

	    it2 = mPDGDecayListID.find(it->first);
            vector<Int_t>& vID = it2->second;
	    for(UInt_t i = 0; i < vdecays.size(); i++){
		vector<Int_t>& d = vdecays[i];
		cout<<"channel  "<<setw(2)<<i<<" : "<<flush;

		TString msg1;
		TString msg2;
		TString msg7 ="       ";
		TString msg5 ="     ";
		TString msg0 ="";

		for(UInt_t j = 0; j < d.size(); j++){
		    msg1+= Form("%4d " ,d[j]);
		    msg2+= Form("%6s ",pdg.GetParticle(d[j])->GetName());
		}
		it2 = mPDGDecayListStat.find(it->first);
		if(d.size() == 2)cout<<msg1<<msg5<<"  ("<<msg2<<") ,"<<msg7<<" source ID = "<<setw(10)<<vID[i]<<"  Stats : "<<it2->second[i]<<endl;
		else             cout<<msg1<<msg0<<"  ("<<msg2<<") ,"<<msg0<<" source ID = "<<setw(10)<<vID[i]<<"  Stats : "<<it2->second[i]<<endl;
	    }
	}
	cout<<"###########################################################################"<<endl;
    }

    Bool_t fillDecayMap(TDatabasePDG& PDG,Int_t pid)
    {
	TParticlePDG* p = PDG.GetParticle(pid);
	if(p){
	    vector<vector<Int_t> > vdecays;
	    vector<Int_t>          vstat;
	    vector<Int_t>          vid;
	    vector<Int_t>          vdgeant;

	    TObjArray* a = p->DecayList();
	    for(Int_t i = 0 ; i < a->GetEntries(); i++){
		TDecayChannel* ch = (TDecayChannel*) a->At(i);
		vector<Int_t> daughters;
		for(Int_t j = 0 ; j < ch->NDaughters(); j ++){
		    Int_t pid_daughter = ch->DaughterPdgCode(j);
		    daughters.push_back(pid_daughter);
		}

                vdgeant.clear();
		Int_t geantID_mother = getPdgPIDtoGeant3(PDG,pid,kTRUE);
		for(UInt_t j = 0 ; j < daughters.size(); j++) vdgeant.push_back(getPdgPIDtoGeant3(PDG,daughters[j], kTRUE));
		Int_t sourceID = createSourceID(geantID_mother,vdgeant);

		vdecays.push_back(daughters);
                vstat  .push_back(0);
                vid    .push_back(sourceID);
	    }
	    mPDGDecayList    [pid] = vdecays;
	    mPDGDecayListStat[pid] = vstat;
	    mPDGDecayListID  [pid] = vid;
	} else {
	    cout<<"fillDecayMap() : ID = "<<pid<<" is not a PDG code!"<<endl;
	    return kFALSE;
	}
	return kTRUE;
    }


    void addDecay(TDatabasePDG& PDG,Int_t pid,vector<Int_t>& d)
    {
        // adds the new decay channel (does not check if the channel exists already)
        vector<Int_t> vgeant;
	Int_t geantID_mother = getPdgPIDtoGeant3(PDG,pid,kTRUE);
	for(UInt_t j = 0 ; j < d.size(); j++) vgeant.push_back(getPdgPIDtoGeant3(PDG,d[j], kTRUE));
	Int_t sourceID = createSourceID(geantID_mother,vgeant);

	it = mPDGDecayList.find(pid);
	if(it == mPDGDecayList.end()) {
             vector<vector<Int_t> > vdecays;
	     vector<Int_t> daughters = d;
	     vector<Int_t> vstat;
	     vector<Int_t> vid;
	     vstat  .push_back(0);
	     vid    .push_back(sourceID);
	     vdecays.push_back(daughters);
	     mPDGDecayList    [pid] = vdecays;
	     mPDGDecayListStat[pid] = vstat;
	     mPDGDecayListStat[pid] = vid;
	} else {
	    vector<Int_t> daughters = d;
	    vector<vector<Int_t> >& vdecays = it->second;
	    it2 = mPDGDecayListStat.find(pid);
            it2->second.push_back(0);
	    it2 = mPDGDecayListID.find(pid);
	    it2->second.push_back(sourceID);
	    vdecays.push_back(daughters);
	}
    }

    void removeDecay(Int_t pid,UInt_t ich)
    {
        // remove the decay channel ich
	it = mPDGDecayList.find(pid);
	if(it == mPDGDecayList.end()) return;
	vector<vector<Int_t> >& vdecays = it->second;
	if(ich < vdecays.size()) {
	    vdecays.erase(vdecays.begin()+ich);
	    it2 = mPDGDecayListStat.find(pid);
	    it2->second.erase(it2->second.begin()+ich);
	    it2 = mPDGDecayListID.find(pid);
	    it2->second.erase(it2->second.begin()+ich);
	}
    }

    void init(TDatabasePDG& PDG){

	//vIds.push_back( 111); // pi0
	//vIds.push_back( 221); // eta
	vIds.push_back( 113); // rho0
	vIds.push_back( 213); // rho+
	vIds.push_back(-213); // rho-
	vIds.push_back( 223); // w
	vIds.push_back( 333); // phi

	for(UInt_t i = 0 ; i < vIds.size(); i++){
            fillDecayMap(PDG,vIds[i]);
	}

	Int_t chw[3][2] = {
	    { 113, 111}, // rho0 pi0
	    { 213,-211}, // rho+ pi-
	    {-213, 211}, // rho- pi+
	};
        vector<Int_t> d;
	for(Int_t i = 0; i < 3; i++){
	    d.clear();
	    for(Int_t j = 0; j < 2; j++){
		d.push_back(chw[i][j]);
	    }
            addDecay(PDG,223,d);
	}

	Int_t chphi [2] = {
	    311, -311 //K0 K0bar
	};

        d.clear();
	for(Int_t j = 0; j < 2; j++){
	    d.push_back(chphi[j]);
	}
	addDecay(PDG,333,d);



	printDecayMap(PDG);

    }

};

//###################### COLLECT GLOABAL OBJECTS #########################
struct objects {

    TDatabasePDG    PDG;
    smash_evt       evt;
    vector<smash_p> particles;
    ofstream        out;
    TString         input;
    TString         output;
    Int_t           nEvts;
    Int_t           debug;

    TFile*        PLUTO_out  ;
    TTree*        PLUTO_tree ;
    TClonesArray* PLUTO_evt  ;
    Int_t         PLUTO_npart;
    Double_t      PLUTO_impact;
    Double_t      PLUTO_evtplane;

    map<Int_t,TString> mProcID;

    mapDecay mDecay;


    void init(){

	mProcID[0 ]="No previous process yet, particle was created at initialization";
	mProcID[1 ]="Elastic scattering";
	mProcID[2 ]="Resonance formation (2 -> 1)";
	mProcID[3 ]="Inelastic binary scattering (2 -> 2)";
	mProcID[4 ]="Inelastic multi-particle scattering (2 -> 3)";
	mProcID[5 ]="Resonance decay";
	mProcID[6 ]="Box wall crossing (due to periodic boundary conditions)";
	mProcID[7 ]="Forced thermalization, many particles are replaced by a thermalized ensemble";
	mProcID[8 ]="Hypersurface crossing, Particles are removed from the evolution and printed to a separate output to serve as initial conditions for hybrid models.";
	mProcID[9 ]="Bremsstrahlung process: a + b -> a + b + photon";
	mProcID[10]="Inelastic multi-particle meson scattering (3 -> 1)";
	mProcID[11]="Inelastic multi-particle scattering (3 -> 2)";
	mProcID[12]="Inelastic multi-particle scattering (5 -> 2)";
	mProcID[13]="Inelastic multi-particle scattering (2 -> 5)";
	mProcID[14]="Inelastic multi-particle scattering (4 -> 2)";
	mProcID[15]="Inelastic multi-particle scattering (2 -> 4)";
	mProcID[41]="Soft string excitation, single diffractive AB -> AX. Both quark and anti-/di-quark taken from B.";
	mProcID[42]="Soft string excitation, single diffractive AB -> XB. Both quark and anti-/di-quark taken from A. It makes sense to distinguish it from AB -> AX, because A and B can be particles of different types, for example, a pion and a proton. It matters then, whether the pion creates a string or the proton.";
	mProcID[43]="Soft string excitation, double diffractive. Two strings are formed, one from A and one from B.";
	mProcID[44]="Soft string N-Nbar annihilation, a special case of baryon-antibaryon annihilation. One pair qqbar annihilates immediately and then two strings are formed.";
	mProcID[45]="Soft string excitation, non-diffractive. Two strings are formed both have ends in A and B.";
	mProcID[46]="Hard string excitation, hard string process involving 2 -> 2 QCD process by PYTHIA. Here quarks do not simply form a string. They actually scatter on parton level first.";
	mProcID[47]="Failed string process, Soft String NNbar annihilation process can fail by lack of energy. This is a tag we add to avoid mislabeling the events.";
	mProcID[90]="Add or remove particle(s) process, which ignores conservation laws. It can be thought of as a 0 -> 1 or a 1 -> 0 process.";

	mDecay.init(PDG);
    }


    void printEvent(TString message) {

        evt.print();

        cout<<"NPart " <<particles.size()<<" "<<message.Data()<<endl;
	for(UInt_t i=0; i< particles.size();i++){
	    particles[i].print(i);
	}
    }
    void boost(){
	for(UInt_t i = 0; i < particles.size(); i++){
	    smash_p& p = particles[i];

	    //-----------------------------------------------------------
	    // boost from center of velocity to lab sys
	    TLorentzVector vc;
	    vc.SetPxPyPzE(p.px,p.py,p.pz,p.p0);
            vc.Boost(0.,0.,evt.beta);

            p.pz = vc.Pz();
            p.p0 = vc.E();
            p.E  = vc.E();

            p.v.SetPz(p.pz);
            p.v.SetE(p.E);
	    //-----------------------------------------------------------

	}
    }
};

static objects obj;

//########################################################################

#endif