// analyze_apx.C
//
// What this macro does (step-by-step, in plain language):
// 1) Opens your Allpix² ROOT file.
// 2) Reads PixelHit for a chosen detector branch (e.g. "T0").
//    - Groups neighbouring pixels (8-connected) into clusters per event.
//    - Fills a histogram of cluster sizes (how many pixels per cluster).
// 3) Reads PixelCharge for the same detector branch.
//    - Sums collected charge (electrons) per event.
//    - Converts to energy using 3.6 eV per electron and fills a histogram in keV.
//
// NOTES:
// - Run it in ROOT *after* loading the Allpix² object dictionary library.
// - It does not need DepositedCharge; if you later write DepositedCharge, we can add a branch to use it.
//
// HOW TO RUN (example):
//   root -l
//   gSystem->AddDynamicPath("/data/Abdul/allpix-squared/lib");
//   gSystem->Load("libAllpixObjects.so");
//   .x analyze_apx.C("root_data_set.root", "T0")
//
// If you don't know your lib path, find it with:
//   !find /data/Abdul/allpix-squared -name libAllpixObjects.so
//
// If you want 4-connected clusters, set eight_connected=false when calling.

#include <TFile.h>
#include <TTree.h>
#include <TBranch.h>
#include <TH1I.h>
#include <TH1D.h>
#include <TCanvas.h>

#include <vector>
#include <unordered_map>
#include <queue>
#include <string>
#include <iostream>
#include <cmath>

// We use the Allpix² classes via the loaded dictionary (no headers needed here)
namespace allpix {
  class PixelHit;
  class PixelCharge;
}

namespace {
  inline long long key_cr(int c, int r) { return ( (long long)c << 32 ) | (unsigned int)r; }
}

