#include <TFile.h>
#include <TTree.h>
#include <iostream>
#include <fstream>
#include <string>
#include <unordered_map>


std::string getParticleName(int pdgCode);

std::string identifyAtom(int pdgCode);

std::unordered_map<std::string, int> countParticles(TTree* tree, const char* branchName);

std::unordered_map<std::string, int> countParticlesPdgp(TTree* tree);
  
void writeCountsToCSV(const std::unordered_map<std::string, int>& counts, const std::string& fileName);

void minimal() 

{
  // Open the ROOT file
  TFile *file = new TFile("out.root");
    
  // Get the tree from the file
  TTree *tree = (TTree*)file->Get("mst");

  // ROOT file for output
  TFile *outFile = new TFile("Processed.root", "RECREATE");

  // put counts of particles to csv
  // List of branches to process
  std::vector<std::string> branches = {"pdgp"};
  // Process each branch
  for (const std::string& branch : branches) {
    std::unordered_map<std::string, int> counts;
    counts = countParticlesPdgp(tree);
    
    writeCountsToCSV(counts, branch + ".csv");
  }
  
  
  // Clean up
  
  // Close the output file
  outFile->Close();
  delete file;
}

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// ooooooooooooooooooooooooooooooooooooooooooooooooooooooooo
// OOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOO
// ooooooooooooooooooooooooooooooooooooooooooooooooooooooooo
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

std::string getParticleName(int pdgCode) {
    switch (pdgCode) {
        case 11:    return "Electron";
        case -11:   return "Positron";
        case 13:    return "Muon-";
        case -13:   return "Muon+";
        case 15:    return "Tau-";
        case -15:   return "Tau+";
        case 22:    return "Photon";
        case 12:    return "Electron neutrino";
        case -12:   return "Electron antineutrino";
        case 14:    return "Muon neutrino";
        case -14:   return "Muon antineutrino";
        case 16:    return "Tau neutrino";
        case -16:   return "Tau antineutrino";
        default:    
            if (pdgCode > 1000000000) { // Looks like an atom's PDG code
                return identifyAtom(pdgCode);
            } else {
                // For unidentified particles, print the PDG code
                std::cout << "Unidentified particle with PDG code: " << pdgCode << std::endl;
                return "Unknown particle";
            }
    }
}

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// ooooooooooooooooooooooooooooooooooooooooooooooooooooooooo
// OOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOO
// ooooooooooooooooooooooooooooooooooooooooooooooooooooooooo
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~



