#include <TROOT.h>
#include <TH1.h>
#include <TH2.h>

#include "IVInputFile.hxx"
#include "ICOMETInput.hxx"
#include "IProfileData.hxx"
#include "IG4Trajectory.hxx"
#include "ICOMETEvent.hxx"
#include "IG4VHit.hxx"
#include "ICOMETEventLoopFunction.hxx"
#include "IHandle.hxx"
#include "cometEventLoop.hxx"

using namespace COMET;

namespace{
    class HistByEvent{
        public:
            HistByEvent(const char* name, const char* title, const char* axis_title, int num_events, double delta_max, double max, double delta_min=0, double min=0):
                fName(name),fTitle(title),fAxisTitle(axis_title),fEvents(0),fLastValue(0){
                    hValue=new TH1F(name,(fTitle+";Event No.;"+fAxisTitle).c_str(),num_events,0,num_events );
                    hValueBinned=new TH1F((fName+"_binned").c_str(),(fTitle+" Binned;"+fAxisTitle).c_str(),400, min,max);
                    hValueDelta=new TH1F((fName+"_delta").c_str(),("Change in "+fTitle+";Event No.;Change in "+fAxisTitle).c_str(),num_events-1,0,num_events-1 );
                    hValueDeltaBinned=new TH1F((fName+"_delta_binned").c_str(),("Change in "+fTitle+" Binned;"+fAxisTitle).c_str(),400,delta_min,delta_max);
                }
            void Fill(double value){
                hValue->SetBinContent(fEvents+1,value);
                hValueBinned->Fill(value);
                if(fEvents>0){
                    hValueDelta->SetBinContent(fEvents,value-fLastValue);
                    hValueDeltaBinned->Fill(value-fLastValue);
                }
                fLastValue=value;
                ++fEvents;
            }
            void Write(){
                hValue->Write(); hValueDelta->Write();hValueBinned->Write(); hValueDeltaBinned->Write();
            }
        private:
            std::string fName, fTitle, fAxisTitle;
            int fEvents;
            double fLastValue;
            TH1F *hValue, *hValueBinned, *hValueDelta,  *hValueDeltaBinned;
    };

    struct EventData_t{
        int numTrajectories;
        int numHits;
        ProcInfo_t beginEvent, endSimEvent, endPersistEvent;
        typedef const ProcInfo_t& c_procInfoRef;
        EventData_t():numTrajectories(-1),numHits(-1){}
        EventData_t(c_procInfoRef beg, c_procInfoRef end_sim, c_procInfoRef end_pers):
            beginEvent(beg),endSimEvent(end_sim),endPersistEvent(end_pers){}
        void ChooseMax(const EventData_t& rhs){
            if(numHits < rhs.numHits) numHits=rhs.numHits;
            if(numTrajectories < rhs.numTrajectories) numTrajectories=rhs.numTrajectories;
        }
    };
}

class ProfileEvents: public COMET::ICOMETEventLoopFunction {
public:
    ProfileEvents():fTotalEventsRead(0),fTotalEventsWithHits(0) {
        fInput=NULL;
    }

    virtual ~ProfileEvents() {};

    void Usage(void) {
    }


    virtual bool SetOption(std::string option,std::string value="") {
        return false;
    }

    bool operator () (COMET::ICOMETEvent& event) {
        return true;
    }

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

    virtual int Process(ICOMETEvent& event, int outputFiles){
        IHandle<IProfileData> begin_event= event.Get<IProfileData>("proc_profile/begin_event");
        IHandle<IProfileData> end_event= event.Get<IProfileData>("proc_profile/end_event");
        IHandle<IProfileData> end_persist= event.Get<IProfileData>("proc_profile/end_persist_event");
        if(end_event && end_persist && begin_event){

            ++fTotalEventsRead;
            EventData_t data(begin_event->GetProcData(), end_event->GetProcData(), end_persist->GetProcData());
            COMET::IHandle<COMET::IG4TrajectoryContainer> trajectories
                = event.Get<COMET::IG4TrajectoryContainer>("truth/G4Trajectories");
            if(trajectories) data.numTrajectories=trajectories->size();
            else data.numTrajectories=-1;
            COMET::IHandle<COMET::IDataVector> hits 
                = event.Get<COMET::IDataVector>("truth/g4Hits");
            if (hits) {
                data.numHits=hits->size();
                ++fTotalEventsWithHits;
            }
            else data.numHits=-1;
            fMaxData.ChooseMax(data);
            fEvents.push_back(data);
        }
        return -1;
    }

