#include "TFile.h"
#include "TTree.h"
#include "TH1F.h"
#include "TLorentzVector.h"
#include <iostream>
#include <cmath>

double deltaR(double eta1, double phi1, double eta2, double phi2) {
    double dEta = eta1 - eta2;
    double dPhi = std::fmod(phi1 - phi2 + M_PI, 2 * M_PI) - M_PI; // Wrap dPhi to [-π, π]
    return std::sqrt(dEta * dEta + dPhi * dPhi);
}

void compareDeltaR(const char* inputFile, const char* outputFile) {
    TFile* input = TFile::Open(inputFile, "READ");
    if (!input || input->IsZombie()) {
        std::cerr << "Error opening input file: " << inputFile << std::endl;
        return;
    }

    TFile* output = TFile::Open(outputFile, "READ");
    if (!output || output->IsZombie()) {
        std::cerr << "Error opening output file: " << outputFile << std::endl;
        input->Close();
        return;
    }

    TTree* inputTree = dynamic_cast<TTree*>(input->Get("CollectionTree"));
    TTree* outputTree = dynamic_cast<TTree*>(output->Get("CollectionTree"));
    if (!inputTree || !outputTree) {
        std::cerr << "Error: CollectionTree not found in one or both files!" << std::endl;
        input->Close();
        output->Close();
        return;
    }

    std::vector<float>* inputPx = nullptr;
    std::vector<float>* inputPy = nullptr;
    std::vector<float>* inputPz = nullptr;
    std::vector<float>* inputE = nullptr;

    inputTree->SetBranchAddress("HLTNav_RepackedFeatures_ParticleAuxDyn.px", &inputPx);
    inputTree->SetBranchAddress("HLTNav_RepackedFeatures_ParticleAuxDyn.py", &inputPy);
    inputTree->SetBranchAddress("HLTNav_RepackedFeatures_ParticleAuxDyn.pz", &inputPz);
    inputTree->SetBranchAddress("HLTNav_RepackedFeatures_ParticleAuxDyn.e", &inputE);

    std::vector<float>* outputPx = nullptr;
    std::vector<float>* outputPy = nullptr;
    std::vector<float>* outputPz = nullptr;
    std::vector<float>* outputE = nullptr;

    outputTree->SetBranchAddress("HLTNav_RepackedFeatures_ParticleAuxDyn.px", &outputPx);
    outputTree->SetBranchAddress("HLTNav_RepackedFeatures_ParticleAuxDyn.py", &outputPy);
    outputTree->SetBranchAddress("HLTNav_RepackedFeatures_ParticleAuxDyn.pz", &outputPz);
    outputTree->SetBranchAddress("HLTNav_RepackedFeatures_ParticleAuxDyn.e", &outputE);

    TH1F* h_deltaR = new TH1F("h_deltaR", "DeltaR between lossless and lossy particles;DeltaR;Count", 100, 0, 0.2);

    // Compare entries in both trees
    Long64_t nEntries = std::min(inputTree->GetEntries(), outputTree->GetEntries());
    for (Long64_t i = 0; i < nEntries; ++i) {
        inputTree->GetEntry(i);
        outputTree->GetEntry(i);

        // Check data validity
        if (!inputPx || !inputPy || !inputPz || !inputE || !outputPx || !outputPy || !outputPz || !outputE) {
            std::cerr << "Error: Missing four-vector information in event " << i << std::endl;
            continue;
        }

        if (inputPx->size() != outputPx->size()) {
            std::cerr << "Warning: Mismatch in container sizes in event " << i << std::endl;
            continue;
        }

        // Loop over particles
        for (size_t j = 0; j < inputPx->size(); ++j) {
            // Construct lossless and lossy four-vectors
            TLorentzVector lossless, lossy;
            lossless.SetPxPyPzE((*inputPx)[j], (*inputPy)[j], (*inputPz)[j], (*inputE)[j]);
            lossy.SetPxPyPzE((*outputPx)[j], (*outputPy)[j], (*outputPz)[j], (*outputE)[j]);

            // Compute DeltaR
            double dR = deltaR(lossless.Eta(), lossless.Phi(), lossy.Eta(), lossy.Phi());
            h_deltaR->Fill(dR);
        }
    }

    // Analyze DeltaR tail
    int tailCount = h_deltaR->Integral(h_deltaR->FindBin(0.1), h_deltaR->GetNbinsX());
    double meanDeltaR = h_deltaR->GetMean();
    std::cout << "Mean DeltaR: " << meanDeltaR << std::endl;
    std::cout << "Number of particles with DeltaR > 0.1: " << tailCount << std::endl;

    TFile outFile("DeltaRComparison.root", "RECREATE");
    h_deltaR->Write();
    outFile.Close();

    delete h_deltaR;
    input->Close();
    output->Close();
    delete input;
    delete output;

    std::cout << "DeltaR comparison completed. Histogram saved to DeltaRComparison.root." << std::endl;
}

int main(int argc, char** argv) {
    if (argc != 3) {
        std::cerr << "Usage: " << argv[0] << " <input file> <output file>" << std::endl;
        return 1;
    }

    compareDeltaR(argv[1], argv[2]);
    return 0;
}

// Compilation and Execution Instructions:
// Compile: g++ -o compareDeltaR compareDeltaR.cxx $(root-config --cflags --libs)
// Execute: ./compareDeltaR <input_file> <output_file>