#include <cassert>
#include <cmath>


#include <TH1F.h>
#include <TFile.h>

#include <p0dAnalysisLoop.hxx>
#include <IP0DReconModule.hxx>
#include <ITruthVerticesModule.hxx>

#include <IReconBase.hxx>
#include "IP0dReactionMode.hxx"


class IUserProcessLoop {
public:
    typedef ITruthVerticesModule::ITruthVertex ITruthVertex; 
    typedef IP0DReconModule::IP0DVertex IP0DVertex;
    typedef IP0DReconModule::IP0DTrack IP0DTrack;
    typedef IP0DReconModule::IP0DParticle IP0DParticle;
    typedef IP0DReconModule::TP0DVeto TP0DVeto;
    IUserProcessLoop() {}

    void Initialize();
    void Finalize(COMET::ICOMETOutput*);

    bool operator()(const TClonesArray* trueVertices,
                    const TClonesArray* vertices,
                    std::size_t EventId);
    

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

    void Usage() {
        std::cout << "Run the P0D event reconstruction." 
                  << std::endl
                  << std::endl;
    }

private:
    TH1F* hmode;
    TH1F* hmodecut;
    TH1F* hmoderaw;
    TH1F* nhits;
    TH1F* npizerohits;
    TH1F* hfiducial;
    TH1F* hmuoncosine;
    TH1F* htrack;
    TH1F* hshower;
    TH1F* hlength;

    TH1F* hfraction;
    
    TH1F* hpizero;
    TH1F* hpizeroraw;
    TH1F* hpizerocccut;
    TH1F* hpizero2dcut;
    TH1F* hpizerofidcut;
    TH1F* hpizeroshwcut;
    TH1F* hpizeroinvcut;
    TH1F* hpizerodcycut;
    TH1F* hpizerocut;
    
    TH1F* hmuonpraw;
    TH1F* hmuonpcccut;
    TH1F* hmuonpfidcut;
    TH1F* hmuonpshwcut;
    TH1F* hmuonpcut;
    
    TH1F* hinvmassraw;
    TH1F* hinvmasscut;

    TH1F* hmuonx;
    TH1F* hmuony;
    TH1F* hmuonz;
    TH1F* hpizerocosine;
    TH1F* hpizerox;
    TH1F* hpizeroy;
    TH1F* hpizeroz;

    TH1F* hpizerocandidate;
    TH1F* hpizero3cut;
    TH1F* hchargedpi3cut;
    TH1F* hqe3cut;
    TH1F* hother3cut;
    TH1F* hcc3cut;
    TH1F* hnc3cut;

    TH1F* hbkccqe;
    TH1F* hbkcc1pi0;
    TH1F* hbkcc1pi;
    TH1F* hbkccmulti;


    
};

int main(int argc, char* argv[])
{
    IUserProcessLoop userCode;
    p0dAnalysisLoop(argc, argv, userCode);
    
}


