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

#include <IIntegerDatum.hxx>
#include <IOARuntimeParameters.hxx>
#include <IDatum.hxx>
#include <IVInputFile.hxx>
#include <ICOMETEvent.hxx>
#include <IG4VHit.hxx>
#include <ICOMETEventLoopFunction.hxx>
#include <IHandle.hxx>
#include <cometEventLoop.hxx>

#include "IAnalysisAssistant.hxx"
#include "IAnalyzerFactory.hxx"
#include "IParticleTagger.hxx"
#include "IEventWeighter.hxx"

using namespace COMET;

class AnalysisLoop: public COMET::ICOMETEventLoopFunction {
    typedef std::set<std::string> DatumNameSet;
    typedef std::vector<COMET::IVAnalyzer*> AnalyzerList;
    typedef std::map<std::string,AnalyzerList> DatumAnalyzerMap;

public:
    AnalysisLoop():fTotalEventsRead(0),fInput(NULL),fOutDir(NULL),fSplitDataVectors(true){
        fDataToAnalyse.insert("truth/G4Trajectories");
        fDataToAnalyse.insert("truth/g4Hits");
        fDataToAnalyse.insert("truth/G4PrimVertex00");
    }

    virtual ~AnalysisLoop() {};

    void Usage(void) {
        std::cout<<"-O parameters=<file> Override the parameter file path."<<std::endl;
        std::cout<<"-O datum=<full_path> Add a name to the list of data to visit."<<std::endl;
        std::cout<<"-O weighting=<weighter> Weight / skip events using the named weighter."<<std::endl;
        std::cout<<"-O use=<analyzer> Add an analyzer to be used."<<std::endl;
        std::cout<<"-O dont-use=<analyzer> Add an analyzer to be ignored."<<std::endl;
        std::cout<<"-O clear-data-list Reset the datum list to analyse"<<std::endl;
        std::cout<<"-O no-split-datavectors Dont split up datavectors"<<std::endl;
        std::cout<<"--- Valid analyzers are:" <<std::endl;
        COMET::IAnalyzerFactory::Instance().PrintKnownTypes();
    }

    virtual bool SetOption(std::string option,std::string value="") {
        if(option=="parameters") COMET::IOARuntimeParameters::Get().ReadParamOverrideFile(value);
        else if(option=="datum") fDataToAnalyse.insert(value);
        else if(option=="clear-data-list") fDataToAnalyse.clear();
        else if(option=="use") return COMET::IAnalyzerFactory::Instance().Enable(value);
        else if(option=="dont-use") return COMET::IAnalyzerFactory::Instance().Disable(value);
        else if(option=="weighting") return COMET::IEventWeighter::Instance().SetWeighter(value);
        else if(option=="no-split-datavectors") fSplitDataVectors=false;
        else return false;
        return true;
    }

    bool operator () (COMET::ICOMETEvent& event) {
        double event_weight=COMET::IEventWeighter::CalculateWeight(event);
        if(event_weight<0) {
            return false;
        }

        bool ok=COMET::IParticleTagger::Instance().TagNewEvent(event);
        //COMET::IParticleTagger::Instance().PrintSummary();
        if(!ok){
            COMETError("Problem trying to tag particles");
            return false;
        }

        for(auto i_datum=fAnalyzers.begin();
                i_datum!=fAnalyzers.end();++i_datum){
            for(auto i_analyzer=i_datum->second.begin();
                    i_analyzer!=i_datum->second.end();++i_analyzer){
                ok&=(*i_analyzer)->EventStart();
            }
        }
        if(!ok) {
            COMETError("Problem during event start");
            return false;
        }

        for(auto i_request=fDataToAnalyse.begin(); 
                i_request!=fDataToAnalyse.end(); ++i_request){
            COMET::IHandle<COMET::IDatum> datum 
                = event.Get<COMET::IDatum>(*i_request);
            if(!datum) {
                static std::map<std::string,int> warn_count;
                int& count=warn_count[*i_request];
                if(count++<10) COMETError("Couldn't find datum called: '"<<*i_request<<"'");
                if(count==10) COMETError("Will no longer warn about missing datum called: '"<<*i_request);
            } else {
                AnalyseDatum(datum);
            }
        }

        for(auto i_datum=fAnalyzers.begin();
                i_datum!=fAnalyzers.end();++i_datum){
            for(auto i_analyzer=i_datum->second.begin();
                    i_analyzer!=i_datum->second.end();++i_analyzer){
                ok&=(*i_analyzer)->EventEnd();
            }
        }
        if(!ok) {
            COMETError("Problem during event end");
            return false;
        }

        ++fTotalEventsRead;
        // Don't write events out so return false always
        return false;
    }

