#include "TString.h"
#include "TROOT.h"
#include "TSystem.h"
#include "TFile.h"
#include "TTree.h"
#include "TDatabasePDG.h"
#include "TLorentzVector.h"

#include <iostream>
#include <iomanip>
#include <fstream>
#include <vector>
using namespace std;

    struct smash_evt {
	Int_t    evtNr;
        Int_t    nParticle;
	Double_t impact;
	Double_t eBeam;
        Double_t flag;
    };



    struct smash_p {

	Int_t    t;
	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;

        // extended
	Int_t    ncoll;
        Double_t form_time;
	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;

        // helpers
	Int_t    geantID;
	Double_t E;
        Double_t weight;
    };

    struct objects {

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


Int_t convert(TString, TString, Int_t);
Int_t convert_ROOT   (objects& obj);
Int_t convert_OSCAR13(objects& obj);
Int_t write_Header(objects& obj);
Int_t write_Event (objects& obj);
Int_t scan_Config (objects& obj);



Int_t convert(TString in, TString config, TString outf, Int_t nevts)
{
    objects obj;
    obj.evt.eBeam = -1.0;   // for the moment

    obj.debug  = kFALSE;
    obj.input  = in;
    obj.config = config;
    obj.output = outf;
    obj.nEvts  = nevts;

    if(gSystem->AccessPathName(obj.input.Data())){
	cerr<<"ERROR: input file : "<<obj.input<<" not found!"<<endl;
	return 1;
    }

    if(gSystem->AccessPathName(obj.config.Data())){
	cerr<<"ERROR: config file : "<<obj.config<<" not found!"<<endl;
	return 1;
    }

    if(gSystem->AccessPathName(gSystem->DirName(obj.output.Data()))){
	cerr<<"ERROR: output dir : "<<gSystem->DirName(obj.output.Data())<<" not found!"<<endl;
	return 1;
    }

    Int_t ret = scan_Config(obj);
    if(ret != 0) return 1;

    obj.out.open(obj.output.Data());
    if(!obj.out.good()){
	cerr<<"ERROR: Could not open outfile : "<<obj.output<<"!"<<endl;
	return 1;
    }
    obj.out.precision(6);
    obj.out.setf( std::ios::fixed, std:: ios::floatfield );
    if(obj.input.EndsWith(".root")){
	return convert_ROOT(obj);
    } else {
	return convert_OSCAR13(obj);
    }

}

Int_t write_Header(objects& obj)
{
    obj.out<<obj.evt.evtNr<<"\t "<<obj.evt.nParticle<<"\t "<<obj.evt.eBeam<<"\t "<<obj.evt.impact<<"\t "<<obj.evt.flag<<endl;
    return 0;
}

Int_t write_Event(objects& obj)
{
    for(UInt_t i=0; i < obj.particles.size(); i++){
	smash_p& p = obj.particles[i];
	obj.out<<p.E<<"\t "<<p.px<<"\t "<<p.py<<"\t "<<p.pz<<"\t "<<p.geantID<<"\t "<<p.weight<<endl;
    }

    obj.particles.clear();
    return 0;
}

Int_t scan_Config(objects& obj)
{


    ifstream inFile(obj.config.Data());
    if(inFile.good()){
	Char_t line[1000];
	TString l;
	while(!inFile.eof())
	{
	    inFile.getline(line,1000);  // file format and vars decription
	    l = line;
	    if(l.Contains("E_Kin:"))
	    {
		//    E_Kin: 1.23
		Int_t n = sscanf(line,"%*s %lf",&obj.evt.eBeam);
		if(n!=1){
		    cerr<<"ERROR: Could not read E_Kin from config file!"<<endl;
		    return 1;
		} else {
		    cout<<"E_Kin: "<< obj.evt.eBeam <<endl;
		}

		return 0;
	    }
	}
    }
    return 1;
}

Int_t convert_ROOT(objects& obj)
{


    TFile*  f = new TFile(obj.input.Data(),"READ");

    if(!f){
       cerr<<"ERROR: Could not open infile : "<<obj.input.Data()<<"!"<<endl;
       return 1;
    }
    f->cd();

    TTree* T = (TTree*) gDirectory->Get("particles");
    if(!T){
	cerr<<"ERROR: Could not reed TTee \"particles\" from infile : "<<obj.input.Data()<<"!"<<endl;
	return 1;
    }


#define NPAR 1000

    //Declaration of leaves types
    Int_t           npart;
    Double_t        impact_b;
    Int_t           ev;
    Int_t           tcounter;
    Int_t           pdgcode[NPAR];
    Double_t        p0[NPAR];
    Double_t        px[NPAR];
    Double_t        py[NPAR];
    Double_t        pz[NPAR];
    Double_t        t [NPAR];
    Double_t        x [NPAR];
    Double_t        y [NPAR];
    Double_t        z [NPAR];

    // Set branch addresses.
    T->SetBranchAddress("npart",&npart);
    T->SetBranchAddress("impact_b",&impact_b);
    T->SetBranchAddress("ev",&ev);
    T->SetBranchAddress("tcounter",&tcounter);
    T->SetBranchAddress("pdgcode",pdgcode);
    T->SetBranchAddress("p0",p0);
    T->SetBranchAddress("px",px);
    T->SetBranchAddress("py",py);
    T->SetBranchAddress("pz",pz);
    T->SetBranchAddress("t",t);
    T->SetBranchAddress("x",x);
    T->SetBranchAddress("y",y);
    T->SetBranchAddress("z",z);


    Long64_t nentries = T->GetEntries();

    if(nentries > obj.nEvts && obj.nEvts != -1) nentries = obj.nEvts;


    Long64_t nbytes = 0;
    smash_p p;
    TLorentzVector v;
    Int_t evtCt     = 0;
    for (Long64_t i=0; i<nentries;i++) {
	nbytes = T->GetEntry(i);
        if(nbytes<=0) break;

	obj.evt.evtNr     = ev+1;
	obj.evt.nParticle = npart;
	obj.evt.impact    = impact_b;

	for(Int_t j = 0; j < npart; j++){
	    p.t      = (Int_t) t[j];
	    p.x      = x[j];
	    p.y      = y[j];
	    p.z      = z[j];
	    p.p0     = p0[j];
	    p.px     = px[j];
	    p.py     = py[j];
	    p.pz     = pz[j];
	    p.pdg    = (Int_t) pdgcode[j];
	    p.geantID= obj.PDG.ConvertPdgToGeant3(p.pdg);
	    TParticlePDG* pdg =  obj.PDG.GetParticle(p.pdg);
            p.mass   = pdg->Mass();
            p.charge = (Int_t)pdg->Charge();
	    p.weight = 1.;
	    p.ID     = -1;

	    v.SetXYZM(p.px,p.py,p.pz,p.mass);

	    p.E = v.E();

	    obj.particles.push_back(p);
	}
	write_Header(obj);
	write_Event (obj);
	evtCt++;
	if(evtCt == obj.nEvts) {
	    break;
	}
    }

    obj.out.close();
    cout<<"Number of events read from file : "<<setw(6)<<evtCt<<" (requested "<<obj.nEvts<<")"<<endl;

    return 0;
}

Int_t convert_OSCAR13(objects& obj)
{

    Int_t evtCt = 0;
    ifstream inFile(obj.input.Data());
    if(inFile.good()){
        TLorentzVector v;
	smash_p   p;
        Char_t line[1000];


        //-----------------------------------------------------------
        // reading  smash header
	/*
	 #!OSCAR2013 particle_lists t x y z mass p0 px py pz pdg ID charge
	 # Units: fm fm fm fm GeV GeV GeV GeV GeV none none e
	 #
          or  extended format

	 #!OSCAR2013Extended particle_lists t x y z mass p0 px py pz pdg ID charge  ncoll form_time xsecfac proc_id_origin proc_type_origin time_last_coll pdg_mother1 pdg_mother2
         # Units: fm fm fm fm GeV GeV GeV GeV GeV none none e none fm none none none fm none none
	 #
	 */

	inFile.getline(line,1000);  // file format and vars decription
	TString l = line;
	if(!l.BeginsWith("#!OSCAR2013")) {
	    cerr<<"ERROR: input file is not OSCAR2013 format!"<<endl;
            return 1;
	}

	Bool_t extended = kFALSE;

	if(!l.BeginsWith("#!OSCAR2013Extended")) {
	    extended = kTRUE;
	}

	inFile.getline(line,1000);  // units
	inFile.getline(line,1000);  // empty
	//-----------------------------------------------------------


	while(!inFile.eof())
	{
	    //-----------------------------------------------------------
            // reading begin event
	    inFile.getline(line,1000);  // # event 1 out 435
            if(!inFile.good()) break;
	    Int_t n = sscanf(line,"%*s %*s %d %*s %d",&obj.evt.evtNr,&obj.evt.nParticle);
            //-----------------------------------------------------------
	    if(n==2){
		if(obj.debug){
		    cout<<"-------------------------------------------------------------------------------"<<endl;
		    cout<<"Begin event :"<<setw(5)<<obj.evt.evtNr<<setw(4)<<" nParticle : "<<obj.evt.nParticle <<endl;
		}
                Bool_t ok=kTRUE;
		//-----------------------------------------------------------
		// particle eof event
		for(Int_t i = 1; i <= obj.evt.nParticle; i++ ){

		    if(!extended) {
			inFile.getline(line,1000);  //  t x y z mass p0 px py pz pdg ID charge
			n = sscanf(line,"%d %lf %lf %lf %lf %lf %lf %lf %lf %d %d %d "
				   ,&p.t,&p.x,&p.y,&p.z,&p.mass,&p.p0,&p.px,&p.py,&p.pz,&p.pdg,&p.ID,&p.charge);

			if(n!=12){
			    cerr<<"ERROR: expected line is not a particle : "<<line<<endl;
			    ok = kFALSE;
			    break;
			}

		    } else {
			inFile>> p.t >> p.x >> p.y >> p.z >> p.mass >> p.p0 >> p.px >> p.py >>p.pz >> p.pdg >> p.ID >> p.charge
			    >> p.ncoll >> p.form_time >> p.xsecfac >> p.proc_id_origin >> p.proc_type_origin >> p.time_last_coll >> p.pdg_mother1 >> p.pdg_mother2;
			if(inFile.fail()){

 			    cerr<<"ERROR: expected line is not a particle!"<<endl;
			    ok = kFALSE;
			    break;
			}

                        n = 20;
		    }



		    TParticlePDG* pdg =  obj.PDG.GetParticle(p.pdg);
		    p.mass   = pdg->Mass();  // to get the same out put for root/ascii (mass not in root file. smash precision of mass is lower (p 0.938 vs 0.9383 )
		    p.geantID = obj.PDG.ConvertPdgToGeant3(p.pdg);
                    p.weight = 1.;
		    v.SetXYZM(p.px,p.py,p.pz,p.mass);
		    p.E = v.E();
                    obj.particles.push_back(p);


		    if(n==12 || n==20){
			if(obj.debug){
			    cout<<setw(4)<<i<<"    : "
				<<p.t<< " "
				<<setw(12)<<p.x<<" "
				<<setw(12)<<p.y<<" "
				<<setw(12)<<p.z<<" "
				<<setw(12)<<p.mass<<" "
				<<setw(12)<<p.p0<<" "
				<<setw(12)<<p.px<<" "
				<<setw(12)<<p.py<<" "
				<<setw(12)<<p.pz<<" "
				<<setw(6)<<p.pdg<<" "
				<<setw(6)<<p.ID<<" "
				<<setw(2)<<p.charge;
				if(extended){
				    cout<<setw(2)<<p.ncoll<<" "
					<<setw(12)<<p.form_time<<" "
					<<setw(2)<<p.xsecfac<<" "
					<<setw(4)<<p.proc_id_origin<<" "
					<<setw(2)<<p.proc_type_origin<<" "
					<<setw(12)<<p.time_last_coll<<" "
					<<setw(6)<<p.pdg_mother1<<" "
					<<setw(6)<<p.pdg_mother2<<" ";
				}
				cout<<endl;
			}
		    }
		}
		//-----------------------------------------------------------

		if(!ok) break;

		//-----------------------------------------------------------
		// reading end event
		inFile.getline(line,1000);  // example:  # event 1 end 0 impact   6.004
		Int_t ev;
		n = sscanf(line,"%*s %*s %d %*s %*d %*s %lf",&ev,&obj.evt.impact);
		//-----------------------------------------------------------
		if(n==2){
		    if(ev!=obj.evt.evtNr){
			cerr<<"ERROR: end event does not match begin event : begin = "<<obj.evt.evtNr<<" end = "<<ev<<endl;
                        break;
		    }
		    if(obj.debug){
			cout<<"End  event :"<<setw(5)<<ev<<setw(4)<<" impact : "<<obj.evt.impact<<endl;
			cout<<"-------------------------------------------------------------------------------"<<endl;
		    }
                    write_Header(obj);
		    write_Event (obj);
		    evtCt++;
		    if(evtCt == obj.nEvts && obj.nEvts != -1 ) {
                        break;
		    }

		} else {
		    cerr<<"ERROR: expected line is not an end event : "<<line<<endl;
		    break;
		}

	    } else {
		cerr<<"ERROR: expected line is not an begin event : "<<line<<endl;
		break;
	    }
	}


    } else {
	cerr<<"ERROR: could not open : "<<obj.input<<"!"<<endl;
	return 1;
    }

    inFile.close();
    obj.out.close();
    cout<<"Number of events read from file : "<<setw(6)<<evtCt<<" (requested "<<obj.nEvts<<")"<<endl;
    return 0;
}

int main(int argc, char **argv)
{
    TROOT Analysis("Analysis","compiled analysis macro");

    // argc is the number of arguments in char* array argv
    // CAUTION: argv[0] contains the progname
    // argc has to be nargs+1
    if(argc>1) cout<<"nArgs          "<<argc-1<<" convert-smash arguments "<<endl;
    if(argc>1) cout<<"infile       = "<<argv[1]<<endl;
    if(argc>2) cout<<"config file  = "<<argv[2]<<endl;
    if(argc>3) cout<<"outfile      = "<<argv[3]<<endl;
    if(argc>4) cout<<"nevents      = "<<argv[4]<<endl;

    TString number ;
    TString nevts ;
    switch (argc)
    {
    case 4:       //  inputfile config outfile
	return convert(TString(argv[1]),TString(argv[2]),TString(argv[3]),-1);
	break;
    case 5:       // inputfile config outfile  nEvents
        nevts  = argv[4];
	return convert(TString(argv[1]),TString(argv[2]),TString(argv[3]),nevts.Atoi());
	break;
    default:
	cerr<<"USAGE:  convert-smash infile config outfile [nevents]"<<endl;
        cerr<<"convert-smash expects  smash OSCAR13 or OSCAR13Extended ascii format or Root file as input."<<endl;
        cerr<<"Output will be in HGeant ascii evt format."<<endl;
	return 1; // fail
    }
}