    virtual void Initialize(void){
    }

    virtual void Finalize(ICOMETOutput * const file){

        int numEvents=fEvents.size();

        // Histograms for memory / timing
        HistByEvent hMemVirt("memory_virtual" , "Virtual Memory Usage","Virtual Memory (MB)" , numEvents,35,6e3);
        HistByEvent hMemResident("memory_resident" , "Resident Memory Usage","Resident Memory (MB)" , numEvents,25,6e3);
        HistByEvent hCpuUser("time_user","CPU User Time"," User Time (s)", numEvents , 2,3000 );
        HistByEvent hCpuSys("time_sys","CPU System Time"," System Time (s)", numEvents , 1,100 );

        // histograms for profiling per trajectory
        HistByEvent hTrajectories("num_trajectories","Number of Trajectories per Event","No. Trajectories", numEvents,6e3, 12000,-6e3,5000);
        HistByEvent hHits("num_hits","Number of Hits per Event","No. Hits", numEvents,100, 200);
        TH2F hResidentVTraj("resident_v_traj","Change in Resident Memory Compared to No. Trajectories per Event"
                ";No. Trajectories; #Delta Resident Memory (MB)",300, 0, 12000,300, 0, 5 );
        TH2F hVirtualVTraj("virtual_v_traj","Change in Virtual Memory Compared to No. Trajectories per Event"
                ";No. Trajectories; #Delta Virtual Memory (MB)",300, 0, 12000,300, 0, 5 );
        TH2F hUserPersistingVTraj("user_persist_v_traj","User Time Spent Persisting Data Compared to No. Trajectories per Event"
                ";No. Trajectories; User Time Persisting (s)",300, 0, 12000,300, 0, 0.14 );
        TH2F hUserSimulatingVTraj("user_simulate_v_traj","User Time Spent Simulating Data Compared to No. Trajectories per Event"
                ";No. Trajectories; User Time Simulating (s)",300, 0, 12000,300, 0, 1 );

        double last_res=0, last_virt=0;
        for (std::vector<EventData_t>::const_iterator i_event=fEvents.begin();
                i_event!=fEvents.end(); ++i_event){
            const int bin=i_event - fEvents.begin();
            const double& mem_virt = i_event->endPersistEvent.fMemVirtual /1024.; // Convert to MB
            const double& mem_res  = i_event->endPersistEvent.fMemResident/1024.; // Convert to MB
            const double& cpu_user_persist = i_event->endPersistEvent.fCpuUser;
            const double& cpu_user_sim = i_event->endSimEvent.fCpuUser;
            const double& cpu_user_begin = i_event->beginEvent.fCpuUser;
            const double& cpu_sys  = i_event->endPersistEvent.fCpuSys;
            const int& nTrajs  = i_event->numTrajectories;
            const int& nHits  = i_event->numHits;

            hMemResident.Fill(mem_res);
            hMemVirt.Fill(mem_virt);
            hCpuUser.Fill(cpu_user_persist);
            hCpuSys.Fill(cpu_sys);

            hTrajectories.Fill(nTrajs);
            hHits.Fill(nHits);
            if(bin>0){
                hResidentVTraj.Fill(nTrajs,mem_res-last_res);
                hVirtualVTraj.Fill(nTrajs,mem_virt-last_virt);
                hUserSimulatingVTraj.Fill(nTrajs,cpu_user_sim - cpu_user_begin);
                hUserPersistingVTraj.Fill(nTrajs,cpu_user_persist - cpu_user_sim);
            }
            last_res=mem_res;
            last_virt=mem_virt;
        }

        hMemVirt.Write();
        hMemResident.Write();
        hCpuUser.Write();
        hCpuSys.Write();
        hTrajectories.Write();
        hHits.Write();
        hResidentVTraj.Write();
        hVirtualVTraj.Write();
        hUserSimulatingVTraj.Write();
        hUserPersistingVTraj.Write();

        COMETLog("Total events read: "<<fTotalEventsRead);
        COMETLog("Total events with hits: "<<fTotalEventsWithHits);
    }

private:
    const IVInputFile* fInput;
    long int fTotalEventsRead;
    long int fTotalEventsWithHits;
    std::vector<EventData_t> fEvents;
    std::vector<ProcInfo_t> fBeginEvents;
    std::vector<ProcInfo_t> fEndPersists;
    EventData_t fMaxData;
};

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