#include <stdio.h>
#include <stdlib.h>
#include <iostream>

//ROOT headers
#include <TSystem.h>
#include <TSystemFile.h>
#include <TDirectory.h>
#include <TFile.h>
#include <TTree.h>
#include <TObjString.h>
//#include <TObjectTable.h>

//G4VMC headers
#include <TGeant4.h>

//VMC headers
#include "PamVMCApplication.h"
#include "PamVMCOptMgr.h"
//This is main VMC routine
using namespace std;

#define DEBUG 0

void PackDump(TString path, TObjString* os);
void UnPackDump(TString path, TObjString* os);

int main(int argc, char **argv)
{

  if(argc < 2) {
    cout<<"ERROR: PamVMCmain - no input file specified"<<endl;
    abort();
  }
  //Option1 montecarlo selection,  default G4
  gSystem->Load("libg4root.so");
  cout<<" Input card file: "<<argv[1]<<endl;

  PamVMCOptMgr * omg = PamVMCOptMgr::Instance();
  omg->ParseFile(argv[1]);
  omg->Print();

  //Checking if cookie and temp_path
  TString working_path = omg->GetOutPath();
  TString cookie = "";
  if( omg->IsUseCookie() ){
    cookie = gSystem->Getenv("PBS_JOBCOOKIE");
    working_path = omg->GetTmpPath();
    working_path += "/";
    working_path += cookie;
  }
  TSystemFile *ff = new TSystemFile();
  if( !ff->IsDirectory(working_path)) gSystem->mkdir(working_path);
  delete ff;

  //if I'm going to store random stuff I need temporary directory to dump g4 engine status
  TString tmpseeddir;
  Int_t save_ev;
  TObjString* ev_dump = NULL;
  TFile * dmp_file = NULL;
  TTree * t_dump = NULL;
  tmpseeddir = working_path;
  tmpseeddir += "/";
  tmpseeddir += omg->GetOutPattern();
  tmpseeddir+="_g4rndm_tmp";
  if (omg->DoWriteRandom()){
  
    cout<<"Making  Geant-4 RNDM tmp directory: "<<tmpseeddir<<endl;
    gSystem->mkdir(tmpseeddir);

    //Changing directory
   
    gSystem->Exec("cd "+working_path);
    gSystem->cd(working_path);
    //gDirectory->cd(working_path);
    
    //Making files and trees for saved random statuses
    
    ev_dump = new TObjString();
    TString file_dump_name_tmp = working_path;
    file_dump_name_tmp += "/";
    file_dump_name_tmp += omg->GetRandomOutFilename();
    dmp_file = new TFile(file_dump_name_tmp,"recreate");
    t_dump = new TTree("dumps","G4 rndm for events");
    t_dump->Branch("evno", &save_ev,"save_ev/I");
    t_dump->Branch("dump_cont","TObjString", &ev_dump);
  }

  //INITIALIZATION OF MC
  TString out_tmp_fname = working_path;
  out_tmp_fname += "/";
  out_tmp_fname += omg->GetOutPattern();
  cout<<out_tmp_fname<<endl;
  PamVMCApplication* appl 
    = new PamVMCApplication("PAMG4_VMC", "PAMELA GEANT4 VMC application",out_tmp_fname,0);
  //setting verbose level
  appl->SetVerboseLevel( omg->GetVerboseLevel() );
  //initialize physycs with specific RunConfiguration, defined in macro
  appl->InitMC( omg->GetG4ConfigC() );
  cout<<"MACRO READED"<<endl;
  //additional flags and controls related to G4 core
  ((TGeant4*)gMC)->ProcessGeantMacro( omg->GetG4ConfigIn() );
  //set maximum allowed steps
  ((TGeant4*)gMC)->SetMaxNStep( omg->GetMaxSteps() );
 
  //generate initial random seed if any
  TFile * read_rndm_file = NULL;
  TTree * read_rndm_tree = NULL;
  Int_t inpev;
  TObjString * rnd_dump = NULL;
  TString tmpfile;
  if ( omg->IsReadRandomMode() ){
    //read random seed from file
    cout<<"Making  Geant-4 RNDM tmp directory: "<<tmpseeddir<<endl;
    gSystem->mkdir(tmpseeddir);
    tmpfile = tmpseeddir+"/currEvent.rndm";
    gSystem->Exec("cd "+working_path);
    gSystem->cd(working_path);
    read_rndm_file = new TFile( omg->GetRandomInFilename() );
    read_rndm_tree = (TTree*)read_rndm_file->Get("dumps");
    Int_t inpev;  
    read_rndm_tree->SetBranchAddress("evno", &inpev);
    read_rndm_tree->SetBranchAddress("dump_cont", &rnd_dump);
  } else {
    Int_t seed1, seed2;
    if( omg->IsUserRandomMode() ){
      omg->GetUserRandomSeed( seed1, seed2 );
    } else {
      TRandom3 rinit1(0);
      TRandom3 rinit2(0);
      seed1 = rinit1.Integer(numeric_limits<int>::max());
      seed2 = rinit2.Integer(numeric_limits<int>::max());
    }
    cout<<"Initial RND seeds for G4 core are: " << seed1<<" , " << seed2<<endl;
    TString command = "/random/setSeeds ";
    command += seed1;
    command += " ";
    command += seed2;
    ((TGeant4*)gMC)->ProcessGeantCommand(command);
  }

  /*seting tmp directory to save status of random gen if any*/
  if( omg->DoWriteRandom() ){
    ((TGeant4*)gMC)->ProcessGeantCommand("/random/setDirectoryName "+tmpseeddir);
  }
 

  Double_t x,y,z,theta,phi,momentum;
  Int_t id,pdg;
  Double_t xl, yl, zl, xh, yh, zh, tl, th, pl, ph;
  Int_t nev_max = omg->GetEvMax();


  //* if I have to read primary kinematics from a file do it *//
  TFile* rfin = NULL;
  TTree* primtree = NULL;
  TClonesArray * prims = NULL;
  PamVMCPrimary *p = NULL;
  if ( omg->GetReadPrimMode() ){
    rfin = new TFile ( omg->GetReadPrimPath(), "read");
    primtree = (TTree*) rfin->Get("hits");
    if(!primtree) cout<<"ERROR: tree 'hits' with primaries kinematics not found"<<endl;
    primtree->SetBranchAddress("PRIM", &prims);
    nev_max = primtree->GetEntries();
  }

  cout<<"OK1"<<endl;

  //MAIN LOOP
  Int_t proc_ev = 0;
  Int_t accept_ev = 0;
  Int_t trig_ev = 0;
  Bool_t stop = kFALSE;
  TRandom3 * aux_rnd = new TRandom3(0);
  PamVMCPrimaryGenerator* gpr = appl->GetPrimaryGenerator();
  gpr->SetRandom(aux_rnd);
 
  //gObjectTable->Print();
  gSystem->Exec(" ps aux | grep PamVMC.exe" );
  while (!stop){
   
   
    if ( omg->GetReadPrimMode() ){
      //read kinematics from file
      rfin->cd();
      primtree->GetEntry(proc_ev);
      p = (PamVMCPrimary*)prims->At(0);
      x = p->fX0;
      y = p->fY0;
      z = p->fZ0;
      theta = p->fTHETA;
      phi = p->fPHI;
      pdg = p->fPDG;
      momentum = p->fP0;
      id = p->fID;
	
   
      gpr->SetParticle(pdg);
      gpr->SetParticleID(id);
      gpr->SetDirection(theta,phi);
      gpr->SetMomentum(momentum);
      gpr->SetPosition(x,y,z);

    } else {
      //generate kinematics
   
      pdg = omg->GetPrimPDG();
      id = proc_ev;
      gpr->SetParticle(pdg);
      gpr->SetParticleID(id);
      
      //generate spatial distribution
      if( omg->IsCustomKinematics() ){
	//vertex
	if( omg->IsBoxVertex() ){
	  //box
	  omg->GetBoxVertex(xl,xh,yl,yh,zl,zh);
	  gpr->GenPosition(xl, xh, yl, yh, zl, zh);
	} else {
	  //point
	  omg->GetFixedPointVertex(x,y,z);
	  gpr->SetPosition(x,y,z);
	}
	//angles
	if( omg->IsIsotropicPlane() ){
	  //isoplane angles
	  omg->GetIsotropicAngles(tl, th, pl, ph);
	  gpr->GenDirection(tl, th, pl, ph);
	} else {
	  //fixed angles
	  omg->GetFixedAngles(theta, phi);
	  gpr->SetDirection(theta,phi);
	}
      } else {
	// GOD mode
	gpr->Gen2PiPosAng();
      }
      gpr->GetPosition(x,y,z);
      gpr->GetDirection(theta, phi);
      

      //generate MOMENTUM
      Momentum_units mu = omg->GetPrimMomUnits();
      if( omg->GetPrimMomMode() == FIXED ){
	momentum = omg->GetMomFixed();

	switch (mu){
	case GeV_T:
	  gpr->SetMomentum(gpr->KinEToMomentum(momentum));
	  break;
	case GV_p:
	  gpr->SetMomentum(momentum);
	case GV_R:
	  gpr->SetMomentum( gpr->RigToMomentum(momentum));
	default:
	  break;
	}
	
      }
      if( omg->GetPrimMomMode() == FLAT ){
	Double_t m1, m2;
	omg->GetMomLimits(m1,m2);
	switch (mu){
	case GeV_T:
	  gpr->GenSpe( m1, m2, kTRUE);
	  break;
	case GV_p:
	  gpr->GenSpe( m1, m2, kFALSE);
	case GV_R:
	  gpr->GenSpe( gpr->RigToMomentum(m1), gpr->RigToMomentum(m2), kFALSE);
	default:
	  break;
	}

      }
      if( omg->GetPrimMomMode() == POWERLAW ){
	Double_t m1, m2;
	omg->GetMomLimits(m1,m2);
	Double_t gamma = omg->GetMomGamma();
	switch (mu){
	case GeV_T:
	  gpr->GenSpe( m1, m2, gamma, kTRUE);
	  break;
	case GV_p:
	  gpr->GenSpe( m1, m2, gamma, kFALSE);
	case GV_R:
	  gpr->GenSpe( gpr->RigToMomentum(m1), gpr->RigToMomentum(m2), gamma, kFALSE);
	default:
	  break;
	}
      }

    }

    /*if (DEBUG){
      cout<<"EVENT: "<<proc_ev<<endl;
      cout<<"ID: "<<id<<endl;
      cout<<"PDG:"<<pdg<<endl
	  <<"X0: "<<x<<" (cm)"<<endl
	  <<"Y0: "<<y<<" (cm)"<<endl
	  <<"Z0: "<<z<<" (cm)"<<endl
	  <<"Theta:"<<theta*180./TMath::Pi()<<" (deg)"<<endl
	  <<"Phi:"<<phi*180./TMath::Pi()<<" (deg)"<<endl
	  <<"Momentum: "<<momentum<<" (GV)"<<endl;
	  } */  

    if ( omg->IsReadRandomMode() ){
      read_rndm_file->cd();
      UnPackDump(tmpfile, rnd_dump);
      TString rndfile = "/random/resetEngineFrom "+tmpfile;
      //cout<<"reading... "<<rndfile<<endl;
      ((TGeant4*)gMC)->ProcessGeantCommand(rndfile);
    }

    if( omg->DoWriteRandom() ){
      //cout<<"saving RND"<<endl;
      ((TGeant4*)gMC)->ProcessGeantCommand("/random/setSavingFlag 1");
    }

    gpr->ExtrapolateTrajectory();        //Added by V.Formato on 2012/06/19
    if(!gpr->ToBeSimulated()){ proc_ev++; continue;}

    /////RUN///////
    appl->RunMC(1);
    /////////////
    proc_ev ++;
    accept_ev = gpr->GetNgood();
    if (appl->GetTrigger()) trig_ev++;

    //    if(DEBUG)
      cout<<"TOT_EV: "<<proc_ev<<" TRIG_EV: "<<trig_ev<<" ACCEPT_EV:"<<accept_ev<<endl;
    
    //gObjectTable->Print();

    //saving dumps if any
    if( omg->DoWriteRandom() ){
      switch( omg->GetSaveCond() ){
      case EVERYTHING:
	dmp_file->cd();
	PackDump(tmpseeddir+"/currentEvent.rndm", ev_dump);
	save_ev = id;
	t_dump->Fill();
	break;
      case TRIG_ONLY:
	if(appl->GetTrigger()){
	  dmp_file->cd();
	  PackDump(tmpseeddir+"/currentEvent.rndm", ev_dump);
	  save_ev = id;
	  t_dump->Fill();
	}
	break;
      case ACCEPT_ONLY:
	if(gpr->Getgood()){
	  dmp_file->cd();
	  PackDump(tmpseeddir+"/currentEvent.rndm", ev_dump);
	  save_ev = id;
	  t_dump->Fill();
	}
	break;
      default:
	break;
      }
    }

    //condition to stop
    switch( omg->GetStopMode() ){
    case NEV_TOT:
      if( proc_ev >= nev_max) stop = kTRUE;
      break;
    case NEV_ACCEPT:
      if( accept_ev >= nev_max) stop = kTRUE;
      break;
    case NEV_TRIGG:
      if( trig_ev >= nev_max) stop = kTRUE;
      break;
    default:
      break;
    }

    if( (omg->GetReadPrimMode()) && ( proc_ev >= nev_max) ) stop = kTRUE;

    //gSystem->Exec(" ps aux | grep PamVMC.exe" );
  }
  gSystem->Exec(" ps aux | grep PamVMC.exe" );


  appl->FinishRun();
  if( omg->DoWriteRandom() ){
     dmp_file->cd();
     t_dump->Write();
     dmp_file->Close();
  }
 
  cout<<"done...."<<endl;
  
  if( read_rndm_file ) read_rndm_file->Close();
  delete read_rndm_file;
  delete appl;

  //delete tmp directory with dumps
  gSystem->Exec("rm -rf "+tmpseeddir);

  //moving good files to good directory
  if( omg->IsUseCookie() ){
    TString source_dir = working_path;
    TString dest_path =  omg->GetOutPath();
    TString out_root_fname = omg->GetOutPattern()+".root";
    TString out_pam_fname = omg->GetOutPattern()+".pam";
    TString out_dump_fname = omg->GetRandomOutFilename();

    cout<<"--->Move "<<source_dir<<"/"<<out_root_fname<<" to "<< dest_path<<endl;
    gSystem->Exec("mv " + source_dir+"/"+out_root_fname + " " + dest_path +"/");
    cout<<"--->Move "<<source_dir<<"/"<<out_pam_fname<<" to "<< dest_path<<endl;
    gSystem->Exec("mv " + source_dir+"/"+out_pam_fname + " " + dest_path +"/");
    if( omg->DoWriteRandom() ){
      cout<<"--->Move "<<out_dump_fname<<" to "<< dest_path<<endl;
      gSystem->Exec("mv " + out_dump_fname + " " + dest_path +"/");
    }
    cout<<"->Cleaning up tmpdir: "<<source_dir<<endl;
    gSystem->Exec("cd "+dest_path);
    gSystem->Exec("rm -rf " + source_dir );
  }
  return 0;
}


void PackDump(TString path, TObjString* os){

  ifstream ffile;
  ffile.open(path.Data(), ios::binary);
  if(!ffile) {
    cout<<"ERROR! wrong filename: "<<path<<endl;
    return;
  } 

  // get length of file:
  ffile.seekg (0, ios::end);
  Int_t length = ffile.tellg();
  ffile.seekg (0, ios::beg);

  // allocate memory:
  char* buffer = new char [length];

  // read data as a block:
  ffile.read (buffer,length);
  ffile.close();

  os->SetString(buffer);

  delete[] buffer;
}

void UnPackDump(TString path, TObjString* os){

  ofstream ffile;
  ffile.open(path.Data(), ios::binary);
  if(!ffile) {
    cout<<"ERROR! wrong filename: "<<path<<endl;
    return;
  } 
  
  // get length of file:
  Int_t length = os->Sizeof();

  // allocate memory:
  const char* buffer = 0;

  TString tmp = os->GetString();
  buffer = tmp.Data();
  
  // read data as a block:
  ffile.write (buffer,length-1);

  cout<<"->Unpack"<<endl;
  ffile.close();
}