#include "IAnalyzeTreesMain.hxx"
#include "ITreeAnalyzerFactory.hxx"
#include "IBeamlineParameterExtractor.hxx"
#include "IParticleTagRegister.hxx"

#include <IOARuntimeParameters.hxx>
#include <IIntegerDatum.hxx>
#include <ICOMETLog.hxx>

#include <TGeoManager.h>
#include <TTree.h>
#include <TFile.h>
#include <TDirectory.h>
#include <TClonesArray.h>
#include <TChain.h>
#include <TBranch.h>

#include <set>
#include <sstream>

COMET::IAnalyzeTreesMain::IAnalyzeTreesMain():
        fInputTreeName("COMETEventsSummary"),
        fOutputFileName("tree_analysis.root"),
        fParamFile(""),
        fEntriesToRead(0),
        fPrintEntry(1000),
        fInputChain(NULL)
{
}

COMET::IAnalyzeTreesMain::~IAnalyzeTreesMain(){
    delete fInputChain;
}

void COMET::IAnalyzeTreesMain::PrintUsage(){
    COMETLog("Usage: \n"
            <<fProgName<<" [options] file_name [file_name_2 ...]\n"
            <<fProgName<<" [options] -I file_list\n"
            "\n"
            "Analyze a root tree (or trees) built by stripping an oaEvent file with a tree_builder analyzer.\n"
            "`file_name` can be a single file or contain a wild-card (escape this if you don't want the shell to expand it).\n"
            "All files and globs will be passed through to a ROOT TChain which will build the list of files\n"
            "\n"
            "Options:\n"
            " -h               Print this guidance\n"
            " -t TREENAME      The name of the tree in the file [currently "<<fInputTreeName <<"]\n"
            " -o OUTFILE       The name of the output file [currently "<<fOutputFileName <<"]\n"
            " -p param_file    Provide a parameter file to obtain histogram parameters\n"
            " -n entries       The number of entries to read.  [default: all entries]\n"
            " -u analyzer      Enable a specific analyzer\n"
            " -d analyzer      Disable a specific analyzer\n"
            " -I file_list     A text file containing a list of files to analyze\n"
            " -E branch_name   Enable a specific branch\n"
            " -D branch_name   Disable a specific branch\n"
            " -P print_count   Print the event number once for multiples of print_count\n"
            " -g geom_file     Provide the name of a file with a description of the geometry\n"
            );
}

std::string COMET::IAnalyzeTreesMain::GetOptionString(){
    return "hI:P:n:t:g:u:d:o:p:E:D:";
}

COMET::IAnalyzeTreesMain::Terminate_t COMET::IAnalyzeTreesMain::ProcessOption(
        int opt, const char* value){
    std::stringstream sstream;
    if(value) sstream.str(value);
    switch (opt) {
        case 'g':
            COMET::IBeamlineParameterExtractor::Load(fBeamline,TGeoManager::Import(value));
            fBeamline.Initialise();
            break;
        case 'P':
            sstream>>fPrintEntry;
            break;
        case 'p':
            fParamFile=optarg;
            break;
        case 'o':
            fOutputFileName=std::string(optarg);
            break;
        case 't':
            fInputTreeName=std::string(optarg);
            break;
        case 'n':
            sstream>>fEntriesToRead;
            break;
        case 'E':
            fBranchWhiteList.insert(value);
            break;
        case 'D':
            fBranchBlackList.insert(value);
            break;
        case 'u':
            if(!COMET::ITreeAnalyzerFactory::Instance().Enable(value)) return kErrorStop;
            break;
        case 'd':
            if(!COMET::ITreeAnalyzerFactory::Instance().Disable(value)) return kErrorStop;
            break;
        case 'I':
            return ParseInputFileList(value);
            break;
        case 'h':
            PrintUsage();
            return kNoErrorStop;
        default:
            std::cout<<"Error: Unknown option "<<opt <<std::endl<<std::endl;
            PrintUsage();
            return kErrorStop;
    }
    return kNoErrorContinue;
}

