#include "KM3SimIO.hh"
#include <stdexcept>
#include <Hit.hh>
#include <Trk.hh>
#include <Evt.hh>
#include <io_ascii.hh>
#include <boost/algorithm/string/predicate.hpp>
#include "G4ParticleTable.hh"
//#include "G4SystemOfUnits.hh"
#include <CLHEP/Units/SystemOfUnits.h>
using namespace CLHEP;

int KM3SimIO::GetNoEventsEvt(std::istream& is){

   int n_events = 0;    
   string w;
   cout<<"Reading number of events..."<<endl;
   while(!is.eof()){
	is>>w;
	if(w == mc_keys::start_event_t && is.good()) n_events += 1;
   }
   cout<<"Number of events to process: "<<n_events<<endl;
   is.clear();
   is.seekg(0);
   return n_events;
   
}


KM3SimIO::KM3SimIO(string infilechar, string outfilechar) {

  EVT_input = false;
  EVT_output= false;
  verbosity = 2;
  header = nullptr;
  evt = nullptr;
  
  if(boost::algorithm::ends_with(outfilechar, ".evt")) EVT_output = true;
  if(boost::algorithm::ends_with(infilechar, ".evt")) EVT_input = true;

  if(EVT_input){
	  header = new Head();
          evt = new Evt();
	  inevt.open(infilechar.c_str(), ifstream::in);
	  if (inevt.fail()){
		THROW(Exception, "The input file does not exist!");  
	  }	  
	  //inevt.seekg(0);
	  read(*header, inevt);
	  //header->print();
  }
  else{
      inputfile = new TFile(infilechar.c_str(), "READ");
      input_tree = (TTree*) inputfile->Get("E");
      input_tree->SetBranchAddress("Evt", &evt);
	  header = (Head*) inputfile->Get("Head");
	  if(!header) header = (Head*) inputfile->Get("Header");
	  //header->print();
  }

  if(EVT_output) outevt.open(outfilechar.c_str(), ofstream::out);  
  else{
      outputfile = new TFile(outfilechar.c_str(), "RECREATE");
      output_tree = new TTree("E", "Evt Tree");
      output_tree->Branch("Evt", &evt, 32000, 4);
  }
  
  //evt = outfile->pevt;

  if(EVT_input) nevents = GetNoEventsEvt(inevt);
  else{
	  nevents = input_tree->GetEntries();
      cout<<"Number of events to process: "<<nevents<<endl;
  }

  //nevents = outfile->size();
  
  evt_id = 0;


  RunHeaderIsWrite = false;
  cout<<"KM3SimIO initialized!"<<endl;
}

KM3SimIO::~KM3SimIO() {
  
  if(EVT_output) outevt.close();
  else{
	  outputfile->Write("", TObject::kOverwrite);
	  outputfile->Close();
  }
  
  if(EVT_input) inevt.close();
  else inputfile->Close(); 
    

}


void KM3SimIO::AddHeaderInfo(string tagname, string tagval) {

    /* old way to write duplicated tags!
    string key = tagname;
    for (int i = 1; header->find(key) != header->end() ; i++)
    {
      key = tagname + "_" + to_string(i);
    }

    (*header)[key] = tagval;
	*/
	
	header->set_line(tagname, tagval);
}

void KM3SimIO::WriteRunHeader() {

  RunHeaderIsWrite = true;
  if(EVT_output) write(*header, outevt);
  else {
     outputfile->WriteObject( header , "Head");
  }
  
}



