#include <map>
#include <cmath>
#include <cstdlib>
#include <sstream>

#include <TROOT.h>
#include <TH1F.h>

#include <IHandle.hxx>
#include <IHitSelection.hxx>
#include <IHit.hxx>
#include <HEPUnits.hxx>

#include <cometEventLoop.hxx>

class ILayerEnergy: public COMET::ICOMETEventLoopFunction {
public:
    ILayerEnergy() {
        fHitSelection = "";
        fEnergyRange = 5*unit::MeV;
        fTimeRange = 5800*unit::ns;
        fErrorCount = 0;
    }

    virtual ~ILayerEnergy() {};

    void Usage(void) {
        std::cout << "    -O hits=<name> The IHitSelection to process."
                  << std::endl;
        std::cout << "    -O range=<val> The histogram range in MeV."
                  << std::endl;
        std::cout << "    -O time=<val> The histogram range in ns."
                  << std::endl;
    }

    virtual bool SetOption(std::string option,std::string value="") {
        if (option == "hits") fHitSelection = value;
        else if (option == "range") {
            std::istringstream input(value);
            input >> fEnergyRange;
        }
        else if (option == "time") {
            std::istringstream input(value);
            input >> fTimeRange;
        }
        else return false;
        return true;
    }

    void Initialize() {
        if (fHitSelection == "") std::exit(1);

        std::string name = fHitSelection+"LayerEnergy";
        std::string title = "Energy deposited in each "+fHitSelection+" layer";
        fLayerEnergy = new TH1F(name.c_str(), title.c_str(),
                                100, 0.0, fEnergyRange);
        fLayerEnergy->SetXTitle("Energy (MeV)");

        name = fHitSelection + "HitTime";
        title = "Time of each " + fHitSelection + " hit";

        int timeBins = int(fTimeRange/(10*unit::ns));

        fHitTimes = new TH1F(name.c_str(), title.c_str(),
                             timeBins, 0.0, fTimeRange);
        fHitTimes->SetXTitle("Time (ns)");
        fHitTimes->SetYTitle("Hits / 10 ns");
    }

    // Histogram the charge in each layer.  This is a touch convoluted since
    // it can't depend on libGeomInfo (which isn't a SimG4 dependency).
    bool operator () (COMET::ICOMETEvent& event) {
        COMET::IHandle<COMET::IHitSelection> hits 
            = event.GetHitSelection(fHitSelection.c_str());
        if (!hits) {
            if (++fErrorCount < 5) {
                COMETError("Missing Hits for " << fHitSelection 
                           << " (" << fErrorCount << ")"
                           << " in Event: " << event.GetContext());
                event.ls();
            }
            return false;
        }
        fErrorCount = 0;

        // The first element of the pair is the Z coordinate.  The second
        // element is the accumulated charge.
        std::map<double,double> layerMap;
        for (COMET::IHitSelection::iterator h = hits->begin();
             h != hits->end();
             ++h) {
            fHitTimes->Fill((*h)->GetTime()/unit::ns);
            double z = (*h)->GetPosition().Z();
            std::map<double,double>::iterator lower 
                = layerMap.lower_bound(z-1*unit::mm);
            std::map<double,double>::iterator upper 
                = layerMap.upper_bound(z+1*unit::mm);
            bool found = false;
            for (std::map<double,double>::iterator i = lower; i != upper; ++i) {
                if (std::abs(i->first-z) < 1*unit::mm) {
                    i->second += (*h)->GetCharge();
                    found = true;
                    break;
                }
            }
            if (!found) layerMap[z] = (*h)->GetCharge();
        }

        for (std::map<double,double>::iterator i = layerMap.begin();
             i != layerMap.end();
             ++i) {
            fLayerEnergy->Fill(i->second);
        }

        return false;
    }

private:
    /// The name of the hit selection to process.
    std::string fHitSelection;

    /// The energy range in a layer
    double fEnergyRange;
    
    // The energy deposited in each layer (actually this is the charge saved
    // in the IMCHit.
    TH1F* fLayerEnergy;

    // The time range for the hits.
    double fTimeRange;

    // The time of the energy deposits.
    TH1F* fHitTimes;

    // Error count
    int fErrorCount;
};

int main(int argc, char **argv) {
    ILayerEnergy userCode;
    cometEventLoop(argc,argv,userCode);
}