#include <map>

#include "TString.h"
#include "TChain.h"
#include "TH2.h"
#include "TMath.h"
#include "TFile.h"

#include "Math/Vector3D.h"

using namespace std;
using namespace ROOT::Math;

static const int maxmult = 100;

struct event_entry {
    int mult;
    float energy[maxmult];
    float x[maxmult];
    float y[maxmult];
    float z[maxmult];

    void reset() {mult=0;}
};

TChain *finputchain = nullptr;

int fevent_bunch_size = 10;

event_entry fcurrent_event_entry;
vector<event_entry> fevents_bunch;

Long64_t fcurrent_entry_number=0;
Long64_t f_entries_to_treat;

map<int,array<TH2F*,2>> fAngCorrHistograms;

int GGBins = 4000;
float GGMin = 0.;
float GGMax = 2000.;

bool get_new_bunch();
void process_angular_correlations();
void process_event_mixing();
void define_new_theta(int theta);

// The global idea is to generate angular correlation matrices.
// Which means, from a list of detectors measuring gamma energies at specific positions (X,Y,Z)
// plotting all the combinations of two measured gammas in the same event, for each angle made by the two fired detectors

// The global issue with this is the normalisation to take into account the geometrical coverage of the detectors for all energy and angle combinations
// the most efficient method is to "correlate uncorrelated gamma-rays" by mixing events together

// the idea in this macro is to read the TTree as input by bunches of 10 entries:
//
// A first standard read of the 10 entries is first done to plot the angular correlation:
// -> for each entry, we loop on all combination of two gamma-rays, calculate the angle and plot the energy of one versus the erergy of the other
// -> This is done in the function: process_angular_correlations()

// The bunch is then read again to generate the event mixing normalisation:
// -> we loop on all combination of two gamma-rays in the 10 entries of the bunch, with asking that gamma1 and gamma2 are coming from different entries
// -> for all combination, we calculate the angle and plot the energy of one versus the erergy of the other
// -> This is done in the function: process_event_mixing()

void macro_example() {

    finputchain = new TChain("mytree");
    finputchain->Add("ExampleTree.root");

    finputchain->SetBranchAddress("Ge_Mult", &fcurrent_event_entry.mult);
    finputchain->SetBranchAddress("Ge_E", fcurrent_event_entry.energy);
    finputchain->SetBranchAddress("Ge_X", fcurrent_event_entry.x);
    finputchain->SetBranchAddress("Ge_Y", fcurrent_event_entry.y);
    finputchain->SetBranchAddress("Ge_Z", fcurrent_event_entry.z);

    f_entries_to_treat = finputchain->GetEntries();

    fevents_bunch.resize(fevent_bunch_size);

    while(true) {
        // until the end of file is not reached, we readn a new bunch of event and store it in fevents_bunch
        bool eof = get_new_bunch();

        // step 1: processing the standard angular correlations
        process_angular_correlations();

        // step 2: event mixing normalisation
        process_event_mixing();

        if(eof) break;
    }

    TFile *fout = new TFile("my_result.root","recreate");
    for(auto &angle: fAngCorrHistograms) {
        angle.second.at(0)->Write();
        angle.second.at(1)->Write();
    }
    fout->Close();
}

bool get_new_bunch()
{
    bool end_of_file = false;

    // reset bunch
    for(auto &evt: fevents_bunch) evt.reset();

    for(int entry=0 ; entry<fevent_bunch_size ; entry++) {
        if(fcurrent_entry_number>=f_entries_to_treat) {
            end_of_file = true;
            break;
        }
        finputchain->GetEntry(fcurrent_entry_number);

        fevents_bunch.at(entry) = fcurrent_event_entry;

        fcurrent_entry_number++;
    }

    return end_of_file;
}

void process_angular_correlations()
{
    for(auto &event: fevents_bunch) {
        for(int i=0 ; i<event.mult ; i++) {
            for(int j=0 ; j<i ; j++) {
                XYZVectorF pos1(event.x[i],event.y[i],event.z[i]);
                XYZVectorF pos2(event.x[j],event.y[j],event.z[j]);
                double angle = TMath::Nint(TMath::ACos(pos1.Dot(pos2) / sqrt(pos1.Mag2() * pos2.Mag2()))*TMath::RadToDeg());

                if(!fAngCorrHistograms.count(angle)) define_new_theta(angle);

                fAngCorrHistograms.at(angle).at(0)->Fill(event.energy[i],event.energy[j]);
                fAngCorrHistograms.at(angle).at(0)->Fill(event.energy[j],event.energy[i]);
            }
        }
    }
}

void process_event_mixing()
{
    for(auto &event_i: fevents_bunch) {
        for(auto &event_j: fevents_bunch) {
            if (&event_i == &event_j) continue;

            for(int i=0 ; i<event_i.mult ; i++) {
                for(int j=0 ; j<event_j.mult ; j++) {
                    XYZVectorF pos1(event_i.x[i],event_i.y[i],event_i.z[i]);
                    XYZVectorF pos2(event_j.x[j],event_j.y[j],event_j.z[j]);
                    double angle = TMath::Nint(TMath::ACos(pos1.Dot(pos2) / sqrt(pos1.Mag2() * pos2.Mag2()))*TMath::RadToDeg());

                    if(!fAngCorrHistograms.count(angle)) define_new_theta(angle);

                    fAngCorrHistograms.at(angle).at(1)->Fill(event_i.energy[i],event_j.energy[j]);
                    fAngCorrHistograms.at(angle).at(1)->Fill(event_j.energy[j],event_i.energy[i]);
                }
            }
        }
    }
}

void define_new_theta(int theta)
{
    TString Name = Form("GxG_Theta_%d",theta);
    TH2F *GGHist = new TH2F(Name,Name,GGBins,GGMin,GGMax,GGBins,GGMin,GGMax);
    TH2F *GGHistNorm = new TH2F(Name.Append("_Norm"),Name,GGBins,GGMin,GGMax,GGBins,GGMin,GGMax);
    fAngCorrHistograms.insert({theta,{GGHist,GGHistNorm}});
}