void KM3SimIO::ReadEvent() {
	
	cout<<"Event "<<evt_id+1<<" being read.."<<endl;
	if(EVT_input){
		read(*evt, inevt);
	}
	else input_tree->GetEntry(evt_id);

	G4int index = 0;
	G4int idbeam = 0;
	light_tracks_indices.clear();

	for(auto &trk: evt->mc_trks){

        	idbeam = trk.type;
			
			if(trk.is_finalstate() && idbeam != 0 && idbeam < 1000000){ //the logic for type is kept, because it's still possible that conversion errors occur in the external software 
				light_tracks_indices.push_back(index);
				//cout<<"Track id: "<<trk.id<<", type: "<<idbeam<<endl; //DEBUG
			}

			// The logic below shows the PDG codes of the particles which are not decayed in Geant4.10 (maybe in the future this will change)
			// I decide to keep them in the input, because they do not slow down the simulation
        	        //else if (fabs(idbeam) == 411 (D+) || fabs(idbeam) == 421 (D0) || fabs(idbeam) == 431 (D+s) ||
        	        //    fabs(idbeam) == 4122 (LAMBDA+c) && fabs(idbeam) == 4212 (SIGMA+c) &&
        	        //    fabs(idbeam) == 4222 (SIGMA++c))
			index +=1;
	}

}


void KM3SimIO::WriteEvent() {

 //cout<<"EventWritten ID: "<<evt_id<<endl; // DEBUG
 if(EVT_output) write(*evt, outevt);
 else output_tree->Fill();
 evt_id+=1;

}

void KM3SimIO::AddHit(int id, int PMTid, double pe, double t, int trackid,
                           int npepure, double ttpure, int creatorProcess) {

  Hit hit;
  hit.id = id;
  hit.pmt_id = PMTid;
  hit.a = pe;
  hit.t = t;
  //hit.pure_t = ttpure; //new format does not have this
  //hit.pure_a = npepure; //new format does not have this
  hit.pattern_flags = creatorProcess;

  
  //G4cout<<"MC tracks size: "<<evt->mc_trks.size()<<G4endl; // DEBUG
  //G4cout<<"Track ID: "<<trackid<<G4endl; // DEBUG
  int Gid;
  int track_id;
  if (trackid <= GetNumberOfParticles()) { //check if the particle is primary or if it's within the number of the light generating particles
	  
	Gid = evt->mc_trks[light_tracks_indices[trackid - 1]].type;
	track_id = evt->mc_trks[light_tracks_indices[trackid - 1]].id;   
	// -1 comes from the fact that in Geant4 TrackID = 0 is reserved for the abstract parent track, the origin of everything
	
  } else {
    Gid = 6;
    trackid = 999999;
  }


  hit.type = Gid;
  hit.origin = track_id;
  evt->mc_trks[light_tracks_indices[trackid - 1]].hit_ids.push_back(id); 
  // -1 comes from the fact that in Geant4 TrackID = 0 is reserved for the abstract parent track, the origin of everything
  evt->mc_hits.push_back(hit);

}

void KM3SimIO::AddTrackLength(int trackid, double track_L){

  evt->mc_trks[light_tracks_indices[trackid - 1]].len =  track_L / m ;
  // -1 comes from the fact that in Geant4 TrackID = 0 is reserved for the abstract parent track, the origin of everything

}


//////////////////////
//		    //
// READER FUNCTIONS //
//		    //
//////////////////////

int KM3SimIO::GetNumberOfParticles(void) {
    return light_tracks_indices.size(); //this corresponds to the number of particles in the primary vertex. Incoming neutrino, the target and recoil nuclei are not taken into account in the simulation
}


void KM3SimIO::GetParticleInfo(int &idbeam, double &xx0, double &yy0,
                                   double &zz0, double &pxx0, double &pyy0,
                                   double &pzz0, double &t0) {
  G4cout<<"Do nothing for now, but maybe will be useful in the km3sim Multi Threading implementation"<<G4endl;
}

