#include "HitOutputManager.hh"

#include "TFile.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TTree.h"

#include <cstring>
#include <iomanip>

namespace {
const char* kRootOutputName = ANASENG4_DATA_DIR "/hits.root";
const char* kTextOutputName = ANASENG4_DATA_DIR "/hits.txt";
}

HitOutputManager& HitOutputManager::Instance()
{
  static HitOutputManager instance;
  return instance;
}

HitOutputManager::HitOutputManager()
: fRootFile(nullptr),
  fTree(nullptr),
  fEventId(0),
  fDetectorId(0),
  fAnodeId(-1),
  fCathodeId(-1),
  fKineticEnergy(0.0),
  fEnergyDeposit(0.0),
  fWireDeltaESinTheta(0.0),
  fAnodeEnergy(0.0),
  fCathodeEnergy(0.0),
  fWireDeltaE(0.0),
  fX(0.0),
  fY(0.0),
  fZ(0.0),
  fVertexX(0.0),
  fVertexY(0.0),
  fVertexZ(0.0),
  fKineticEnergyHist(nullptr),
  fEnergyDepositHist(nullptr),
  fWireDeltaEHist(nullptr),
  fWireDeltaESinThetaVsDetEHist(nullptr),
  fHitZHist(nullptr),
  fVertexZHist(nullptr),
  fDetectorIdHist(nullptr),
  fHitXYHist(nullptr),
  fVertexXYHist(nullptr),
  fIncludeElectrons(true)
{
  ResetBuffers();
}

HitOutputManager::~HitOutputManager()
{
  Close();
}

void HitOutputManager::Open()
{
  Close();

  fRootFile = TFile::Open(kRootOutputName, "RECREATE");
  fTree = new TTree("hits", "Sensitive detector hits");
  fTree->Branch("eventID", &fEventId, "eventID/I");
  fTree->Branch("particleType", fParticleType, "particleType/C");
  fTree->Branch("kineticEnergy", &fKineticEnergy, "kineticEnergy/D");
  fTree->Branch("energyDeposit", &fEnergyDeposit, "energyDeposit/D");
  fTree->Branch("wireDeltaESinTheta", &fWireDeltaESinTheta, "wireDeltaESinTheta/D");
  fTree->Branch("detectorType", fDetectorType, "detectorType/C");
  fTree->Branch("detectorID", &fDetectorId, "detectorID/I");
  fTree->Branch("anodeID", &fAnodeId, "anodeID/I");
  fTree->Branch("cathodeID", &fCathodeId, "cathodeID/I");
  fTree->Branch("anodeEnergy", &fAnodeEnergy, "anodeEnergy/D");
  fTree->Branch("cathodeEnergy", &fCathodeEnergy, "cathodeEnergy/D");
  fTree->Branch("wireDeltaE", &fWireDeltaE, "wireDeltaE/D");
  fTree->Branch("volumeName", fVolumeName, "volumeName/C");
  fTree->Branch("x", &fX, "x/D");
  fTree->Branch("y", &fY, "y/D");
  fTree->Branch("z", &fZ, "z/D");
  fTree->Branch("vertexX", &fVertexX, "vertexX/D");
  fTree->Branch("vertexY", &fVertexY, "vertexY/D");
  fTree->Branch("vertexZ", &fVertexZ, "vertexZ/D");

  fKineticEnergyHist = new TH1D("hKineticEnergy", "Hit kinetic energy;Kinetic energy [MeV];Counts", 200, 0.0, 200.0);
  fEnergyDepositHist = new TH1D("hEnergyDeposit", "Step energy deposit;Energy deposit [MeV];Counts", 200, 0.0, 50.0);
  fWireDeltaEHist = new TH1D("hWireDeltaE", "Energy loss between anode and cathode crossings;#DeltaE [MeV];Counts", 200, 0.0, 20.0);
  fWireDeltaESinThetaVsDetEHist =
      new TH2D("hWireDeltaESinThetaVsDetE",
           "PC #DeltaE #times sin(#theta_{z}) vs detector deposited energy;Detector deposited energy [MeV];#DeltaE_{PC} #times sin(#theta_{z}) [MeV]",
               200,
               0.0,
               14.0,
               200,
               0.0,
               20.0);
  fHitZHist = new TH1D("hHitZ", "Hit z position;z [mm];Counts", 200, -200.0, 200.0);
  fVertexZHist = new TH1D("hVertexZ", "Track vertex z position;z_{vertex} [mm];Counts", 200, -250.0, 250.0);
  fDetectorIdHist = new TH1D("hDetectorID", "Detector ID occupancy;Detector ID;Counts", 64, -0.5, 63.5);
  fHitXYHist = new TH2D("hHitXY", "Hit position;x [mm];y [mm]", 200, -200.0, 200.0, 200, -200.0, 200.0);
  fVertexXYHist = new TH2D("hVertexXY", "Track vertex position;x_{vertex} [mm];y_{vertex} [mm]", 200, -200.0, 200.0, 200, -200.0, 200.0);

  fTextFile.open(kTextOutputName, std::ios::out | std::ios::trunc);
  fTextFile << "eventID,particleType,kineticEnergy,energyDeposit,wireDeltaESinTheta,detectorType,detectorID,anodeID,cathodeID,anodeEnergy,cathodeEnergy,wireDeltaE,volumeName,x,y,z,vertexX,vertexY,vertexZ\n";
}

