#include <iostream>
#include <sstream>
#include <cassert>
#include <map>
#include <cmath>

#include <TTree.h>
#include <TChain.h>
#include <TClonesArray.h>
#include <TFile.h>
#include <TH1F.h>
#include <TCanvas.h>


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


/// Application that reads the dst output of RunOAAnalysis.exe and tries to
/// select pi0 candidates.
/// Yoshi: This will be moved to the macros directory soon.
using namespace COMET;

/// A utility class used inside the P0D analysis readout application.
class ICounterForP0DAnalysisReadoutApp {
public:
    ICounterForP0DAnalysisReadoutApp() : fSignalCount(0), fCount(0) { }
    
    void SignalCount() { ++fSignalCount; }
    void Count() { ++fCount; }

    unsigned int fSignalCount;
    unsigned int fCount;
};


int main(int argc, const char* argv[])
{
    if (argc < 2) { std::cout << "Usage: P0D.exe file.root. Wildcard supported"
                              << std::endl; exit(1); }
    

    std::string fileName(argv[1]);
    std::cout << "Processing file " << argv[1] << std::endl;

    TChain* eventTree = new TChain("TruthDir/Vertices");
    TChain* reconTree = new TChain("ReconDir/P0D");
    
    eventTree->Add(fileName.c_str());
    reconTree->Add(fileName.c_str());
    
    std::cout << "Entries " << eventTree->GetEntries() << std::endl;

        // Access the true vertex to get interaction mode
    unsigned int NVtxP0D = 0;
    TClonesArray* trueVertices
        = new TClonesArray("COMET::ITruthVerticesModule::ITruthVertex");
    
    eventTree->SetBranchAddress("NVtxP0D", &NVtxP0D);
    eventTree->SetBranchAddress("VtxP0D", &trueVertices);
    
        // Access the reconstruction result
    unsigned int NHits = 0;
    double totalCharge = 0.0;
    double usedCharge = 0.0;
    unsigned int NShowerHits = 0;
    double showerCharge = 0.0;
    unsigned int NShowers = 0;
    unsigned int NVertices = 0;
    unsigned int NTracks = 0;
    double energy = 0.0;
    TVector3* momentum = new TVector3(0.0, 0.0, 0.0);
 
 
    unsigned int EventId = 0;
    TClonesArray* tracks 
        = new TClonesArray("COMET::IP0DReconModule::IP0DTrack");
    
    TClonesArray* vertices
        = new TClonesArray("COMET::IP0DReconModule::IP0DVertex");

    TClonesArray* showers
        = new TClonesArray("COMET::IP0DReconModule::IP0DShower");
    
    reconTree->SetBranchAddress("TotalHits", &NHits);
    reconTree->SetBranchAddress("TotalCharge", &totalCharge);
    reconTree->SetBranchAddress("UsedCharge", &usedCharge);
    reconTree->SetBranchAddress("ShowerHits", &NShowerHits);
    reconTree->SetBranchAddress("ShowerCharge", &showerCharge);
    reconTree->SetBranchAddress("NShowers", &NShowers);
    reconTree->SetBranchAddress("NVertices", &NVertices);
    reconTree->SetBranchAddress("NTracks", &NTracks);
    reconTree->SetBranchAddress("EventID", &EventId);
    reconTree->SetBranchAddress("Tracks", &tracks);
    reconTree->SetBranchAddress("Vertices", &vertices);
    reconTree->SetBranchAddress("Showers", &showers);
    reconTree->SetBranchAddress("Energy", &energy);
    reconTree->SetBranchAddress("Momentum", &momentum);

    TH1F* hmode = new TH1F("mode", "Reaction code", 60, -5, 55);
    TH1F* nhits = new TH1F("hits", "Hits #", 100, 0, 300);
    TH1F* npizerohits = new TH1F("pi0hits", "pi0 hits #", 100, 0, 100);
    TH1F* htotalQ = new TH1F("totalQ", "Total Q", 100, 0.0, 2000.0);
    TH1F* husedQ = new TH1F("usedQ", "Used Q", 100, 0.0, 2000.0);
    TH1F* hratioQ = new TH1F("ratioQ", "Ratio Q", 100, 0.0, 1.0);
    TH1F* hfiducial = new TH1F("fiducial", "Fiducial", 100, -1000.0, 1000.0);
    TH1F* hinvariantmass = new TH1F("inv", "Invariant mass", 25, 0.0, 500.0);
    TH1F* hcosine = new TH1F("cosine", "Cosine", 25, -1.0, 1.0);

    std::map<unsigned int, ICounterForP0DAnalysisReadoutApp> counter;
    counter[0] = ICounterForP0DAnalysisReadoutApp();
    counter[1] = ICounterForP0DAnalysisReadoutApp();
    counter[2] = ICounterForP0DAnalysisReadoutApp();
    counter[3] = ICounterForP0DAnalysisReadoutApp();
    counter[4] = ICounterForP0DAnalysisReadoutApp();
    
    for (unsigned int entry = 0; entry < reconTree->GetEntries(); ++entry) {
        reconTree->GetTree()->GetEntry(entry);
        eventTree->GetTree()->GetEntry(entry);
        std::cout << "Entry " << entry
                  << " Tracks " << NTracks
                  << std::endl;

            // Select single interaction event.
        if (NVtxP0D != 1) continue;
        if (NVertices != 1) continue;

        typedef ITruthVerticesModule::ITruthVertex ITruthVertex;
        ITruthVertex* first = dynamic_cast<ITruthVertex*>(trueVertices->At(0));
        std::stringstream name(first->ReactionCode);
        std::string generator;
        unsigned int code = 0;
        name >> generator >> code;
        std::cout << "\t reaction code " << code << std::endl;

              
        bool singleNCpi0 = code == 31 || code == 32 || code == 36;
        
           
        TObject* object = vertices->At(0);
        typedef IP0DReconModule::IP0DVertex IP0DVertex;
        IP0DVertex* vertex = dynamic_cast<IP0DVertex*>(object);
        assert(vertex);
        double fiducial = vertex->Fiducial;

        bool ccEvent = false;
        for (unsigned int trackCount = 0; trackCount < NTracks; ++trackCount){
            TObject* object = tracks->At(trackCount);
            assert(object);
            typedef IP0DReconModule::IP0DTrack IP0DTrack;
            IP0DTrack* track = dynamic_cast<IP0DTrack*>(object);
            double protonLikelihood = track->ProtonLikelihood;
            double muonLikelihood = track->MuonLikelihood;
            double deltaLikelihood = muonLikelihood - protonLikelihood;
            if (deltaLikelihood > 5.0) { ccEvent = true; break; }
        }
        
        double showerToVertex = 0.0*unit::m;
        for (unsigned int showerCount = 0;
             showerCount < NShowers; ++showerCount) {
            typedef IP0DReconModule::IP0DShower IP0DShower;
            TObject* object = showers->At(showerCount);
            IP0DShower* shower = dynamic_cast<IP0DShower*>(object);
            assert(shower);
            TVector3 vertexPos = vertex->Position.Vect();
            TVector3 showerPos = shower->Position.Vect();
            double distance = (vertexPos - showerPos).Mag();
            showerToVertex = std::max(showerToVertex, distance);
        }

        double p = momentum->Mag();
        double invariantMass = 0.35 * std::sqrt(energy*energy - p*p);
        double cosine = momentum->Z()/momentum->Mag();
        
        
        nhits->Fill(NHits);
        htotalQ->Fill(totalCharge);
        husedQ->Fill(usedCharge);
        hratioQ->Fill(usedCharge/totalCharge);
        hfiducial->Fill(fiducial);
        
        assert(usedCharge/totalCharge <= 1.0);
        
        counter[0].Count();
        if (singleNCpi0 && first->Fiducial > 0.0) counter[0].SignalCount();
        
        if (totalCharge < 200.0) continue;
        if (usedCharge/totalCharge < 0.8) continue;
        counter[1].Count();
        if (singleNCpi0 && first->Fiducial > 0.0) counter[1].SignalCount();
        
        if (fiducial < 0.0) continue; 
        counter[2].Count();
        if (singleNCpi0 && first->Fiducial > 0.0) counter[2].SignalCount();
        
        if (ccEvent) continue;
        if (NTracks != 0) continue; 
        counter[3].Count();
        if (singleNCpi0 && first->Fiducial > 0.0) counter[3].SignalCount();

        if (NShowers < 3 || NShowers > 5) continue;
        if (showerToVertex < 10*unit::cm) continue;

        if (invariantMass < 40.0 || invariantMass > 200.0) continue;
        
        counter[4].Count();
        if (singleNCpi0 && first->Fiducial > 0.0) counter[4].SignalCount();
        
        
        
        hmode->Fill(code);
        npizerohits->Fill(NHits);
        hinvariantmass->Fill(invariantMass);
        hcosine->Fill(cosine);
        
    }

    std::cout << "Start pizero # " << counter[0].fSignalCount << std::endl;
    std::cout << "Stop pizero # " << counter[4].fSignalCount << std::endl;
    std::cout << "\t Efficiency " << "      Purity " << std::endl;

    unsigned int totalNpi0 = counter[0].fSignalCount;
    std::cout.precision(2);
    std::cout.setf(std::ios_base::fixed);
    for (std::map<unsigned int, ICounterForP0DAnalysisReadoutApp>::iterator c = counter.begin();
         c != counter.end(); ++c) {
        unsigned int Npi0 = c->second.fSignalCount;
        unsigned int Nany = c->second.fCount;
        std::cout.width(6);
        std::cout << "\t " << 100.0 * Npi0/totalNpi0
                  << "\t        " << 100.0 * Npi0/Nany << std::endl;
    }
    
    TH1F* hefficiency = new TH1F("efficiency", "", 10, 0, 10);
    TH1F* hpurity = new TH1F("purity", "", 10, 0, 10);
    
    for (std::map<unsigned int, ICounterForP0DAnalysisReadoutApp>::const_iterator c = counter.begin();
         c != counter.end(); ++c) {
        unsigned int Npi0 = c->second.fSignalCount;
        unsigned int Nany = c->second.fCount;
        double efficiency = 100.0 * Npi0/totalNpi0;
        double purity = 100.0 * Npi0/Nany;
        hefficiency->SetBinContent(c->first + 1, efficiency);
        hpurity->SetBinContent(c->first + 1, purity);
    }
    
    TFile f("output.root", "recreate");

    hmode->Write();
    nhits->Write();
    npizerohits->Write();
    htotalQ->Write();
    husedQ->Write();
    hratioQ->Write();
    hfiducial->Write();
    hinvariantmass->Write();
    hcosine->Write();

    TCanvas canvas("efficiency & purity");

    hefficiency->SetLineColor(kBlue);
    hefficiency->SetLineWidth(2);
    hpurity->SetLineColor(kRed);
    hpurity->SetLineWidth(2);
    hefficiency->Draw();
    hpurity->Draw("same");

    canvas.Update();
    canvas.Write();

    
    f.Close();
    
    exit(0);
}