#include "CRYGenerator.h"
#include "CRYSetup.h"

#include "TROOT.h"
#include "TFile.h"
#include "TMath.h"
#include "TNtuple.h"
#include "TH1F.h"

#include <iostream>
#include <string>
#include <vector>
#include <fstream>
#include <map>
#include <math.h>
#include <stdlib.h>  // For Ubuntu Linux

// Choose what detector(s) to focus on
bool CDC = true, FDC = false, BCAL = true, TOF = false, FCAL = false;

//These are needed to define "box" we project the tracks on (set artifically high to start)
float xMin = 10, xMax = -10, yMin = 10, yMax = -10, zPlane = 0;
 
int main( int argc, const char *argv[]) {

  // This step sets the box dimensions according to whether specific detectors are enabled. (tracks with no hits are eventually rejected by HDGEANT)
  // Remember that the coordinate system we use is different (see comments near the end of this code.)
  if(CDC){
	float xMinLocal = -0.05, xMaxLocal = 1.77, yMinLocal = -0.925, yMaxLocal = 0.925, zLocal = 0.925;
	if(xMin > xMinLocal) xMin = xMinLocal;
	if(xMax < xMaxLocal) xMax = xMaxLocal;
        if(yMin > yMinLocal) yMin = yMinLocal;
        if(yMax < yMaxLocal) yMax = yMaxLocal;
	if(zPlane < zLocal) zPlane = zLocal;
  }
  if(FDC){
        float xMinLocal = 1.85, xMaxLocal = 3.60, yMinLocal = -0.925, yMaxLocal = 0.925, zLocal = 0.925;
        if(xMin > xMinLocal) xMin = xMinLocal;
        if(xMax < xMaxLocal) xMax = xMaxLocal;
        if(yMin > yMinLocal) yMin = yMinLocal;
        if(yMax < yMaxLocal) yMax = yMaxLocal;
        if(zPlane < zLocal) zPlane = zLocal;
  }
  if(BCAL){
        float xMinLocal = 0.17, xMaxLocal = 4.10, yMinLocal = -0.925, yMaxLocal = 0.925, zLocal = 0.925;
        if(xMin > xMinLocal) xMin = xMinLocal;
        if(xMax < xMaxLocal) xMax = xMaxLocal;
        if(yMin > yMinLocal) yMin = yMinLocal;
        if(yMax < yMaxLocal) yMax = yMaxLocal;
        if(zPlane < zLocal) zPlane = zLocal;
  }
  if(FCAL){
        float xMinLocal = 6.25, xMaxLocal = 7.25, yMinLocal = 1.2, yMaxLocal = 1.2, zLocal = 1.2;
        if(xMin > xMinLocal) xMin = xMinLocal;
        if(xMax < xMaxLocal) xMax = xMaxLocal;
        if(yMin > yMinLocal) yMin = yMinLocal;
        if(yMax < yMaxLocal) yMax = yMaxLocal;
        if(zPlane < zLocal) zPlane = zLocal;
  }
  if(TOF){
        float xMinLocal = 6.0, xMaxLocal = 6.40, yMinLocal = -1.5, yMaxLocal = 1.5, zLocal = 1.5;
        if(xMin > xMinLocal) xMin = xMinLocal;
        if(xMax < xMaxLocal) xMax = xMaxLocal;
        if(yMin > yMinLocal) yMin = yMinLocal;
        if(yMax < yMaxLocal) yMax = yMaxLocal;
        if(zPlane < zLocal) zPlane = zLocal;
  }

  if (!TROOT::Initialized()) {
    static TROOT root("RooTuple", "I am the Gatekeeper");
  }


  if ( argc < 2 ) {
    std::cout << "usage " << argv[0] << " <setup file name> <simulated time (seconds)>\n";
    std::cout << "simulated time = 5 sec by default\n";
    return 0;
  }

  TFile *outputROOTFile=new TFile("gluexCRYgen.root","RECREATE");
  
  TNtuple * ntuple = new TNtuple("nt1","NTuple of thrown muons","ke:x:y:z:u:v:w");
  TH1F* costhe       = new TH1F("costhe",  "costhe plot;cos(theta)",   200, -1., 0);
  TH1F* theta       = new TH1F("theta",  "theta plot;theta",   200, 0, 90);
  bool fillHistograms = true;

  //int nEv = 100000;
  float totalTime = 5.0;

  //if (argc > 2 ) nEv = atoi(argv[2]);
  if (argc > 2 ) totalTime = atof(argv[2]);

  // Read the setup file into setupString
  std::ifstream inputFile;
  inputFile.open(argv[1],std::ios::in);
  char buffer[1000];

  std::string setupString("");
  while (!inputFile.getline(buffer,1000).eof()) {
    setupString.append(buffer);
    setupString.append(" ");
  }

  // Parse the contents of the setup file
  CRYSetup *setup=new CRYSetup(setupString,"../data");

  // Setup the CRY event generator
  CRYGenerator gen(setup);

  // Initialize some counters
  int nMuon = 0, nMuonAccept = 0, i=0;

  // Initialize the simulated particle
  std::vector<CRYParticle*> *ev=new std::vector<CRYParticle*>;

  // Initialize the output file
  int nEventsPerFile=100000;
  std::ofstream outputFile;
  const char * outputFileName = Form("cryoutput%i.ascii",nMuonAccept/nEventsPerFile);
  outputFile.open(outputFileName,std::ios::trunc);

  while(gen.timeSimulated()<totalTime){
    //for (int i = 0; i < nEv; i++) {

    ev->clear();
    gen.genEvent(ev);
    
    // Print even number every 100k events
    i++;
    if (i % 100000 == 0) std::cout << "Event: " << i << std::endl;

    // Loop over secondary particles (Only one secondary particle => this loop isn't needed)
    //for (unsigned j = 0; j < ev->size(); j++) {
    unsigned j = 0;
      CRYParticle *part=(*ev)[j];
 
      // Quick check that generated particle is a muon. Should be set in configuration file 

      if (part->id() != CRYParticle::Muon) continue;
	nMuon++;

	//CRY plane location
	//float zPlane = 1.88;
	//float zPlane = 0.925;

	//The particles are returned from CRY with an origin location (x,y,z) and their 
	//direction cosines (u,v,w). These are used to project the paths onto a box surrounding the region of the detector.
	//this will reject tracks not passing near the detector and will also decrease the length of track propagtion in GEANT.
	
	bool goodEvent = false;

	//Initialize the origin locations that will be set for each particle
	float xOr, yOr, zOr;

	// Plane above detectors is at the CRY plane. All tracks accepted that originate here.
       	if (part->x() > xMin && part->x() < xMax && part->y() > yMin && part->y() < yMax) {
	  goodEvent = true;
	  xOr = part->x(); yOr = part->y(); zOr = zPlane;
	}

	// First the plane at xMin
	if (!goodEvent){
	  xOr = xMin;
	  float xDist = part->x() - xOr;
	  float pathLength = xDist / part->u();
	  yOr = part->y() + pathLength*part->v();
	  zOr = zPlane + pathLength*part->w();
	  //The check on the direction cosine ensures the path is pointing into the region of the detector
	  if(yOr > yMin && yOr < yMax && zOr > zPlane*(-1) && zOr < zPlane && part->u() > 0) goodEvent = true;
	}

	//Next the plane at xMax
        if (!goodEvent){
	  xOr = xMax;
	  float xDist = part->x() - xOr;
          float pathLength = xDist / part->u();
          yOr = part->y() + pathLength*part->v();
          zOr = zPlane + pathLength*part->w();
	  if(yOr > yMin && yOr < yMax && zOr > zPlane*(-1) && zOr < zPlane && part->u() < 0) goodEvent = true;
        }

	  //Now the plane at yMin
	  if (!goodEvent){
	    yOr = yMin;
	    float yDist = part->y()-yOr;
	    float pathLength = yDist / part->v();
	    xOr = part->x() + pathLength*part->u();
	    zOr = zPlane + pathLength*part->w();
	    if(xOr > xMin && xOr < xMax && zOr > zPlane*(-1) && zOr < zPlane && part->v() > 0) goodEvent = true;
	  }

	  //Finally the plane at yMax
          if (!goodEvent){
	    yOr = yMax;
	    float yDist = part->y()-yOr;
            float pathLength = yDist / part->v();
            xOr = part->x() + pathLength*part->u();
            zOr = zPlane + pathLength*part->w();
            if(xOr > xMin && xOr < xMax && zOr > zPlane*(-1) && zOr < zPlane && part->v() < 0) goodEvent = true;
          }

	if (!goodEvent) {
		delete (*ev)[j];
		continue;
	}

	nMuonAccept++;
	
	//Fill the following ntuple and histograms of the accepted tracks
	if (fillHistograms == true){
      		ntuple->Fill(part->ke(), xOr,yOr,zOr,part->u(),part->v(),part->w());
      		theta->Fill((TMath::Pi()-TMath::ACos(part->w()))*TMath::RadToDeg());
      		costhe->Fill(part->w());
	}      

      //Of course all of the units are different, and the coordinate system is different between CRY and GEANT
      //
      //    Units                     Coordinates
      // CRY  |  HDGEANT             CRY  | GEANT
      //================             ============
      // MeV  |  GeV                  x   |   z
      //  m   |  cm                   y   |   x
      //                              z   |   y
      //
      //Also GEANT wants the momentum and not the kinetic energy, so we need to make all of the conversions
      
      float xGe, yGe, zGe, uGe, vGe, wGe, px, py, pz;
      float mMuon = 105.66; // MeV/c^2
      float mMuonGe = .10566; // Gev/c^2
      float E = part->ke()/1000 + mMuonGe; // GeV
      float pTot = TMath::Sqrt( E * E - mMuonGe * mMuonGe); // GeV/c
      xGe = yOr*100; yGe = zOr*100; zGe = xOr*100; // cm 
      uGe = part->v(); vGe = part->w(); wGe = part->u();
      px = pTot * uGe; py = pTot * vGe; pz = pTot * wGe;

      //Output to ascii file
      int runNumber = 1, nParticles = 1;
	if(nMuonAccept % nEventsPerFile == 0){
		if(part->charge() == 1){
                        outputFile << runNumber << " " << nEventsPerFile << " " << nParticles << std::endl;
                        outputFile << "1 " << "mu+ " << mMuonGe << " " << part->charge() << " ";
                        outputFile << xGe << " " << yGe << " " << zGe << " " << px << " " << py << " " << pz << " " << E << std::endl;
                }
                else{
                        outputFile << runNumber << " " << nEventsPerFile << " " << nParticles << std::endl;
                        outputFile << "1 " << "mu- " << mMuonGe << " " << part->charge() << " ";
                        outputFile << xGe << " " << yGe << " " << zGe << " " << px << " " << py << " " << pz << " " << E << std::endl;
                }
                outputFile.close();
                outputFileName = Form("cryoutput%i.ascii",nMuonAccept/nEventsPerFile);
                outputFile.open(outputFileName,std::ios::trunc);
  	}
	else{
		if(part->charge() == 1){
      			outputFile << runNumber << " " << nMuonAccept % nEventsPerFile << " " << nParticles << std::endl;
      			outputFile << "1 " << "mu+ " << mMuonGe << " " << part->charge() << " ";
      			outputFile << xGe << " " << yGe << " " << zGe << " " << px << " " << py << " " << pz << " " << E << std::endl;
		}	
        	else{
                	outputFile << runNumber << " " << nMuonAccept % nEventsPerFile << " " << nParticles << std::endl;
                	outputFile << "1 " << "mu- " << mMuonGe << " " << part->charge() << " ";
                	outputFile << xGe << " " << yGe << " " << zGe << " " << px << " " << py << " " << pz << " " << E << std::endl;
       		}
	}
      delete (*ev)[j];
  if(nMuonAccept % nEventsPerFile == 0){
    outputFile.close();
    outputFileName = Form("cryoutput%i.ascii",nMuonAccept/nEventsPerFile);
    outputFile.open(outputFileName,std::ios::trunc);
  }
 
  }

  std::cout << "Run completed.\n";
  std::cout << "Total muons: " << nMuon << " muons\n";
  std::cout << "Total muons accepted: " << nMuonAccept << " muons\n";
  std::cout << "Total time simulated: " << gen.timeSimulated() << " seconds\n";
  double muonsPerSecondPerm2=nMuon/(30*30*gen.timeSimulated());
  std::cout << "Muons per second per m2 (assuming 30x30 m)" << muonsPerSecondPerm2 << std::endl;

  // Write the ROOT file
  outputROOTFile->Write();
  outputROOTFile->Close();

  delete setup;

  return 0;
}