// ---- main function ----
void analyze_allpix(const char* filename = "root_data_set.root",
                 const char* det_branch = "T0",
                 bool eight_connected = true) {
  // 0) Open file
  TFile* f = TFile::Open(filename);
  if(!f || f->IsZombie()){
    std::cout << "[ERR] Cannot open: " << filename << "\n";
    return;
  }
  std::cout << "[OK ] Opened: " << filename << "\n";

  // 1) Get PixelHit tree + branch
  TTree* thit = (TTree*)f->Get("PixelHit");
  if(!thit){
    std::cout << "[ERR] No 'PixelHit' tree in file.\n";
    f->GetListOfKeys()->Print();
    return;
  }
  TBranch* bhit = thit->FindBranch(det_branch);
  if(!bhit){
    std::cout << "[ERR] 'PixelHit' has no branch '" << det_branch << "'. Available:\n";
    thit->GetListOfBranches()->Print();
    return;
  }
  std::vector<allpix::PixelHit*>* hits = nullptr;
  bhit->SetObject(&hits);

  // 2) Get PixelCharge tree + branch (for energy)
  TTree* tchg = (TTree*)f->Get("PixelCharge");
  if(!tchg){
    std::cout << "[WARN] No 'PixelCharge' tree. Energy histogram will be empty.\n";
  }
  TBranch* bchg = tchg ? tchg->FindBranch(det_branch) : nullptr;
  std::vector<allpix::PixelCharge*>* chgs = nullptr;
  if(bchg) bchg->SetObject(&chgs);
  else     std::cout << "[WARN] 'PixelCharge' has no branch '" << det_branch << "'.\n";

  // 3) Histograms
  // Cluster size: 0..30 pixels (tweak if you expect larger clusters)
  TH1I* h_cluster = new TH1I("h_cluster_size",
                             "Cluster size (8-connected);pixels per cluster;clusters",
                             31, -0.5, 30.5);

  // Event energy from collected charge (keV)
  TH1D* h_Eevt = new TH1D("h_event_energy",
                          "Collected energy per event;E_{col} [keV];events",
                          400, 0, 400); // adjust range as needed

  // 4) Event loop (we assume both trees have same number of entries; if not, use min)
  const Long64_t nEvents_hit = thit->GetEntries();
  const Long64_t nEvents_chg = (tchg ? tchg->GetEntries() : 0);
  const Long64_t nEvents = std::max(nEvents_hit, nEvents_chg);

  std::cout << "[.. ] Events (PixelHit)   : " << nEvents_hit << "\n";
  if(tchg) std::cout << "[.. ] Events (PixelCharge): " << nEvents_chg << "\n";

  // neighbour offsets
  const int dc8[8] = {-1, 0, 1, -1, 1, -1, 0, 1};
  const int dr8[8] = {-1,-1,-1,  0, 0,  1, 1, 1};
  const int dc4[4] = { 0,-1, 1,  0};
  const int dr4[4] = {-1, 0, 0,  1};

  // Loop
  for(Long64_t ievt=0; ievt<nEvents; ++ievt){
    if(ievt % 100 == 0) {
      std::cout << "  - event " << ievt << "/" << nEvents << "\n";
    }

    // --- A) Cluster size from PixelHit ---
    if(ievt < nEvents_hit){
      thit->GetEntry(ievt);

      // Build map (col,row) -> compact index
      std::unordered_map<long long,int> idx;
      std::vector<int> cols, rows;
      cols.reserve(hits ? hits->size() : 0);
      rows.reserve(hits ? hits->size() : 0);

      if(hits){
        idx.reserve(hits->size()*2);
        for(size_t i=0;i<hits->size();++i){
          // get pixel indices (integer column,row)
          auto& pix = (*hits)[i]->getPixel();
          int c = pix.getIndex().x();
          int r = pix.getIndex().y();
          long long k = key_cr(c,r);
          // avoid duplicates (should not happen for hits, but be safe)
          if(idx.find(k) == idx.end()){
            idx[k] = (int)cols.size();
            cols.push_back(c);
            rows.push_back(r);
          }
        }
      }

      // BFS for connected components
      std::vector<char> used(cols.size(),0);
      for(size_t i=0;i<cols.size();++i){
        if(used[i]) continue;
        int csize = 0;
        std::queue<int> q; q.push((int)i); used[i]=1;
        while(!q.empty()){
          int k = q.front(); q.pop(); ++csize;
          int c = cols[k], r = rows[k];
          if(eight_connected){
            for(int s=0;s<8;++s){
              long long kk = key_cr(c + dc8[s], r + dr8[s]);
              auto it = idx.find(kk);
              if(it != idx.end()){
                int j = it->second;
                if(!used[j]){ used[j]=1; q.push(j); }
              }
            }
          } else {
            for(int s=0;s<4;++s){
              long long kk = key_cr(c + dc4[s], r + dr4[s]);
              auto it = idx.find(kk);
              if(it != idx.end()){
                int j = it->second;
                if(!used[j]){ used[j]=1; q.push(j); }
              }
            }
          }
        }
        h_cluster->Fill(csize);
      }
    }

    // --- B) Event energy from PixelCharge (sum electrons -> keV) ---
    if(bchg && ievt < nEvents_chg){
      tchg->GetEntry(ievt);
      double Ne_sum = 0.0;
      if(chgs){
        for(const auto* pc : *chgs){
          // PixelCharge::getCharge() returns number of electrons collected in that pixel
          Ne_sum += pc->getCharge();
        }
      }
      const double E_keV = Ne_sum * 3.6e-3; // 3.6 eV per e-  -> keV
      h_Eevt->Fill(E_keV);
    }
  }

  // 5) Draw results
  TCanvas* c1 = new TCanvas("c_cluster","Cluster size",800,600);
  h_cluster->SetLineWidth(2);
  h_cluster->Draw();

  TCanvas* c2 = new TCanvas("c_energy","Event collected energy",800,600);
  h_Eevt->SetLineWidth(2);
  h_Eevt->Draw();

  std::cout << "[OK ] Cluster mean: " << h_cluster->GetMean() << "\n";
  std::cout << "[OK ] Event energy mean (keV): " << h_Eevt->GetMean() << "\n";
}