bool COMET::IAnalyzeTreesMain::Initialise(int argc, char* argv[]){
    fProgName=argv[0];

    Terminate_t status=ProcessAllOptions(argc,argv);
    if(Terminate(status)) return false;

    fInputChain=InitChain(fInputFileNames, fInputTreeName);
    if(!fInputChain) return false;
    COMETLog(fInputChain->GetNtrees()<<" files will be studied");

    if(!fParamFile.empty()){
        COMET::IOARuntimeParameters::Get().ReadParamOverrideFile(fParamFile);
    }

    std::map<std::string,std::string> data;
    status=PrepareData(data);
    if(Terminate(status)) return false;

    status=SetupAnalyzers(data);
    if(Terminate(status)) return false;

    DisableUnusedBranches();

    return true;
}

COMET::IAnalyzeTreesMain::Terminate_t COMET::IAnalyzeTreesMain::ProcessAllOptions(int argc,char* argv[]){
    for (;;) {
        int opt = getopt(argc, argv, GetOptionString().c_str());
        if (opt<0) break;
        Terminate_t status=ProcessOption(opt,optarg);
        if(Terminate(status)) return status;
    }
    for(; optind<argc;++optind){
        fInputFileNames.push_back(argv[optind]);
    }
    if (fInputFileNames.empty()) {
        COMETError("Error: No input file" << std::endl);
        PrintUsage();
        return kErrorStop;
    }
    return kNoErrorContinue;
}

COMET::IAnalyzeTreesMain::Terminate_t COMET::IAnalyzeTreesMain::ParseInputFileList(const char* filename){
    std::ifstream infile(filename);
    if(!infile){
        COMETError("Unable to open file, '"<<filename<<"' to read in the input file list");
        return kErrorStop;
    }
    std::string line;
    int n_files=0;
    while(std::getline(infile,line)){
        fInputFileNames.push_back(line);
        ++n_files;
    }
    COMETLog("Added "<<n_files<<" file(s) according to list '"<<filename<<"'");
    return kNoErrorContinue;
}

TChain* COMET::IAnalyzeTreesMain::InitChain(const std::vector<std::string>& file_list, const std::string& tree_name){
    // Open the input file(s)
    TChain* inChain=new TChain(tree_name.c_str());
    int n_added=0;
    for(size_t i=0; i<file_list.size(); ++i){
        int added=inChain->Add(file_list.at(i).c_str());
        if(added<=0){
            COMETError("Filename (or glob): "<<file_list.at(i)<<" is invalid.");
        }
        n_added+=added;
    }
    if(n_added<=0 || inChain->GetNtrees()<=0 || inChain->LoadTree(0)){
        COMETError("No valid input filename(s) provided.");
        return NULL;
    }
    return inChain;
}

COMET::IAnalyzeTreesMain::Terminate_t COMET::IAnalyzeTreesMain::PrepareData(
        std::map<std::string,std::string>& data){
    // It seems we need to do both of these for the chain to work with more than one
    // input file.  Doesn't seem to be documented anywhere though...
    const long int entries=fInputChain->GetEntries();
    fInputChain->LoadTree(0);
    COMETLog("Input file(s) contain :"<<entries<<" entries");

    // Set up each branch and the corresponding ClonesArray
    TIter next(fInputChain->GetListOfBranches());
    TObject* object;
    while((object=next())){
        TBranch* branch=(TBranch*)object;
        if(!branch) continue;

        std::string name=branch->GetName();
        if(branch->GetClassName()!=std::string("TClonesArray")) continue;
        if(!fBranchWhiteList.empty()){
            if(fBranchWhiteList.find(branch->GetName())==fBranchWhiteList.end()) continue;
        }else if (fBranchBlackList.find(branch->GetName())!=fBranchBlackList.end()) continue;

        COMET::analysis::ListOfClonesArrays::iterator i_data=
            fClonesArrays.insert(std::make_pair(name,COMET::analysis::TreeData())).first;
        const TClonesArray*& array=i_data->second.array;
        fInputChain->SetBranchAddress(name.c_str(),&array,&i_data->second.branch);
        branch->SetAddress(&array);

        for(int i=0; !array && i<100;++i)branch->GetEntry(i);
        if(!array){
            COMETError("Problem identifying data type for branch: "<<branch->GetName());
            return kErrorStop;
        }
        data[branch->GetName()]=array->GetClass()->GetName();
        COMETLog("Analyzing: "<<branch->GetName()<<" of type: "<<array->GetClass()->GetName());
    }

    fInputChain->LoadTree(0);
    bool ok=COMET::IParticleTagRegister::InitFromFile(fInputChain->GetFile());
    if(!ok) return kErrorStop;
    COMETLog("Valid tags are:");
    COMET::IParticleTagRegister::Instance().PrintAllTags();

    return kNoErrorContinue;
}