void HitOutputManager::Close()
{
  if (fRootFile != nullptr) {
    fRootFile->cd();
    if (fTree != nullptr) {
      fTree->Write("", TObject::kOverwrite);
    }
    if (fKineticEnergyHist != nullptr) {
      fKineticEnergyHist->Write("", TObject::kOverwrite);
    }
    if (fEnergyDepositHist != nullptr) {
      fEnergyDepositHist->Write("", TObject::kOverwrite);
    }
    if (fWireDeltaEHist != nullptr) {
      fWireDeltaEHist->Write("", TObject::kOverwrite);
    }
    if (fWireDeltaESinThetaVsDetEHist != nullptr) {
      fWireDeltaESinThetaVsDetEHist->Write("", TObject::kOverwrite);
    }
    if (fHitZHist != nullptr) {
      fHitZHist->Write("", TObject::kOverwrite);
    }
    if (fVertexZHist != nullptr) {
      fVertexZHist->Write("", TObject::kOverwrite);
    }
    if (fDetectorIdHist != nullptr) {
      fDetectorIdHist->Write("", TObject::kOverwrite);
    }
    if (fHitXYHist != nullptr) {
      fHitXYHist->Write("", TObject::kOverwrite);
    }
    if (fVertexXYHist != nullptr) {
      fVertexXYHist->Write("", TObject::kOverwrite);
    }
    fRootFile->Close();
    delete fRootFile;
    fRootFile = nullptr;
    fTree = nullptr;
    fKineticEnergyHist = nullptr;
    fEnergyDepositHist = nullptr;
    fWireDeltaEHist = nullptr;
    fWireDeltaESinThetaVsDetEHist = nullptr;
    fHitZHist = nullptr;
    fVertexZHist = nullptr;
    fDetectorIdHist = nullptr;
    fHitXYHist = nullptr;
    fVertexXYHist = nullptr;
  }

  if (fTextFile.is_open()) {
    fTextFile.close();
  }
}

void HitOutputManager::SetIncludeElectrons(bool includeElectrons)
{
  fIncludeElectrons = includeElectrons;
}

bool HitOutputManager::GetIncludeElectrons() const
{
  return fIncludeElectrons;
}

bool HitOutputManager::ShouldRecordParticle(const std::string& particleType) const
{
  if (fIncludeElectrons) {
    return true;
  }

  return particleType != "e-" && particleType != "e+";
}

void HitOutputManager::RecordHit(int eventId,
                                 const std::string& particleType,
                                 double kineticEnergy,
                                 double energyDeposit,
                                 double wireDeltaESinTheta,
                                 const std::string& detectorType,
                                 int detectorId,
                                 int anodeId,
                                 int cathodeId,
                                 double anodeEnergy,
                                 double cathodeEnergy,
                                 double wireDeltaE,
                                 const std::string& volumeName,
                                 const G4ThreeVector& position,
                                 const G4ThreeVector& vertexPosition)
{
  if (fRootFile == nullptr || fTree == nullptr || !fTextFile.is_open()) {
    return;
  }

  ResetBuffers();

  fEventId = eventId;
  fDetectorId = detectorId;
  fAnodeId = anodeId;
  fCathodeId = cathodeId;
  fKineticEnergy = kineticEnergy;
  fEnergyDeposit = energyDeposit;
  fWireDeltaESinTheta = wireDeltaESinTheta;
  fAnodeEnergy = anodeEnergy;
  fCathodeEnergy = cathodeEnergy;
  fWireDeltaE = wireDeltaE;
  fX = position.x();
  fY = position.y();
  fZ = position.z();
  fVertexX = vertexPosition.x();
  fVertexY = vertexPosition.y();
  fVertexZ = vertexPosition.z();

  std::strncpy(fParticleType, particleType.c_str(), sizeof(fParticleType) - 1);
  std::strncpy(fDetectorType, detectorType.c_str(), sizeof(fDetectorType) - 1);
  std::strncpy(fVolumeName, volumeName.c_str(), sizeof(fVolumeName) - 1);

  fTree->Fill();
  if (fKineticEnergyHist != nullptr) {
    fKineticEnergyHist->Fill(fKineticEnergy);
  }
  if (fEnergyDepositHist != nullptr) {
    fEnergyDepositHist->Fill(fEnergyDeposit);
  }
  if (fWireDeltaEHist != nullptr && fAnodeId >= 0 && fCathodeId >= 0) {
    fWireDeltaEHist->Fill(fWireDeltaE);
  }
  if (fWireDeltaESinThetaVsDetEHist != nullptr && fAnodeId >= 0 && fCathodeId >= 0) {
    fWireDeltaESinThetaVsDetEHist->Fill(fEnergyDeposit, 100.0 * fWireDeltaESinTheta); //2D histogram programmed
  }
  if (fHitZHist != nullptr) {
    fHitZHist->Fill(fZ);
  }
  if (fVertexZHist != nullptr) {
    fVertexZHist->Fill(fVertexZ);
  }
  if (fDetectorIdHist != nullptr) {
    fDetectorIdHist->Fill(fDetectorId);
  }
  if (fHitXYHist != nullptr) {
    fHitXYHist->Fill(fX, fY);
  }
  if (fVertexXYHist != nullptr) {
    fVertexXYHist->Fill(fVertexX, fVertexY);
  }

  fTextFile << fEventId << ','
            << fParticleType << ','
            << std::setprecision(12) << fKineticEnergy << ','
            << fEnergyDeposit << ','
            << fWireDeltaESinTheta << ','
            << fDetectorType << ','
            << fDetectorId << ','
            << fAnodeId << ','
            << fCathodeId << ','
            << fAnodeEnergy << ','
            << fCathodeEnergy << ','
            << fWireDeltaE << ','
            << fVolumeName << ','
            << fX << ','
            << fY << ','
            << fZ << ','
            << fVertexX << ','
            << fVertexY << ','
            << fVertexZ << '\n';
}

void HitOutputManager::ResetBuffers()
{
  fParticleType[0] = '\0';
  fDetectorType[0] = '\0';
  fVolumeName[0] = '\0';
}