#include "ROOTHeaders.h"
#include "CPPHeaders.h"
#include "Constants.h"

using namespace std;

double calculateKineticEnergy(double Px, double Py, double Pz);

void lateralDist(){

    TFile *fEpos = new TFile("epos_100tev_proton.root", "READ");

    TProfile *pEpos = new TProfile("EPOS-LHC", "", 30, 0, 150);

    TTree *tEpos = (TTree *) fEpos->Get("cor");

    TLeaf *lEposPId = tEpos->GetLeaf("PId");

    int entEpos   = tEpos->GetEntries();

    double kineticEnergy, muPosition, numDensityEpos;
    int muNumberEpos, showerRadiusEpos;

    map<int, int> particleInSideByCircleRadiusEpos;
    for (int k = 0; k <= 195; k += 5) {
        particleInSideByCircleRadiusEpos[k] = 0;
    }

    for (int i = 0; i < entEpos; i++) {
        tEpos->GetEntry(i);

        int leafEntry = lEposPId->GetNdata();

        for (int iPart = 0; iPart < leafEntry; iPart++) {
            int pID = lEposPId->GetValue(iPart);

            if( pID == 5 || pID == 6) {
                double Px = tEpos->GetLeaf("Px")->GetValue(iPart);
                double Py = tEpos->GetLeaf("Py")->GetValue(iPart);
                double Pz = tEpos->GetLeaf("Pz")->GetValue(iPart);

                kineticEnergy = calculateKineticEnergy(Px, Py, Pz);

                if (kineticEnergy >= 1) {
                    double muX = tEpos->GetLeaf("X")->GetValue(iPart);
                    double muY = tEpos->GetLeaf("Y")->GetValue(iPart);

                    muPosition = std::sqrt((pow((muX / 100), 2)) + (pow((muY / 100), 2)));
                }
            }   // Muon particle loop
        }   // All particles loop

        // counts all muons within the radius at every 5m circle
        for (int k = 0; k <= 195; k += 5){
            if (muPosition > k && muPosition <= (k+5)) {
                particleInSideByCircleRadiusEpos[k]++;
            }
        }

    }   // Shower loop


    for (int k = 0; k <= 195; k += 5) {
        muNumberEpos = particleInSideByCircleRadiusEpos[k];
        int radiusSquare = (pow((k+5), 2) - pow(k, 2));
        double area = M_PI * radiusSquare;
        numDensityEpos = muNumberEpos / area;
        showerRadiusEpos = k;

        pEpos->Fill(showerRadiusEpos, numDensityEpos);
    }

    TCanvas *canvas = new TCanvas("final_canvas", " ", 1200, 800);
    canvas->SetLogy();
    pEpos->SetMarkerColor(kBlue);
    pEpos->SetMarkerSize(2.5);
    pEpos->SetMarkerStyle(33);
    pEpos->SetStats(1);
    pEpos->GetXaxis()->SetTitle("#bf{Shower Radius (m)}");
    pEpos->GetXaxis()->SetTitleSize(0.05);
    pEpos->GetXaxis()->SetRangeUser(0, 150);
    pEpos->GetYaxis()->SetTitle("#bf{#rho (m^{-2})}");
    pEpos->GetYaxis()->SetTitleSize(0.05);
    pEpos->GetYaxis()->SetTitleOffset(0.55);
//    pEpos->SetTitle("Lateral Distributions of Muons from 100 TeV Proton Shower");
    pEpos->Draw("HIST P");

    TCanvas *c1 = new TCanvas("c1", " ", 1200, 800);
    c1->SetLogy();
    pEpos->SetMarkerColor(kBlue);
    pEpos->SetMarkerSize(2.5);
    pEpos->SetMarkerStyle(33);
    pEpos->SetStats(1);
    pEpos->GetXaxis()->SetTitle("#bf{Shower Radius (m)}");
    pEpos->GetXaxis()->SetTitleSize(0.05);
    pEpos->GetXaxis()->SetRangeUser(0, 5);
    pEpos->GetYaxis()->SetTitle("#bf{#rho (m^{-2})}");
    pEpos->GetYaxis()->SetTitleSize(0.05);
    pEpos->GetYaxis()->SetTitleOffset(0.55);
    pEpos->Draw("HIST P");


    TCanvas *c2 = new TCanvas("c2", " ", 1200, 800);
    c2->SetLogy();
    pEpos->SetMarkerColor(kBlue);
    pEpos->SetMarkerSize(2.5);
    pEpos->SetMarkerStyle(33);
    pEpos->SetStats(1);
    pEpos->GetXaxis()->SetTitle("#bf{Shower Radius (m)}");
    pEpos->GetXaxis()->SetTitleSize(0.05);
    pEpos->GetXaxis()->SetRangeUser(5, 10);
    pEpos->GetYaxis()->SetTitle("#bf{#rho (m^{-2})}");
    pEpos->GetYaxis()->SetTitleSize(0.05);
    pEpos->GetYaxis()->SetTitleOffset(0.55);
    pEpos->Draw("HIST P");


    TCanvas *c3 = new TCanvas("c3", " ", 1200, 800);
    c3->SetLogy();
    pEpos->SetMarkerColor(kBlue);
    pEpos->SetMarkerSize(2.5);
    pEpos->SetMarkerStyle(33);
    pEpos->SetStats(1);
    pEpos->GetXaxis()->SetTitle("#bf{Shower Radius (m)}");
    pEpos->GetXaxis()->SetTitleSize(0.05);
    pEpos->GetXaxis()->SetRangeUser(10, 15);
    pEpos->GetYaxis()->SetTitle("#bf{#rho (m^{-2})}");
    pEpos->GetYaxis()->SetTitleSize(0.05);
    pEpos->GetYaxis()->SetTitleOffset(0.55);
    pEpos->Draw("HIST P");

    TCanvas *c4 = new TCanvas("c4", " ", 1200, 800);
    c4->SetLogy();
    pEpos->SetMarkerColor(kBlue);
    pEpos->SetMarkerSize(2.5);
    pEpos->SetMarkerStyle(33);
    pEpos->SetStats(1);
    pEpos->GetXaxis()->SetTitle("#bf{Shower Radius (m)}");
    pEpos->GetXaxis()->SetTitleSize(0.05);
    pEpos->GetXaxis()->SetRangeUser(15, 20);
    pEpos->GetYaxis()->SetTitle("#bf{#rho (m^{-2})}");
    pEpos->GetYaxis()->SetTitleSize(0.05);
    pEpos->GetYaxis()->SetTitleOffset(0.55);
    pEpos->Draw("HIST P");
}

double calculateKineticEnergy(double Px, double Py, double Pz) {
    double PSquared = (pow(Px, 2) + pow(Py, 2) + pow(Pz, 2));
    return std::sqrt(PSquared + pow(MASS_MUON, 2)) - MASS_MUON;
}