bool KM3SimIO::GetNeutrinoInfo(int &idneu, int &idtarget, double &xneu,
                                   double &yneu, double &zneu, double &pxneu,
                                   double &pyneu, double &pzneu, double &t0) {
  idneu = 0;
  idtarget = 0;
  xneu = 0.0;
  yneu = 0.0;
  zneu = 0.0;
  pxneu = 0.0;
  pyneu = 0.0;
  pzneu = 0.0;
  //if (!isneutrinoevent) return;
  Trk neutrino = evt->mc_trks[0];
  if(!neutrino.is_neutrino()) return false;
  idneu = neutrino.type;
  //if (argnumber == 15)
  //  idtarget = (int)args[14];
  //else
  idtarget = 1445;
  xneu = neutrino.pos.x * CLHEP::m / CLHEP::cm; // convert to cm
  yneu = neutrino.pos.y * CLHEP::m / CLHEP::cm;
  zneu = neutrino.pos.z * CLHEP::m / CLHEP::cm;
  double pmom = neutrino.E;
  pxneu = neutrino.dir.x * pmom;
  pyneu = neutrino.dir.y * pmom;
  pzneu = neutrino.dir.z * pmom;
  t0 = neutrino.t;

  if(verbosity>2) G4cout<<"Neutrino info: x:"<<xneu<<" y:"<<yneu<<" z: "<<zneu<<" Px:"<<pxneu<<" Py:"<<pyneu<<" Pz:"<<pzneu<<endl;
  return true;
}


// following is for hepevt interface
#include "G4PrimaryVertex.hh"
#include "G4PrimaryParticle.hh"
//#include "G4ThreeVector.hh"

//////////////////////////////////

void KM3SimIO::GeneratePrimaryVertex(G4Event *anEvent) {

    int idbeam;
    double xx0, yy0, zz0;
    double pxx0, pyy0, pzz0;
    double t0;
	

    //cout<<"Number of tracks for event: "<<evt->mc_trks.size()<<endl; //DEBUG
    for(auto &trk: evt->mc_trks){


	  idbeam = trk.type;
      
	  if( !trk.is_finalstate() || idbeam == 0 || idbeam > 1000000) continue;


      //if (fabs(idbeam) == 411 (D+) || fabs(idbeam) == 421 (D0) || fabs(idbeam) == 431 (D+s) ||
      //    fabs(idbeam) == 4122 (LAMBDA+c) && fabs(idbeam) == 4212 (SIGMA+c) &&
      //    fabs(idbeam) == 4222 (SIGMA++c))
      // these particles are not defined or have no decay modes in GEANT4
      // they are kept in the simulation, because they don't do any harm and their occurace in theorey can be counted




      //cout<<"Track type: "<<idbeam<<endl; //DEBUG
      xx0 = trk.pos.x * m ; // convert to cm
      yy0 = trk.pos.y * m ;
      zz0 = trk.pos.z * m ;
	  t0 = trk.t * ns;

		

      double totenergy = trk.E;
      G4ParticleDefinition* definition = G4ParticleTable::GetParticleTable()->FindParticle(idbeam);
      double pmass = definition->GetPDGMass() / GeV;
      //cout<<"PMass :"<<pmass<<" GeV"<<endl; //DEBUG
      double pmom = totenergy * totenergy - pmass * pmass;
      if (pmom > 0.0)
          pmom = sqrt(totenergy * totenergy - pmass * pmass);
      else
          pmom = 0.0;    
      //G4cout<<"Pmom is: "<<pmom<<G4endl; //DEBUG

      pxx0 = trk.dir.x * pmom;
      pyy0 = trk.dir.y * pmom;
      pzz0 = trk.dir.z * pmom;

      G4PrimaryVertex *vertex = new G4PrimaryVertex(G4ThreeVector(xx0, yy0, zz0), t0);
      G4PrimaryParticle *particle = new G4PrimaryParticle(idbeam);
      particle->SetMomentum(pxx0 * GeV, pyy0 * GeV, pzz0 * GeV);
      vertex->SetPrimary(particle);
	  // Put the vertex to G4Event object
      anEvent->AddPrimaryVertex(vertex);
    }

}