std::string identifyAtom(int pdgCode) {
  // Extract ZZZ and AAA from the PDG code
  int ZZZ = (pdgCode / 10000) % 1000; // Extracts the atomic number
  int AAA = (pdgCode / 10) % 1000; // Extracts the mass number

  // Map from atomic number to element symbol for all known elements
  std::unordered_map<int, std::string> elementSymbols = {
    {1, "H"}, {2, "He"}, {3, "Li"}, {4, "Be"}, {5, "B"},
    {6, "C"}, {7, "N"}, {8, "O"}, {9, "F"}, {10, "Ne"},
    {11, "Na"}, {12, "Mg"}, {13, "Al"}, {14, "Si"}, {15, "P"},
    {16, "S"}, {17, "Cl"}, {18, "Ar"}, {19, "K"}, {20, "Ca"},
    {21, "Sc"}, {22, "Ti"}, {23, "V"}, {24, "Cr"}, {25, "Mn"},
    {26, "Fe"}, {27, "Co"}, {28, "Ni"}, {29, "Cu"}, {30, "Zn"},
    {31, "Ga"}, {32, "Ge"}, {33, "As"}, {34, "Se"}, {35, "Br"},
    {36, "Kr"}, {37, "Rb"}, {38, "Sr"}, {39, "Y"}, {40, "Zr"},
    {41, "Nb"}, {42, "Mo"}, {43, "Tc"}, {44, "Ru"}, {45, "Rh"},
    {46, "Pd"}, {47, "Ag"}, {48, "Cd"}, {49, "In"}, {50, "Sn"},
    {51, "Sb"}, {52, "Te"}, {53, "I"}, {54, "Xe"}, {55, "Cs"},
    {56, "Ba"}, {57, "La"}, {58, "Ce"}, {59, "Pr"}, {60, "Nd"},
    {61, "Pm"}, {62, "Sm"}, {63, "Eu"}, {64, "Gd"}, {65, "Tb"},
    {66, "Dy"}, {67, "Ho"}, {68, "Er"}, {69, "Tm"}, {70, "Yb"},
    {71, "Lu"}, {72, "Hf"}, {73, "Ta"}, {74, "W"}, {75, "Re"},
    {76, "Os"}, {77, "Ir"}, {78, "Pt"}, {79, "Au"}, {80, "Hg"},
    {81, "Tl"}, {82, "Pb"}, {83, "Bi"}, {84, "Po"}, {85, "At"},
    {86, "Rn"}, {87, "Fr"}, {88, "Ra"}, {89, "Ac"}, {90, "Th"},
    {91, "Pa"}, {92, "U"}, {93, "Np"}, {94, "Pu"}, {95, "Am"},
    {96, "Cm"}, {97, "Bk"}, {98, "Cf"}, {99, "Es"}, {100, "Fm"},
    {101, "Md"}, {102, "No"}, {103, "Lr"}, {104, "Rf"}, {105, "Db"},
    {106, "Sg"}, {107, "Bh"}, {108, "Hs"}, {109, "Mt"}, {110, "Ds"},
    {111, "Rg"}, {112, "Cn"}, {113, "Nh"}, {114, "Fl"}, {115, "Mc"},
    {116, "Lv"}, {117, "Ts"}, {118, "Og"}
  };

  // Check if the element symbol exists for the atomic number
  if (elementSymbols.find(ZZZ) == elementSymbols.end()) {
    return "Unknown element";
  }

  // Construct and return the result as "Symbol-A"
  return elementSymbols[ZZZ] + "-" + std::to_string(AAA);
}

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// ooooooooooooooooooooooooooooooooooooooooooooooooooooooooo
// OOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOO
// ooooooooooooooooooooooooooooooooooooooooooooooooooooooooo
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~


// Function to read PDG IDs from pdgp branch and count particle occurrences
std::unordered_map<std::string, int> countParticlesPdgp(TTree* tree) {
  std::cout << "Counting particle IDs for pdgp.\n";
  std::unordered_map<std::string, int> particleCounts;

  // Variables to hold branch data
  Int_t* pdgp = nullptr; // Pointer to the array of PDG codes
  Int_t np = 0; // Number of PDG codes in pdgp for each entry

  // Automatically manage the branch-to-variable linkage
  tree->SetBranchAddress("np", &np);
  tree->SetBranchAddress("pdgp", &pdgp);

  Long64_t nEntries = tree->GetEntries();
  for (Long64_t i = 0; i < nEntries; ++i) {
    std::cout << "Before GetEntry, np: " << np << ", pdgp pointer is null? " << (pdgp==nullptr) << std::endl;
    tree->GetEntry(i);

    std::cout << "The pointer to pdgp is: " << pdgp << "\n";
    
    std::cout << "After GetEntry, np: " << np << ", pdgp pointer is null? " << (pdgp==nullptr) << std::endl;

    // Now iterate over the pdgp array based on np
    for (int j = 0; j < np; ++j) {
      std::cout << "1\n";
      if(np>0) std::cout << np << ", " << pdgp[0] << ".\n";
      std::string particleName = getParticleName(pdgp[j]);
      std::cout << "2\n";
      particleCounts[particleName]++;
    }

    if (i % 10000 == 0) {
      std::cout << "Processed " << i << " of " << nEntries << " entries.\n";
    }
  }

  return particleCounts;
}


// Function to write counts to a CSV file
void writeCountsToCSV(const std::unordered_map<std::string, int>& counts, const std::string& fileName) {
  std::ofstream file(fileName);
  for (const auto& pair : counts) {
    file << pair.first << "," << pair.second << "\n";
  }
}