void IUserProcessLoop::Initialize()
{
    hmode = new TH1F("mode", "Reaction code", 60, -5, 55);
    hmodecut = new TH1F("modecut", "Reaction code", 60, -5, 55);
    hmoderaw = new TH1F("moderaw", "Reaction code", 60, -5, 55);
    nhits = new TH1F("hits", "Hits #", 100, 0, 100);
    npizerohits = new TH1F("pi0hits", "pi0 hits #", 100, 0, 100);
    hfiducial = new TH1F("fiducial", "Fiducial", 100, -1000.0, 1000.0);
    hmuoncosine = new TH1F("cosine", "Cosine", 25, -1.0, 1.0);
    htrack = new TH1F("trkCount", "Track #", 10, -0.5, 9.5);
    hshower = new TH1F("shwCount", "Shower #", 10, -0.5, 9.5);
    hlength = new TH1F("length", "Track length", 100, 0.0, 1000.0);

    hfraction = new TH1F("frac", "Fraction", 100, -5.0, 5.0);
    
    hpizero = new TH1F("pizero", "Momentum", 25, 0.0, 1000.0);
    hpizeroraw = new TH1F("p3raw", "Momentum", 25, 0.0, 1000.0);
    hpizerocccut = new TH1F("p3cccut", "Momentum", 25, 0.0, 1000.0);
    hpizero2dcut = new TH1F("p32dcut", "Momentum", 25, 0.0, 1000.0);
    hpizerofidcut = new TH1F("p3fidcut", "Momentum", 25, 0.0, 1000.0);
    hpizeroshwcut = new TH1F("p3shwcut", "Momentum", 25, 0.0, 1000.0);
    hpizeroinvcut = new TH1F("p3invcut", "Momentum", 25, 0.0, 1000.0);
    hpizerodcycut = new TH1F("p3dcycut", "Momentum", 25, 0.0, 1000.0);
    hpizerocut = new TH1F("p3cut", "Momentum", 25, 0.0, 1000.0);
    
    hmuonpraw = new TH1F("muonpraw", "Momentum", 25, 0.0, 1000.0);
    hmuonpcccut = new TH1F("muonpcccut", "Momentum", 25, 0.0, 1000.0);
    hmuonpfidcut = new TH1F("muonpfidcut", "Momentum", 25, 0.0, 1000.0);
    hmuonpshwcut = new TH1F("muonpshwcut", "Momentum", 25, 0.0, 1000.0);
    hmuonpcut = new TH1F("muonpcut", "Momentum", 25, 0.0, 1000.0);
    
    hinvmassraw = new TH1F("invraw", "Invariant mass", 50, 0.0, 1000.0);
    hinvmasscut = new TH1F("invcut", "Invariant mass", 50, 0.0, 1000.0);

    hmuonx = new TH1F("x", "vertex x", 50, -1200.0, 1200.0);
    hmuony = new TH1F("y", "vertex y", 50, -1200.0, 1200.0);
    hmuonz = new TH1F("z", "vertex z", 100, -3500.0, -500.0);
    hpizerocosine = new TH1F("pi0cosine", "Cosine", 25, -1.0, 1.0);
    hpizerox = new TH1F("pi0x", "vertex x", 50, -1200.0, 1200.0);
    hpizeroy = new TH1F("pi0y", "vertex y", 50, -1200.0, 1200.0);
    hpizeroz = new TH1F("pi0z", "vertex z", 100, -3500.0, -500.0);

    hpizerocandidate = new TH1F("all3cut", "Energy", 25, 0.0, 1000.0);
    hpizero3cut = new TH1F("3cut", "Energy", 25, 0.0, 1000.0);
    hchargedpi3cut = new TH1F("chargedpi3cut", "Energy", 25, 0.0, 1000.0);
    hqe3cut = new TH1F("qe3cut", "Energy", 25, 0.0, 1000.0);
    hother3cut = new TH1F("other3cut", "Energy", 25, 0.0, 1000.0);
    hcc3cut = new TH1F("cc3cut", "Energy", 25, 0.0, 1000.0);
    hnc3cut = new TH1F("nc3cut", "Energy", 25, 0.0, 1000.0);

    hbkccqe = new TH1F("ccqe", "Energy", 25, 0.0, 1000.0);
    hbkcc1pi0 = new TH1F("cc1pi0", "Energy", 25, 0.0, 1000.0);
    hbkcc1pi = new TH1F("cc1pi", "Energy", 25, 0.0, 1000.0);
    hbkccmulti = new TH1F("ccmulti", "Energy", 25, 0.0, 1000.0);

    
}

