#include "DetectorConstruction.h"
#include "RunAction.h"

#include "G4AccumulableManager.hh"
#include "G4Run.hh"
#include "G4RunManager.hh"
#include "G4SystemOfUnits.hh"

RunAction::RunAction() {
    G4RunManager::GetRunManager()->SetPrintProgress(50000);

    const DetectorConstruction* detConstruction =
        static_cast<const DetectorConstruction*>(G4RunManager::GetRunManager()->GetUserDetectorConstruction());
    sensitiveList = detConstruction->GetSensitiveList();

    G4AccumulableManager* accumulableManager = G4AccumulableManager::Instance();
    for (int i = 0; i < sensitiveList.size(); i++) {
        G4String name = sensitiveList[i];
        accumulableManager->CreateAccumulable<G4double>(name + "Edep", 0);
        accumulableManager->CreateAccumulable<G4double>(name + "Edep2", 0);
    }
}

RunAction::~RunAction() {}

void RunAction::BeginOfRunAction(const G4Run*) {
    G4RunManager::GetRunManager()->SetRandomNumberStore(false);

    G4AccumulableManager* accumulableManager = G4AccumulableManager::Instance();
    accumulableManager->Reset();
}

void RunAction::EndOfRunAction(const G4Run* run) {
    G4int nofEvents = run->GetNumberOfEvent();
    if (nofEvents == 0) return;

    G4AccumulableManager* accumulableManager = G4AccumulableManager::Instance();
    accumulableManager->Merge();

    std::map<G4String, G4double> edep;
    std::map<G4String, G4double> edep2;
    std::map<G4String, G4double> rms;

    for (int i = 0; i < sensitiveList.size(); i++) {
        G4String name = sensitiveList[i];
        G4Accumulable<G4double>* fEdep = accumulableManager->GetAccumulable<G4double>(name + "Edep");
        G4Accumulable<G4double>* fEdep2 = accumulableManager->GetAccumulable<G4double>(name + "Edep2");
        edep[name] = fEdep->GetValue();
        edep2[name] = fEdep2->GetValue();
        rms[name] = edep2[name] - edep[name] * edep[name] / nofEvents;
        if (rms[name] > 0.)
            rms[name] = std::sqrt(rms[name]);
        else
            rms[name] = 0.;
    }

    if (IsMaster()) {
        DetectorConstruction* detConstruction =
            (DetectorConstruction*)(G4RunManager::GetRunManager()->GetUserDetectorConstruction());
        std::ofstream edepFile("result/edep_" + detConstruction->GetParticleType() + ".txt");

        for (int i = 0; i < sensitiveList.size(); i++) {
            G4String name = sensitiveList[i];
            edepFile << "bodyPart: " << std::setw(30) << name << " | eDep: " << std::setw(15) << edep[name] / MeV
                     << " MeV | eDep2: " << std::setw(15) << edep2[name] / MeV << " MeV | rms: " << std::setw(15)
                     << rms[name] / MeV << " MeV" << G4endl;
        }

        G4cout << G4endl << "--------------------End of Global Run-----------------------" << G4endl;
    } else
        G4cout << G4endl << "--------------------End of Local Run------------------------" << G4endl;
}

void RunAction::AddEdep(std::map<G4String, G4double> edep) {
    G4AccumulableManager* accumulableManager = G4AccumulableManager::Instance();
    for (int i = 0; i < sensitiveList.size(); i++) {
        G4String name = sensitiveList[i];
        G4double value = edep[name] * MeV;
        G4Accumulable<G4double>* fEdep = accumulableManager->GetAccumulable<G4double>(name + "Edep");
        G4Accumulable<G4double>* fEdep2 = accumulableManager->GetAccumulable<G4double>(name + "Edep2");
        *(fEdep) += value;
        *(fEdep2) += value * value;
    }
}