COMET::IAnalyzeTreesMain::Terminate_t COMET::IAnalyzeTreesMain::SetupAnalyzers(
        std::map<std::string,std::string>& data){
    // Create TreeAnalyzers for the input data
    COMET::ITreeAnalyzerFactory::Instance().GetAnalyzers(data,fAnalyzers);
    if(fAnalyzers.empty()){
        COMETError("No valid analyzers were created");
        COMET::ITreeAnalyzerFactory::Instance().PrintKnownTypes();
        return kErrorStop;
    }

    fOutputFile=TFile::Open(fOutputFileName.c_str(),"CREATE");
    if(!fOutputFile || fOutputFile->IsZombie()){
        COMETError("Problem opening output file");
        return kErrorStop;
    }

    bool ok=true;
    for(AnalyzerList::const_iterator i_analyzer=fAnalyzers.begin();
            i_analyzer!=fAnalyzers.end(); ++i_analyzer){
        ok&=(*i_analyzer)->Initialise(fClonesArrays,*this);
    }
    if(!ok) return kErrorStop;

    return kNoErrorContinue;
}

void COMET::IAnalyzeTreesMain::DisableUnusedBranches(){
    std::set<std::string> used_data;
    for(AnalyzerList::const_iterator i_analyzer=fAnalyzers.begin();
            i_analyzer!=fAnalyzers.end(); ++i_analyzer){
        const std::set<std::string>* targets=COMET::ITreeAnalyzerFactory::Instance().GetTargetDataNames((*i_analyzer)->GetName());
        if(targets && !targets->empty()) used_data.insert(targets->begin(), targets->end());
        else used_data.insert((*i_analyzer)->GetDatumName());
    }

    fInputChain->SetBranchStatus("*",0);
    for(COMET::analysis::ListOfClonesArrays::const_iterator i_data=fClonesArrays.begin();
            i_data!=fClonesArrays.end(); ++i_data){
        if(used_data.find(i_data->first)==used_data.end()) continue;
        fInputChain->SetBranchStatus((i_data->first).c_str(),1);
        fInputChain->SetBranchStatus((i_data->first+".*").c_str(),1);
        COMETLog("Enabled branch: "<<i_data->first);
    }
}

bool COMET::IAnalyzeTreesMain::Process(){
    TDirectory::TContext context(0);
    bool ok=true;

    for(fCurrentEntry=0;
            (fEntriesToRead<1 || fCurrentEntry<fEntriesToRead);
                ++fCurrentEntry){
        int read=fInputChain->LoadTree(fCurrentEntry);
        if(read<0) break;
        if(read==0){
            // New file opened
            COMET::IParticleTagRegister::InitFromFile(fInputChain->GetFile());
        }
        read=fInputChain->GetEntry(fCurrentEntry);
        if(read<0){
            COMETError("Problem reading in event "<<fCurrentEntry<<". Skipping...");
            continue;
        }
        if(fCurrentEntry%fPrintEntry==0) COMETNamedLog("IAnalyzeTreesMain::Process","Looking at entry: "<<fCurrentEntry);

        for(AnalyzerList::const_iterator i_analyzer=fAnalyzers.begin();
                i_analyzer!=fAnalyzers.end(); ++i_analyzer){
            ok&=(*i_analyzer)->Analyze();
        }
        if(!ok) break;
    }
    return ok;
}

bool COMET::IAnalyzeTreesMain::Finalize(){
    bool ok=true;
    for(AnalyzerList::const_iterator i_analyzer=fAnalyzers.begin();
            i_analyzer!=fAnalyzers.end(); ++i_analyzer){
        ok&=(*i_analyzer)->Finalize();
    }
    //fOutputFile->ls();
    COMETLog("Read "<<fCurrentEntry<<" entries");
    COMET::IIntegerDatum* events_read=new COMET::IIntegerDatum("events_read",Form("%lu",fCurrentEntry),fCurrentEntry);
    events_read->Write();
    fOutputFile->Write();
    fOutputFile->Close();
    return ok;
}