bool IUserProcessLoop::operator()(const TClonesArray* trueVertices,
                                 const TClonesArray* vertices,
                                 std::size_t EventId)
{
    ITruthVertex* first = dynamic_cast<ITruthVertex*>(trueVertices->At(0));
    bool fiducial = first->Fiducial > 0.0;
    IP0dReactionMode primVertex(first);
    TVector3 vertexPosition = primVertex.Position().Vect();
    double vertexTime = primVertex.Position().T();
    size_t code = primVertex.ReactionCode();
    size_t pizeroCount = primVertex.PizeroCount();
    double pizeroMomentum = primVertex.PizeroMomentum().P();
    double pizeroEnergy = primVertex.PizeroMomentum().E();
    size_t muonCount = primVertex.MuonCount();
    TLorentzVector p4 = primVertex.MuonMomentum();
    double muonMomentum = p4.P();
    size_t chargedpiCount = primVertex.ChargedPiCount();
    size_t chargedKaonCount = primVertex.ChargedKaonCount();
    const size_t chargedMesonCount = chargedpiCount + chargedKaonCount;
    bool noChargedParticle = muonCount == 0 && chargedMesonCount == 0;
    bool singleNCpi0 = pizeroCount == 1 && noChargedParticle && fiducial;
    
    hmoderaw->Fill(code);
    if (singleNCpi0 && fiducial) hpizero->Fill(pizeroEnergy);
    
    
    const std::size_t vertexCount = vertices->GetEntriesFast();
    for (std::size_t vertexIndex = 0;
         vertexIndex < std::min(vertexCount,std::size_t(1)); ++vertexIndex) {
        TObject* object = vertices->At(vertexIndex);
        IP0DVertex* vertex = dynamic_cast<IP0DVertex*>(object);
        assert(vertex);
        
        unsigned int NShowers = vertex->NShowers;
        unsigned int Nparticles = vertex->Nparticles;
        unsigned int Ndecay = vertex->Ndecay;
        double fiducialRecon = vertex->Fiducial;
        
        bool ccEvent = false;
        size_t unmatchedCount = 0;
        size_t protonCount = 0;
        std::vector<IP0DTrack>& tracks = vertex->Tracks;
        for (std::vector<IP0DTrack>::iterator t = tracks.begin();
             t != tracks.end(); ++t) {
            if (t->TrackName == "unmatched") {
                const double length = t->Length;
                if (length > 400.0) ++unmatchedCount;
                hlength->Fill(t->Length);
            }
            else {
                double protonLikelihood = t->ProtonLikelihood;
                double muonLikelihood = t->MuonLikelihood;
                double deltaLikelihood = muonLikelihood - protonLikelihood;
                
                    // For track that failed track fitting, => muon-like.
                if (t->Status != COMET::IReconBase::kSuccess)
                        deltaLikelihood = 10.0;
                if (deltaLikelihood > 5.0) { ccEvent = true;  break; }
                else { ++protonCount; }
            }
            
        }

        
            // Loop over all particles, but now there is only pizero
        double momentum = 0.0;
        double energy = 0.0;
        std::vector<IP0DParticle>& particles = vertex->Particles;
        std::vector<IP0DParticle>::iterator pt = particles.begin();
        for ( ; pt != particles.end(); ++pt) { 
            momentum = pt->ParticleMomentum.Mag();
            energy = pt->ParticleEnergy;
            break;
        }
        const double k = 0.23;
        energy *= k;
        momentum *= k;
        double mgg = std::sqrt(energy*energy - momentum*momentum);
            //std::cout << "\t Invariant mass " << mgg << std::endl;
        
        if (singleNCpi0) hpizeroraw->Fill(pizeroMomentum);
        if (muonCount == 1) hmuonpraw->Fill(muonMomentum);
        
            // Make pi0 selection
            // Reject event with muon-like track
        if (ccEvent) continue;
        if (singleNCpi0) hpizerocccut->Fill(pizeroMomentum);
        
        if (unmatchedCount != 0) continue;
        if (singleNCpi0) hpizero2dcut->Fill(pizeroMomentum);
        
        if (muonCount == 1) hmuonpcccut->Fill(muonMomentum);

            // Reject event with reconstructed vertex outside
            // the fiducial vol.
        if (fiducialRecon < 0.0) continue;
        
        if (singleNCpi0) hpizerofidcut->Fill(pizeroMomentum);
        if (muonCount == 1) hmuonpfidcut->Fill(muonMomentum);
        
        hshower->Fill(NShowers);
            // Reject event with too few or too many showers
        if (NShowers < 3 || NShowers > 5) continue;
        
        if (singleNCpi0) hpizeroshwcut->Fill(pizeroMomentum);
        if (muonCount == 1) hmuonpshwcut->Fill(muonMomentum);
        
        double& epi0 = pizeroEnergy;
        if (singleNCpi0 && energy)  hfraction->Fill((epi0-energy)/epi0);
        
        hinvmassraw->Fill(mgg);
        if (singleNCpi0) hinvmasscut->Fill(mgg);
        
            // Reject event using the pi0 reconstructed invariant mass.
        if (mgg < 70.0 || mgg > 200.0) continue;
        if (singleNCpi0) hpizeroinvcut->Fill(pizeroMomentum);
        
        if (Ndecay >= 3) continue;
        if (singleNCpi0) hpizerodcycut->Fill(pizeroMomentum);


        
        TP0DVeto veto = vertex->VetoObject;
            //if (veto.tpcNHits > 2000) continue;
        
        hmodecut->Fill(code);
        if (singleNCpi0) hpizerocut->Fill(pizeroMomentum);
        if (muonCount == 1)  {
            hmuonpcut->Fill(muonMomentum);
            
            hmuoncosine->Fill(p4.Pz()/p4.P());
            hmuonx->Fill(vertexPosition.X());
            hmuony->Fill(vertexPosition.Y());
            hmuonz->Fill(vertexPosition.Z());
        }

        TLorentzVector pi04 = primVertex.PizeroMomentum();
        if (singleNCpi0) {
            hpizerocosine->Fill(pi04.Pz()/pi04.P());
            hpizerox->Fill(vertexPosition.X());
            hpizeroy->Fill(vertexPosition.Y());
            hpizeroz->Fill(vertexPosition.Z());
        }

        
        hpizerocandidate->Fill(energy);
                //if (singleNCpi0) hpizero3cut->Fill(energy);
        if (pizeroCount == 1 && noChargedParticle)
                hpizero3cut->Fill(energy);
        else if (muonCount == 1 &&
                 (chargedpiCount == 1 || pizeroCount == 1))
                hchargedpi3cut->Fill(energy);
        else if (muonCount == 1 && chargedpiCount == 0 && pizeroCount == 0)
                hqe3cut->Fill(energy);
        else
                hother3cut->Fill(energy);


        bool background = !(pizeroCount == 1 && noChargedParticle);
        bool cc = muonCount == 1;
        bool nc = muonCount == 0;
        size_t piCount = pizeroCount + chargedpiCount;
        if (background && cc) {
            if (pizeroCount == 0 && chargedpiCount == 0 &&
                chargedKaonCount == 0) {
                std::cout << "muon " << muonCount
                          << " pizero " << pizeroCount
                          << " pi+/- " << chargedpiCount
                          << " code " << code
                          << std::endl;
                hbkccqe->Fill(energy);
            }
            else if (pizeroCount == 1 &&  chargedpiCount == 0
                     && chargedKaonCount == 0) {
                hbkcc1pi0->Fill(energy);
            }
            else if (pizeroCount == 0 && chargedpiCount == 1 &&
                     chargedKaonCount == 0) {
                hbkcc1pi->Fill(energy);
            }
            else {
                hbkccmulti->Fill(energy);
            }
        }
                    
        if (background && nc)
                std::cout << "ncbackground " << EventId 
                          << "     " << static_cast<std::size_t>(vertexTime)
                          << "   " <<  first->LeptonPDGs.size()
                          << std::endl;

        
        if (background && cc) hcc3cut->Fill(energy);
        else if (background && nc) hnc3cut->Fill(energy);
        else {}


    }
        

    return false;
}