    void AnalyseDatum(const COMET::IHandle<COMET::IDatum> datum){
        static std::string data_vector_class="COMET::IDataVector";
        if(data_vector_class==datum->ClassName() && fSplitDataVectors){
            COMET::IHandle<COMET::IDataVector> data_vector=datum->Get<COMET::IDataVector>(".");
            if(!data_vector){
                COMETError("Funny behaviour: datum '"<<datum->GetName()<<"' thinks it's an IDataVector but won't cast as such");
                return;
            }
            for(auto i_datum =data_vector->begin(); 
                    i_datum!=data_vector->end(); ++i_datum){
                AnalyseDatum(COMET::IHandle<COMET::IDatum>(*i_datum,false));
            }
        }else{
            AnalyzerList& analyzers=GetAnalyzers(datum);
            for(AnalyzerList::iterator i_anal=analyzers.begin();
                    i_anal!=analyzers.end(); ++i_anal){
                (*i_anal)->Analyze(datum);
            }
        }
    }

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

    virtual void Initialize(){
        fOutDir=TDirectory::CurrentDirectory();
        if (!fOutDir->IsWritable()) {
            COMETError("Output directory is not writable.\n" << 
                    "Did you forget to specify an output filename?");
            exit(-1);
        }
    }

    virtual void Finalize(ICOMETOutput * const file){
        FinalizeAnalyzers();
        // Log the various stats from this run:
        COMET::IIntegerDatum* events_read=new COMET::IIntegerDatum("events_read",Form("%ld",fTotalEventsRead),fTotalEventsRead);
        events_read->Write();
        COMETLog("Total events read: "<<fTotalEventsRead);
    }

    AnalyzerList& GetAnalyzers(const COMET::IHandle<IDatum> datum){
        // Previously visited this datum
        DatumAnalyzerMap::iterator i_list=fAnalyzers.find(datum->GetName());
        if(i_list!=fAnalyzers.end()) return i_list->second;

        // New datum set-up analyzers
        AnalyzerList& analyzers=fAnalyzers[datum->GetName()];
        COMET::IAnalyzerFactory::Instance().GetAnalyzers(&*datum,analyzers);
        COMETNamedTrace("AnalysisLoop","# analyzers: "<<analyzers.size());
        for(AnalyzerList::iterator i_anal=analyzers.begin();
                i_anal!=analyzers.end(); ++i_anal){
            COMET::IVAnalyzer* analyzer=*i_anal;
            if(!analyzer)continue;
            std::string dir_name=datum->GetName();
            dir_name+="_"+analyzer->GetName();
            COMETNamedTrace("AnalysisLoop","Adding analyzer, '"<<analyzer->GetName()<<" for "<<datum->GetName()<<", plots to: "<<dir_name);
            TDirectory* out_dir=fOutDir->GetDirectory(dir_name.c_str());
            if(!out_dir)out_dir=fOutDir->mkdir(dir_name.c_str());
            if(out_dir) analyzer->SetOutputDirectory(out_dir);
            analyzer->Initialise(datum->GetName());
            analyzer->EventStart();
        }
        return analyzers;
    }

    void FinalizeAnalyzers(){
        for(auto i_list=fAnalyzers.begin();
                i_list!=fAnalyzers.end(); ++i_list){
            for(auto i_anal=i_list->second.begin();
                    i_anal!=i_list->second.end(); ++i_anal){
                (*i_anal)->Finalize();
            }
        }
        //Write the tag register to disk
        COMET::IParticleTagger::SaveRegister();
    }

private:
    long int fTotalEventsRead;
    const IVInputFile* fInput;
    DatumAnalyzerMap fAnalyzers;
    DatumNameSet fDataToAnalyse;
    TDirectory* fOutDir;
    bool fSplitDataVectors;
};

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