#include <cometEventLoop.hxx>
#include <memory>
#include <iostream>

#include <TApplication.h>
#include <TFile.h>
#include <TTree.h>
#include <TVector3.h>

#include "IGeomVisitor.hxx"
#include "IGeomInfo.hxx"
#include "ICTHGeom.hxx"
#include "IDetectorSolenoidGeom.hxx"

class IGeomTestCTH: public COMET::ICOMETEventLoopFunction {

public:

  IGeomTestCTH()
  :fCTHGeom(NULL)
  ,fDSGeom(NULL)
  ,fGeomInit(false)
  ,fOFName("geomTestCTHoutput")
  ,fOutputFile(NULL)
  ,fTree(NULL)
  ,fEvent(0)
  {
  }
  
  virtual ~IGeomTestCTH() {};

  void Usage(void) {
    std::cout<< "    -O ofile     name of the output file" << std::endl;
  }

  virtual bool SetOption(std::string option,std::string value="") {
    if (option == "ofile") fOFName = value;
    return true;
  }

  void Initialize() {
    std::string rootname = fOFName+".root";
    fOutputFile = new TFile(rootname.c_str(),"RECREATE");
    if (!fOutputFile->IsOpen()) {
      std::cout << "couldn't open output file " << rootname << std::endl;
      return;
    }
    fTree = new TTree("GeomCTH", "CHT geometry information");
    fTree->Branch("event",  &fEvent,  "event/I");
    fTree->Branch("channel",&fChannel,"channel/I");
    fTree->Branch("counter",&fCounter,"counter/I");
    fTree->Branch("module", &fModule, "module/I");
    fTree->Branch("segment",&fSegment,"segment/I");
    fTree->Branch("scint",  &fScint,  "scint/I");
    fTree->Branch("global", "TVector3", &fGlobal);
    fTree->Branch("local",  "TVector3", &fLocal);
  }

  void Finalize(COMET::ICOMETOutput* const output) {
    std::string rootname = fOFName+".root";
    if (!fOutputFile->IsOpen()) { std::cout << " no output file " << rootname << std::endl ; return;}
    fOutputFile->cd();
    fTree->Write();
  }

  void GeomInit() {
    COMETLog("Init geometry");

    fDSGeom = new COMET::IDetectorSolenoidGeom();
    fDSGeom->Fill();
    fDSGeom->ls("");

    fCTHGeom = new COMET::ICTHGeom();
    fCTHGeom->Fill();
    fCTHGeom->ls("");

    COMETLog("GetNumberOfModules()   = " << fCTHGeom->GetNumberOfModules());
    COMETLog("GetNumberOfLayers()    = " << fCTHGeom->GetNumberOfLayers());
    COMETLog("GetNumberOfSegments()  = " << fCTHGeom->GetNumberOfSegments());

    fGeomInit = true;
  }
  
  bool operator () (COMET::ICOMETEvent& event) {
    COMET::IOADatabase::Get().UpdateGeometry();
    COMET::IGeomInfo::Get();
    COMETLog("  Start ev "<<fEvent);

    if (!fGeomInit) GeomInit();

    TString mod[2] = {"US array","DS array"};
    TString lay[2] = {"Cherenkov","Scintillator"};
    for (UInt_t i = 0; i < (UInt_t)fCTHGeom->GetNumberOfModules(); i++) {
      for (UInt_t j = 0; j < (UInt_t)fCTHGeom->GetNumberOfSegments(); j++) {
        for (UInt_t k = 0; k < (UInt_t)fCTHGeom->GetNumberOfLayers(); k++) { /// Cherenkov:0, Scintilaltor:1
          fCounter = fCTHGeom->GetCounterId(i,k,j,0);
          fChannel = fCTHGeom->CounterIdToChannel(fCounter);
          COMETLog("For "<<lay[k].Data()<<", "<<j<<"-th segment @"
                         <<mod[i]<<" (counterId="<<fCounter<<")");
          if (!fCTHGeom->ChannelToGlobalPosition(fChannel, fGlobal)) {
            COMETWarn("ChannelToGlobalPosition() returns false!");
            return false;
          }
          if(!fDSGeom->GetDetPositionInDSCoordinate(fGlobal, fLocal)) {
            COMETWarn("Cannot find the position in DS coordinate!");
            return false;
          }
          fModule  = i;
          fSegment = j;
          fScint   = k;
          COMETLog("  Global position = ("<<fGlobal.X()<<", "<<fGlobal.Y()<<", "<<fGlobal.Z()<<")");
          COMETLog("  Local  position = ("<<fLocal.X() <<", "<<fLocal.Y() <<", "<<fLocal.Z()<<")");
          fTree->Fill();
        }
      }
    }
    COMETLog("  DONE ev "<<fEvent);
    fEvent++;
    return true;
  }

  COMET::ICTHGeom *fCTHGeom;  
  COMET::IDetectorSolenoidGeom *fDSGeom;  

private:
  bool        fGeomInit;

  std::string fOFName;
  TFile*      fOutputFile; 
  TTree*      fTree; 

  UInt_t      fEvent;
  UInt_t      fChannel;
  UInt_t      fCounter;
  UInt_t      fModule;
  UInt_t      fSegment;
  UInt_t      fScint;
  TVector3    fGlobal;
  TVector3    fLocal;
};


int main(int argc, char **argv) {
  IGeomTestCTH userCode;

  cometEventLoop(argc,argv,userCode,1);
}