void IUserProcessLoop::Finalize(COMET::ICOMETOutput*)
{
    TFile f("output.root", "recreate");

    nhits->Write();
    npizerohits->Write();
    hfiducial->Write();
    
    htrack->Write();
    hshower->Write();
    hlength->Write();

    hfraction->Write();
    
    hmode->Write();
    hmoderaw->Write();
    hmodecut->Write();
    
    hpizeroraw->Write();
    hpizerocccut->Write();
    hpizero2dcut->Write();
    hpizerofidcut->Write();
    hpizeroshwcut->Write();
    hpizeroinvcut->Write();
    hpizerodcycut->Write();
    hpizerocut->Write();
    
    hmuonpraw->Write();
    hmuonpcccut->Write();
    hmuonpfidcut->Write();
    hmuonpshwcut->Write();
    hmuonpcut->Write();
    hinvmassraw->Write();
    hinvmasscut->Write();
    
    hmuoncosine->Write();
    hmuonx->Write();
    hmuony->Write();
    hmuonz->Write();
    hpizerocosine->Write();
    hpizerox->Write();
    hpizeroy->Write();
    hpizeroz->Write();


    hpizerocandidate->Write();
    hpizero3cut->Write();
    hchargedpi3cut->Write();
    hqe3cut->Write();
    hother3cut->Write();
    hcc3cut->Write();
    hnc3cut->Write();


    hbkccqe->Write();
    hbkcc1pi0->Write();
    hbkcc1pi->Write();
    hbkccmulti->Write();
    
    f.Close();
  
}