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

#include <TROOT.h>
#include <TH1F.h>
#include <TVector3.h>
#include <TGeoManager.h>
#include <TGeoVolume.h>
#include <TGeoShape.h>
#include <TGeoBBox.h>
#include <TProfile.h>

#include <IHandle.hxx>
#include <IHitSelection.hxx>
#include <IHit.hxx>
#include <IMCHit.hxx>
#include <IG4HitSegment.hxx>
#include <HEPUnits.hxx>
#include <IGeomInfo.hxx>
#include <IP0DGeom.hxx>
#include <IOADatabase.hxx>
#include <IGeomIdManager.hxx>

#include <cometEventLoop.hxx>

class IP0DAtten: public COMET::ICOMETEventLoopFunction {
public:
    IP0DAtten() {
        fErrorCount = 0;
        fTrackCount = 0;
    }

    virtual ~IP0DAtten() {};

    void Usage(void) { }

    void Initialize() {
        int bins = 25;
        fP0DXAtten = new TProfile("p0dXAtten", 
                                  "X Layer: Charge vs Distance",
                                  bins, 0.0, 250*unit::cm,
                                  0.0, 1000.0, "e");
        fP0DXAtten->SetYTitle("Mean Number of Avalanches");
        fP0DXAtten->SetXTitle("Distance from MPPC");

        fP0DYAtten = new TProfile("p0dYAtten", 
                                  "Y Layer: Charge vs Distance",
                                  bins, 0.0, 250*unit::cm,
                                  0.0, 1000.0, "e");
        fP0DYAtten->SetYTitle("Mean Number of Avalanches");
        fP0DYAtten->SetXTitle("Distance from MPPC");
    }

    // Histogram the charge in each bar.
    bool operator () (COMET::ICOMETEvent& event) {
        COMET::IHandle<COMET::IHitSelection> hits 
            = event.GetHitSelection("p0d");
        COMET::IOADatabase::Get().Geometry();
        if (!hits) {
            if (++fErrorCount < 5) {
                COMETError("Missing Hits for P0D (" << fErrorCount << ")"
                           << " in Event: " << event.GetContext());
            }
            return false;
        }
        fErrorCount = 0;
        ++fTrackCount;

        for (COMET::IHitSelection::iterator h = hits->begin();
             h != hits->end();
             ++h) {
            double q = (*h)->GetCharge();
            COMET::IHandle<COMET::IMCHit> mcHit = (*h);
            const COMET::IMCHit::ContributorContainer& g4Hits 
                = mcHit->GetContributors();
            TVector3 pos(0,0,0);
            if (g4Hits.size() < 1) continue;
            for (COMET::IMCHit::ContributorContainer::const_iterator g 
                     = g4Hits.begin();
                 g != g4Hits.end();
                 ++g) {
                COMET::IG4HitSegment* seg 
                    = dynamic_cast<COMET::IG4HitSegment*>(*g);
                pos.SetX(pos.X() + seg->GetStartX() + seg->GetStopX());
                pos.SetY(pos.Y() + seg->GetStartY() + seg->GetStopY());
                pos.SetZ(pos.Z() + seg->GetStartZ() + seg->GetStopZ());
            }
            pos = (0.5/(g4Hits.size()))*pos;
            COMET::IGeometryId geomId;
            COMET::IOADatabase::Get().GeomId().GetGeometryId(pos.X(),
                                                          pos.Y(), 
                                                          pos.Z(), geomId);
            COMET::IOADatabase::Get().GeomId().CdId(geomId);
            TGeoBBox* shape 
                = dynamic_cast<TGeoBBox*>(
                    gGeoManager->GetCurrentVolume()->GetShape());
            Double_t master[3]; pos.GetXYZ(master);
            Double_t local[3];
            gGeoManager->MasterToLocal(master, local);
            double dist = shape->GetDZ() - local[2];
            if (COMET::IGeomInfo::P0D().GetBar(geomId).GetLayer()) {
                fP0DYAtten->Fill(dist,q);
            }
            else {
                fP0DXAtten->Fill(dist,q);
            }
        }

        return false;
    }

    void Finalize(COMET::ICOMETOutput * const file) {
        // Correct for oversampling of the muons.
        double factor = std::sqrt(fP0DXAtten->GetEntries()/fTrackCount);
        for (int i = 0; i<fP0DXAtten->GetNbinsX(); ++i) {
            double e = fP0DXAtten->GetBinError(i);
            fP0DXAtten->SetBinError(i,e*factor);
            std::cout << e << std::endl;
        }
        for (int i = 0; i<fP0DYAtten->GetNbinsX(); ++i) {
            double e = fP0DYAtten->GetBinError(i);
            fP0DYAtten->SetBinError(i,e*factor);
        }
    }

private:
    // The light yeild as a function of distance from sensor.
    TProfile* fP0DXAtten;
    TProfile* fP0DYAtten;
    int fTrackCount;

    int fErrorCount;
};

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