#include <map>
#include <string>
#include <TDirectory.h>
#include <TParticlePDG.h>
#include <TDatabasePDG.h>
#include <TGeoNode.h>
#include <TROOT.h>
#include <TGeoManager.h>
#include <IOADatabase.hxx>

#include <HEPUnits.hxx>
#include "IDatum.hxx"
#include "IVInputFile.hxx"
#include "ICOMETEvent.hxx"
#include "IHandle.hxx"
#include "IRooTrackerFile.hxx"

COMET::IRooTrackerFile* fOutFile=NULL;
const std::string defaultTreeName("RooTrackerTree");
const std::string defaultOutFileName("output.rootracker");
int monitorID = 0;

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

void InitRooTracker(const char *file_rootracker){
  return;
}

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

void AddParticle(int evt,double x[3],double p[3], double mass, int pid){
  //brEvtNum=evt;
  double energy =std::sqrt(p[0]*p[0]+p[1]*p[1]+p[2]*p[2]+mass*mass);

  int particleNum=fOutFile->NewParticle();
  fOutFile->SetStdHepX4(      particleNum, x[0], x[1], x[2], 0.0);
  fOutFile->SetStdHepP4(      particleNum, p[0], p[1], p[2], energy);
  fOutFile->SetStdHepPdg(     particleNum, pid);
  fOutFile->SetStdHepTrackId( particleNum, 1);
  // Assume unity for the weight
  fOutFile->SetWeight(        particleNum, 1);
  // TODO have this support more than one monitor
  fOutFile->SetMonitorId(     particleNum, monitorID);
  // set to 1 since that is what SimG4 is expecting for particles that it can track
  fOutFile->SetStdHepStatus(  particleNum, 1);
  // Set the event number
  fOutFile->SetEvtNum(evt);

  return;
}

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

void PrintUsage(const char* progName){
    std::cout<<"Usage: "<<progName<<": [options] InputFile [InputFile2 ...]\n"
               "Convert text file to RooTracker ROOT tree\n"
               "Input text format is:\n"
               "x1 x2 x3 p1 p2 p3\n"
               "PDG codes, particle mass and monitor names are identical for all particles\n"
               "If OutputFile is not provided, histograms are written to '"<<defaultOutFileName<<"'\n"
               "Options:\n"
               "-t TREENAME      Specify the name of the tree in the file [defaults to "<<defaultTreeName <<"]\n"
               "-o OUTFILE       Specify the name of the output file [defaults to "<<defaultOutFileName <<"]\n"
               "-m monitor       Specify the name of monitor for the all entries\n"
               "-M mass          Specify the mass for all entries in GeV/c^2\n"
               "-n particles     Specify the number of particles in total\n"
               "-N number        Specify the number of partciles per event\n"
               "-p pid           Specify the PID for all entries\n"
               "-h               Print this help and exit\n"
               <<std::endl;
}

int main(int argc, char** argv){
    std::string outFileName=defaultOutFileName;
    std::string outTreeName=defaultTreeName;
    std::vector<std::string> inFileNames;
    std::string outMonitor="ProtonBeam";
    double outMass = 938.2720; // proton mass in MeV
    int outPID = 2212; // proton PID
    int totalNParticles = -1;
    int freqParticles = 1000;
    for (;;) {
        int c = getopt(argc, argv, "ht:o:m:M:n:N:p:");
        if (c<0) break;
        switch (c) {
            case 't':
                outTreeName=std::string(optarg);
                break;
            case 'o':
                outFileName=std::string(optarg);
                break;
            case 'm':
                outMonitor=optarg;
                break;
            case 'M':
                outMass=std::atof(optarg);
                break;
            case 'n':
                totalNParticles=std::atoi(optarg);
                break;
            case 'N':
                freqParticles=std::atoi(optarg);
                break;
            case 'p':
                outPID=std::atoi(optarg);
                break;
            case 'h':
                PrintUsage(argv[0]);
                return 0;
            default:
                std::cout<<"Error: Unknown option "<<c <<std::endl<<std::endl;
                PrintUsage(argv[0]);
                return 1;
        }
    }
    if (argc<=optind) {
        std::cerr << "Error: No input file" << std::endl << std::endl;
        PrintUsage(argv[0]);
        return 1;
    }else {
        for(int i=optind; i<argc;++i){
            inFileNames.push_back(argv[i]);
        }
        ++optind;
    }


    fOutFile = &*COMET::IRooTrackerFile::Open(outFileName.c_str(), 
                                              "create",
                                              outTreeName);
    monitorID = fOutFile->AddMonitor(outMonitor.c_str());
    int numberOfParciles = 0;

    for (std::vector<std::string>::iterator inFile = inFileNames.begin();
                                            inFile != inFileNames.end();
                                            ++inFile){
        FILE *fpin=fopen(inFile->c_str(),"r");
        int event_number=0;
        while(1){
            char line[1024];
            if(fgets(line,sizeof(line),fpin)==0) break;
            if(line[0]=='#' || line[0]=='\n') continue;

            double x[3],p[3];
            sscanf(line,"%lf %lf %lf %lf %lf %lf",&x[0],&x[1],&x[2],&p[0],&p[1],&p[2]);
            AddParticle(event_number,x,p, outMass, outPID);

            numberOfParciles++;
            if (numberOfParciles%freqParticles == 0) {
                fOutFile->Fill();
                event_number++;
            }
            if (numberOfParciles%10000 == 0) {
                COMETLog("Processed " << numberOfParciles << " paricles");
            }
            if (totalNParticles > 0 && numberOfParciles > totalNParticles) break;
        }
        fclose(fpin);
    }

    fOutFile->Close();

    return 0;
}

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