#include <map>
#include <string>
#include <TFile.h>
#include <TDirectory.h>
#include <TROOT.h>

#include "IOARuntimeParameters.hxx"
#include "IDatum.hxx"
#include "IVInputFile.hxx"
#include "ICOMETEvent.hxx"
#include "IMCHit.hxx"
#include "ICOMETEventLoopFunction.hxx"
#include "IHandle.hxx"
#include "cometEventLoop.hxx"
#include "IMCHitTreeBuilder.hxx"
#include "ISumHitTreeBuilder.hxx"
using namespace COMET;

class TrackingTree: public COMET::ICOMETEventLoopFunction {
public:
    TrackingTree():fTotalEventsRead(0),
                   fInput(NULL),
                   fOutDir(NULL),
                   fCDCTreeBuilder(NULL),
                   fCTHTreeBuilder(NULL),
                   fSetupAnalyzers(false)
    {
    }

    virtual ~TrackingTree() {};

    void Usage(void) {
        std::cout<<"-O parameters=<file> Override the parameter file path."<<std::endl;
    }

    virtual bool SetOption(std::string option,std::string value="") {
        if(option=="parameters") COMET::IOARuntimeParameters::Get().ReadParamOverrideFile(value);
        else return false;
        return true;
    }

    bool operator () (COMET::ICOMETEvent& event) {
        // Initialize out analysers
        InitAnalyzers();
        // Record the event as having started
        fCDCTreeBuilder->EventStart();
        fCTHTreeBuilder->EventStart();
        // Get the CDC Hits
        COMET::IHandle<COMET::IHitSelection> cdcHits
            = event.Get<COMET::IHitSelection>("hits/mcCDC");
        // Get the CDC Hits
        COMET::IHandle<COMET::IHitSelection> cthHits
            = event.Get<COMET::IHitSelection>("hits/mcCTH");
        // Check if we need to run the analyser
        if (cdcHits) fCDCTreeBuilder->Analyze(cdcHits);
        if (cthHits) fCTHTreeBuilder->Analyze(cthHits);
        // End the event
        fCDCTreeBuilder->EventEnd();
        fCTHTreeBuilder->EventEnd();
        // Incriment total events and return false to keep the event loop 
        // running
        ++fTotalEventsRead;
        return false;
    }

    virtual void BeginFile(IVInputFile *const input){
        fInput=input;
    }

    virtual void Initialize(void){
        // Set up output directories
        fOutDir=TDirectory::CurrentDirectory();
    }

    virtual void Finalize(ICOMETOutput * const file){
        FinalizeAnalyzers();
        COMETLog("Total events read: "<<fTotalEventsRead);
    }

    void InitAnalyzers(){
        if (fSetupAnalyzers) return;
        // Initialize CDC Tree Builder
        fCDCTreeBuilder = new ISumHitTreeBuilder();
        fCDCTreeBuilder->SetEventTreeName("CDCHitTree");
        std::string dir_name_cdc="mcCDC";
        dir_name_cdc+="_"+fCDCTreeBuilder->GetName();
        TDirectory* out_dir_cdc=fOutDir->GetDirectory(dir_name_cdc.c_str());
        if(!out_dir_cdc)out_dir_cdc=fOutDir->mkdir(dir_name_cdc.c_str());
        if(out_dir_cdc) fCDCTreeBuilder->SetOutputDirectory(out_dir_cdc);
        fCDCTreeBuilder->Initialise("mcCDC");
        // Initialize CTH Tree Builder
        fCTHTreeBuilder = new ISumHitTreeBuilder();
        fCTHTreeBuilder->SetEventTreeName("CTHHitTree");
        std::string dir_name_cth="mcCTH";
        dir_name_cth+="_"+fCTHTreeBuilder->GetName();
        TDirectory* out_dir_cth=fOutDir->GetDirectory(dir_name_cth.c_str());
        if(!out_dir_cth)out_dir_cth=fOutDir->mkdir(dir_name_cth.c_str());
        if(out_dir_cth) fCTHTreeBuilder->SetOutputDirectory(out_dir_cth);
        fCTHTreeBuilder->Initialise("mcCTH");
        // Flag this as done 
        fSetupAnalyzers = true;
    }

    void FinalizeAnalyzers(){
        fCDCTreeBuilder->Finalize();
        fCTHTreeBuilder->Finalize();
    }

private:
    long int fTotalEventsRead;
    const IVInputFile* fInput;
    typedef std::map<std::string,IVAnalyzer*> AnalyzerList;
    TDirectory* fOutDir;
    IMCHitTreeBuilder* fCTHTreeBuilder;
    ISumHitTreeBuilder* fCDCTreeBuilder;
    bool fSetupAnalyzers;
};

int main(int argc, char **argv) {
    TrackingTree userCode;
    COMET::cometEventLoop(argc,argv,userCode,1);
}