// M Hentz, 2018

#include "RunAccumulator.hh"

#include "G4SDManager.hh"
#include "G4ThreeVector.hh"
#include "G4VSensitiveDetector.hh"

#include "PhaseSpaceSD.hh"
#include "PhaseSpaceHit.hh"

#include <sys/stat.h> // to check if data directory exists

RunAccumulator::RunAccumulator( G4int bufferSize, G4double detPosition )
    :G4Run(),
     fDumpCount(0),
     fBufferSize(bufferSize),
     fDetPosition(detPosition),
     fDumpSingle(true)
{
  // Create and initialise output files
  InitialiseOutputFiles(fDumpSingle);

  // Number of parametrised copies in parallel world
//  G4int noOfCopies = (G4int) fDetDistance/fDetSpacing;

  // Instantiate hits vector and save it the run vector
  std::vector<PhaseSpaceHit> hitsVector;
  fRunVector.push_back(hitsVector);

  fDumpInterval = fBufferSize;
}

RunAccumulator::RunAccumulator( G4int bufferSize, G4double detDistance, G4double defSpacing )
    :G4Run(),
     fDumpCount(0),
     fBufferSize(bufferSize),
     fDetDistance(detDistance),
     fDetSpacing(defSpacing),
     fDumpSingle(false)
{
  // Create and initialise output files
  InitialiseOutputFiles(fDumpSingle);

  // Number of parametrised copies in parallel world
  G4int noOfCopies = (G4int) fDetDistance/fDetSpacing;

  // Instantiate noOfCopies hits vectors and save them in a vector of hits vectors
  for (G4int copy = 0; copy < noOfCopies; copy++) {
    std::vector<PhaseSpaceHit> hitsVector;
    fRunVector.push_back(hitsVector);
  }

  fDumpInterval = fBufferSize;
}

RunAccumulator::~RunAccumulator()
{
  fRunVector.clear();
}


void RunAccumulator::RecordEvent(const G4Event* event)
{
  G4HCofThisEvent* HCE = event->GetHCofThisEvent();
  if (!HCE) { return; }

  PhaseSpaceHitsCollection* eventCollection = (PhaseSpaceHitsCollection*) HCE->GetHC(0);
  if (!eventCollection) { return; }

  G4int nHits = eventCollection->entries();

  for ( G4int n = 0; n < nHits; n++ ){
    PhaseSpaceHit* hitPointer = (PhaseSpaceHit*) eventCollection->GetHit(n);

    // Use copy constructor and dereference operator to create hit object that
    // will be saved in vector of hits
    PhaseSpaceHit hit = PhaseSpaceHit(*hitPointer);

    // Insert hit into corresponding vector
    // - The position of the vector within the run vector corresponds to
    //   the volume the hit was registered in
    G4int copyNo = hit.GetCopyNo();
    fRunVector[copyNo].push_back(hit);
  }

  // Dump every fDumpInterbal events and make sure first event is not missed
  bool dump = (numberOfEvent + 1)%fDumpInterval == 0;
  if (dump) { Dump(); }

  // Increment event counter
  numberOfEvent++;
}


void RunAccumulator::Dump()
{
  // Increment dump counter
  fDumpCount++;

  // Avoid warning about comparison of signed and unsigned integers in for-loop below
  G4int noOfVectors = fRunVector.size();

  for (G4int i = 0; i < noOfVectors; i++) {
    if (fRunVector[i].size()) {
      // Only dump if there are hits in the vector
      WritePhaseSpaceHitsToFile(fRunOutputFilenames[i], fRunVector[i]);
      fRunVector[i].clear();
    }
  }

}

