#include "IResampler.hh"
#include "cometEventLoop.hxx"
#include "IRooTrackerFile.hxx"

using namespace COMET;

class IBackwardMCLoop: public ICOMETEventLoopFunction {
private:
    IResampler* fResampler;
    std::string fGeometryMacro;
    std::string fFieldmapDirectory;
    std::string fOutputName;
    IHandle<IRooTrackerFile> fOutputFile;
    int fNResamplings;
    bool fDetailedTransport;
    bool fDebugOutput;

public:
    IBackwardMCLoop(): 
        fResampler(nullptr), fOutputName("output.rootracker"), fNResamplings(1000),
        fDetailedTransport(true), fDebugOutput(false) {

    }

    ~IBackwardMCLoop() {
        delete fResampler;
        fOutputFile->Close();

        /// This output description should only appear if the run succeeded.
        std::cout << 
            "\n\nThe estimated flux for each primary particle has been written\n"
            "to the output file as a RooTracker entry. For each event, the flux\n"
            "is given in units of 1 / (s sr mm^2 MeV) as the 'EvtProb' branch.\n\n"
            "This number varies between resamplings, so 'EvtProb' represents\n"
            "the average flux, while the 'EvtWght' branch contains the inverse\n"
            "variance of the flux over backward MC samplings.\n\n"
            "To achieve better accuracy, you can increase the number of resamplings\n"
            "to perform per event with the '-O n-resamplings=[N]' option.\n\n"
            "In order to calculate the rate of events from the value of 'EvtProb',\n"
            "please use the formula derived on this wiki page:\n"
            "https://gitlab.in2p3.fr/comet/ICEDUST_packages/-/wikis/Backward-Monte-Carlo-sampling-with-SimBackwardMC#rate-calculation\n"
            << std::endl;
    }

    void Initialize() {
        if (fGeometryMacro == "") {
            COMETError("No geometry specified. Please use the \"-O geometry-macro\" "
                    "command-line argument.");
            throw std::runtime_error("No geometry specified");
        }
        fOutputFile = IRooTrackerFile::Open(fOutputName, "recreate");
        fResampler = new IResampler(fGeometryMacro, fFieldmapDirectory,
                                    fNResamplings, fDetailedTransport, fDebugOutput);
    }
   
    bool operator() (ICOMETEvent& event) {
        fResampler->ComputeFlux(event, fOutputFile);
        return true;
    }

    void Usage() {
        std::cout << "Custom options: \n"
            "-O geometry-macro=[name.macro]   Filename of the master geometry macro to load\n"
            "-O fieldmap-directory=[path]   Path of the fieldmap to load "
                    "(default=$(ICEDUST_FIELDMAPS))\n"
            "-O output-file=[file]          Name of output RooTracker file "
                    "(default=output.rootracker)\n"
            "-O n-resamplings=[N]           Number of backward MC resamplings to perform "
                    "to estimate the flux of each primary particle (default=1000)\n"
            "-O detailed-transport=true/false   Switches between PUMAS's detailed and hybrid "
            "transport mode. When disabled, the simulation will use hybrid energy loss and "
            "longitudinal-only scattering (default=true)\n"
            "-O debug-output                Write backward MC steps to IResampler_debug.csv "
                    "(not recommended for large n-resamplings)\n"
            << std::endl;
        
        std::cout << "Description: \n"
            "Re-sample SimG4 muon events using PUMAS (https://niess.github.io/pumas-pages/).\n"
            "Muons are backward sampled up to an initial model of the atmospheric muons\n"
            "flux. The initial model used is a tabulation of CORSIKA simulation results\n"
            "sampled at an altitude of 1600 m a.s.l. using a Polygonato primary composition.\n"
            "In addition the initial flux model has been normalized to the BESS-TeV\n"
            "measurements (Haino et al., Phys. Lett. B, 594, 2004).\n"
            << std::endl;
    }
    bool SetOption(std::string option, std::string value="") {
        if (option == "debug-output") {
            fDebugOutput = true;
            return true;
        }
        else if (option == "geometry-macro") {
            fGeometryMacro = value;
            return true;
        }
        else if (option == "fieldmap-directory") {
            fFieldmapDirectory = value;
            return true;
        }
        else if (option == "n-resamplings") {
            fNResamplings = std::stoi(value);
            return true;
        }
        else if (option == "detailed-transport") {
            if (value == "true")
                fDetailedTransport = true;
            else if (value == "false")
                fDetailedTransport = false;
            else
                return false;
            return true;
        }
        else if (option == "output-file") {
            fOutputName = value;
            return true;
        }
        return false;
    }
};

int main(int argc, char** argv) {
    IBackwardMCLoop userCode;
    cometEventLoop(argc, argv, userCode);
    return 0;
}