void RunAccumulator::InitialiseOutputFiles( G4bool dumpSingle )
{
  // Populates vector containing filenames that Run writes to in
  // WritePhaseSpaceHitsToFile(std::vector<PhaseSpaceHit>) and
  // creates the corresponding files with an appropriate header.

  // Check data directory exists
  struct stat info;
  if( stat( "data", &info ) != 0 ){
    throw std::runtime_error( "Directory data/ doesn't exist!" );
  }

  // If only detecting in single detector
  if( dumpSingle ) {
    G4double z = fDetPosition;

    std::stringstream stream;
    stream << std::fixed << "psf_z" << std::setprecision(0) << z << ".txt";
    G4String filename = stream.str();

    // Populate filename vector
    fRunOutputFilenames.push_back( filename );

    // Create the output files
    std::remove( filename );
    psfile = new std::ofstream;
    psfile->open( "data/" + filename );
    G4cout << " Created file " << filename << G4endl;

    // Insert the file header
    G4String header = "# parentID, name, x [mm], y [mm], z [mm], mom_x, mom_y, mom_z, ke [MeV]";
    (*psfile) << header << G4endl;

    psfile->close();
    return;
  }

  // If dumping to multiple files
  G4double zStart = 0.;
  G4double zEnd = fDetDistance;
  G4double zIncrement = fDetSpacing;

  G4cout << "Initialising phase space files in the range z in ["
         << zStart << ", " << zEnd << "] with an increment"
         << zIncrement << G4endl;

  for( G4double z = zStart; z < zEnd; z += zIncrement ){
    // Format filenames using std::stringstream to set the precision of G4double z
    std::stringstream stream;
    stream << std::fixed << "psf_z" << std::setprecision(0) << z << ".txt";
    G4String filename = stream.str();

    // Populate filename vector
    fRunOutputFilenames.push_back( filename );

    // Create the output files
    std::remove( filename );
    psfile = new std::ofstream;
    psfile->open( "data/" + filename );
    G4cout << " Created file " << filename << G4endl;

    // Insert the file header
    G4String header = "# parentID, name, x [mm], y [mm], z [mm], mom_x, mom_y, mom_z, ke [MeV]";
    (*psfile) << header << G4endl;

    psfile->close();
  }
}

void RunAccumulator::WritePhaseSpaceHitsToFile( G4String filename, std::vector<PhaseSpaceHit> hits )
{
  // Dumps hits on PhaseSpaceSD saved in their corresponding vector fRunVector[i].
  // The iteration variable i corresponds to the volume the particles were located at
  // in the ParallelWorldConstruction when the hit occured.

  // Define variables to be extracted from hit and dumped to file
  PhaseSpaceHit   hit;
  G4int           parentID;
  G4String        particleName;
  G4ThreeVector   pos;
  G4double        posX, posY, posZ;
  G4ThreeVector   mom;
  G4double        momX, momY, momZ;
  G4double        energy;

  // Dump
  psfile = new std::ofstream;
  psfile->open( "data/" + filename, std::ofstream::app );

  // Set precision of variables to be dumped
  (*psfile)
      << std::setiosflags(std::ios::fixed)
      << std::setprecision(6);

  // Iterate through hits and dump each hit
  G4int noOfHits = hits.size();

  for ( G4int j = 0; j < noOfHits; j++ ){
    parentID = hits[j].GetParentID();
    particleName = hits[j].GetParticleName();

    pos  = hits[j].GetPos(); // position
    posX = pos.x() / CLHEP::mm;
    posY = pos.y() / CLHEP::mm;
    posZ = pos.z() / CLHEP::mm;

    // Shift the position in z such that the source is at 0 mm
    G4double roomSizeZ = 8400.*CLHEP::mm;
    G4double zFromSource = posZ + roomSizeZ/2;

    mom = hits[j].GetMom(); // unit-length momentum direction
    momX = mom.x();
    momY = mom.y();
    momZ = mom.z();

    energy = hits[j].GetEnergy() / CLHEP::MeV;

    // Write quantities to file
    (*psfile)
        << parentID     << ","
        << particleName << ","
        << posX         << ","
        << posY         << ","
        << std::setprecision(1)
        << zFromSource  << ","
        << std::setprecision(6)
        << momX         << ","
        << momY         << ","
        << momZ         << ","
        << energy
        << G4endl;
  }
  psfile->close();
}