/*
 * Copyright (C) 2018 Université Clermont Auvergne, CNRS/IN2P3, LPC
 * Author: Valentin NIESS (niess@in2p3.fr)
 *
 * A Geant4 particle source of atmospheric muons using GOUPIL
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>
 */

#include "IGoupilSource.hh"
#include "IGoupilSettings.hh"
#include "IGoupilUserAction.hh"
#include "goupil.h"
#include "goupil-eole.h"
#include "goupil-gull.h"
#include "goupil-primal.h"
#include "goupil-primate.h"
#include "goupil-turtle.h"

/* Geant4 includes */
#include "G4Event.hh"
#include "G4ExceptionHandler.hh"
#include "G4Field.hh"
#include "G4FieldManager.hh"
#include "G4Material.hh"
#include "G4MuonMinus.hh"
#include "G4MuonPlus.hh"
#include "G4TransportationManager.hh"
#include "Randomize.hh"
#include "G4PVPlacement.hh"

/* STD lib includes */
#include <iostream>
#include <fstream>

/* Global data and settings */
static struct {
        /* Reference count of the number of sources instances */
        int references;
        /* The Geant4 geometry layer */
        struct goupil_geometry_layer * g4layer;
        /* The Geant4 frame transform */
        struct goupil_coordinates_frame g4frame;
} gGlobals = {0, NULL, {{0, 0, 0, 0, 0, 0, 0, 0, 0}, {0, 0, 0}}};

/* Error handler for GOUPIL */
static void handle_goupil(const char * message)
{
        G4ExceptionDescription description;
        description << G4endl << message << G4endl;
        G4Exception("IGoupilSource", "A GOUPIL library error occurred",
            FatalException, description);
}

/* Encapsulation of the Geant4 random engine */
static double random_uniform01(struct goupil_random *)
{
        return G4UniformRand();
}

static int random_initialise_sampler(struct goupil_sampler_module * module)
{
        struct goupil_random * random = (struct goupil_random *)module;
        random->uniform01 = &random_uniform01;
        return EXIT_SUCCESS;
}

static const char * goupil_signature_random(void)
{
        return "G4UniformRand";
}

static int random_register(void)
{
        return goupil_random_register(
            &goupil_signature_random, NULL, &random_initialise_sampler);
}

/* Map for translating a Geant4 volume to a GOUPIL medium. TODO: optimize
 * this, if relevant */
typedef std::map<G4VPhysicalVolume *, struct goupil_medium *> VolumeProxy;

/* History data of the geometry stepping */
struct g4_volume_history {
        struct goupil_medium * medium[2];
        VolumeProxy proxy;
};

/* Geant4 specific data for GOUPIL media */
struct medium_data {
        G4VPhysicalVolume * physical;
        const G4Field * field;
};

/* Update the map of Geant4 to GOUPIL volumes & the stepping history */
static struct goupil_medium * UpdateVolumes(
    struct g4_volume_history * history, G4VPhysicalVolume * physical)
{
        /* Check the history */
        if ((G4VPhysicalVolume *)history->medium[0] != NULL) {
                struct medium_data * data = (struct medium_data *)
                    goupil_medium_data(history->medium[0]);
                if (data->physical == physical)
                        return history->medium[0];
        }
        if ((G4VPhysicalVolume *)history->medium[1] != NULL) {
                struct medium_data * data = (struct medium_data *)
                    goupil_medium_data(history->medium[1]);
                if (data->physical == physical) {
                        struct goupil_medium * m = history->medium[1];
                        history->medium[1] = history->medium[0];
                        history->medium[0] = m;
                        return m;
                }
        }

        /* Look down the tree for the physical volume */
        struct goupil_medium *& medium = history->proxy[physical];
        if (medium == NULL) {
                /* Fetch the medium data */
                const char * name = physical->GetName().c_str();
                G4LogicalVolume * logical = physical->GetLogicalVolume();
                G4Material * material = logical->GetMaterial();
                const char * material_name = material->GetName().c_str();
                double density = material->GetDensity() / CLHEP::kg * CLHEP::m3;

                /* Allocate the medium */
                goupil_medium_create(&medium, name, material_name, density,
                    sizeof (struct medium_data));

                /* Link the Geant4 physical volume */
                struct medium_data * data = (struct medium_data *)
                    goupil_medium_data(medium);
                data->physical = physical;

                /* Link any electromagnetic field */
                G4FieldManager * mgr = logical->GetFieldManager();
                if (mgr == NULL) {
                        mgr = G4TransportationManager::
                            GetTransportationManager()->GetFieldManager();
                }
                if (mgr != NULL)
                        data->field = mgr->GetDetectorField();
                else
                        data->field = NULL;
        }

        /* Update the history */
        history->medium[1] = history->medium[0];
        history->medium[0] = medium;

        return medium;
}

/* Encapsulation of the Geant4 geometry navigator, for GOUPIL */
static double geometry_get_medium(struct goupil_geometry * geometry,
    const double position[3], const double direction[3],
    struct goupil_medium ** medium)
{
        G4ThreeVector r(position[0] * CLHEP::m, position[1] * CLHEP::m,
            position[2] * CLHEP::m);
        G4ThreeVector u(direction[0], direction[1], direction[2]);
        G4Navigator * navigator =
            G4TransportationManager::GetTransportationManager()
                ->GetNavigatorForTracking();

        /* Fetch the physical volume */
        if (medium != NULL) {
                navigator->ResetStackAndState();
                G4VPhysicalVolume * physical =
                    navigator->LocateGlobalPointAndSetup(r, &u);
                if (physical == NULL) {
                        *medium = NULL;
                        return 0.;
                }

                /* Fetch or register the volume if new */
                *medium = UpdateVolumes((struct g4_volume_history *)
                    geometry->module.data, physical);
        }

        /* Compute the step length */
        G4double safety = 0.;
        return navigator->ComputeStep(r, u, kInfinity, safety) / CLHEP::m;
}

/* Encapsulation of the Geant4 magnetic field, for GOUPIL */
static double geometry_get_magnet(struct goupil_geometry *,
    struct goupil_medium * medium, const double * position,
    const double * direction, double * magnet)
{
        /* Unpack the Geant4 field */
        struct medium_data * data = (struct medium_data *)
            goupil_medium_data(medium);
        if ((data == NULL) || (data->field == NULL)) return -1.;

        /* Use the Geant4 field getter in order to compute the magnetic field */
        double point[4] = { position[0] * CLHEP::m,
            position[1] * CLHEP::m, position[2] * CLHEP::m, 0. };
        double value0[6], value1[6];
        data->field->GetFieldValue(point, value0);

        /* Fill in the magnetic field components and update the position */
        const double ds = 1. * CLHEP::mm;
        int i;
        for (i = 0; i < 3; i++) {
                magnet[i] = value0[i] / CLHEP::tesla;
                point[i] += ds * direction[i];
        }

        /* Compute the magnetic field gradient */
        data->field->GetFieldValue(point, value1);

        double dB = 0., B = 0.;
        for (i = 0; i < 3; i++) {
                dB += fabs(value0[i] - value1[i]);
                const double tmp = 0.5 * (value0[i] + value1[i]);
                B += tmp * tmp;
        }
        if ((B > 0.) && (dB > 0.)) {
                B = sqrt(B);
                return 1E-02 * ds * B / (dB * CLHEP::m);
        } else {
                /* TODO: Geometry size? */
                return 1.;
        }
}

static const char * goupil_signature_geometry(void)
{
        return "geometry";
}

/* Clean the sampler */
static void geometry_clean(struct goupil_sampler_module * module)
{
        delete (struct g4_volume_history *)module->data;
        module->data = NULL;
}

/* Initialise the geometry data for a new sampler, i.e. a thread local map
 * of G4PhysicalVolume's and stepping history.
 */
static int geometry_initialise_sampler(struct goupil_sampler_module * module)
{
        module->clean = &geometry_clean;
        struct g4_volume_history * history = new struct g4_volume_history;
        history->medium[1] = history->medium[0] = NULL;
        module->data = history;

        struct goupil_geometry * geometry = (struct goupil_geometry *)module;
        geometry->get_medium = &geometry_get_medium;
        geometry->get_magnet = &geometry_get_magnet;

        return EXIT_SUCCESS;
}

static int geometry_register(void)
{
        return goupil_geometry_register(&goupil_signature_geometry, NULL);
}

static G4double to_mm(G4double value)
{
        return round(value * CLHEP::m / CLHEP::mm * 1E+03) * 1E-03;
}

static G4double to_MeV(G4double value)
{
        return round(value * CLHEP::GeV / CLHEP::MeV * 1E+03) * 1E-03;
}

/* Map for translating GOUPIL outer volumes to Geant4 ones */
typedef std::map<struct goupil_medium *, G4VPhysicalVolume *> OuterVolume;

/* Stepping history for recording */
struct SteppingHistory {
        int verbose;
        std::ostream * stream;
        IGoupilState * ancester;

        G4int index;
        struct goupil_geometry * geometry;
        G4double kinetic_energy;
        G4double distance;

        OuterVolume outer;
};

/* Map a GOUPIL event to a name string */
static const char * event_to_str(enum goupil_event event)
{
        if (event & GOUPIL_EVENT_START)
                return "Initialisation";
        else if (event & GOUPIL_EVENT_LIMIT_KINETIC)
                return "KineticLimit";
        else if (event & GOUPIL_EVENT_MEDIUM)
                return "Transportation";
        else if (event & GOUPIL_EVENT_VERTEX_BREMSSTRAHLUNG)
                return "Bremsstrahlung";
        else if (event & GOUPIL_EVENT_VERTEX_PAIR_CREATION)
                return "PairCreation";
        else if (event & GOUPIL_EVENT_VERTEX_PHOTONUCLEAR)
                return "Photonuclear";
        else if (event & GOUPIL_EVENT_VERTEX_DELTA_RAY)
                return "DeltaRay";
        else if (event & GOUPIL_EVENT_VERTEX_COULOMB)
                return "CoulombScattering";
        else
                return "Ionisation";
}

static void recorder_dump(struct goupil_recorder * recorder,
    struct goupil_recorder_frame * frame)
{
        SteppingHistory * history = (SteppingHistory *)recorder->module.data;
        if (frame->event & GOUPIL_EVENT_START) {
                *history->stream                                    << G4endl;
                *history->stream << std::string(105, '*')           << G4endl;
                *history->stream << "* IGoupilSource Information:" << G4endl;
                *history->stream << std::string(105, '*')           << G4endl;
        }

        IGoupilSettings * settings = IGoupilSettings::GetInstance();
        if (settings->GetVerboseLevel() < 2)
                return;

        if (frame->geometry != history->geometry) {
                /* Dump a header. TODO: skip if same than previous */
                const goupil_geometry_layer * layer =
                    goupil_geometry_get_layer(frame->geometry);
                const char * tag[4];
                if (layer == gGlobals.g4layer) {
                        tag[0] = "X(mm)";
                        tag[1] = "Y(mm)";
                        tag[2] = "Z(mm)";
                        tag[3] = "Volume";
                } else {
                        tag[0] = "Lat(deg)";
                        tag[1] = "Lon(deg)";
                        tag[2] = "H(m)";
                        tag[3] = "Medium";
                }

                *history->stream << G4endl
                       << std::setw( 5) << "Step#"      << " "
                       << std::setw( 8) << tag[0]       << " "
                       << std::setw( 8) << tag[1]       << " "
                       << std::setw( 8) << tag[2]       << " "
                       << std::setw( 9) << "KinE(MeV)"  << " "
                       << std::setw( 8) << "dE(MeV)"    << " "
                       << std::setw( 8) << "StepLeng"   << " "
                       << std::setw( 9) << "TrackLeng"  << " "
                       << std::setw(11) << tag[3]       << " "
                       << std::setw( 8) << "ProcName"   << G4endl;
        }

        /* Dump the step summary */
        char line[120], * cursor = line;
        cursor += sprintf(cursor, "%5d ", history->index);

        if (frame->geometry != NULL) {
                const goupil_coordinates_frame * f =
                goupil_geometry_get_frame(frame->geometry);
                goupil_coordinates_cartesian * r = frame->position;
                goupil_sampler_coordinates_transform(
                    recorder->module.sampler, f, r, r);

                cursor += sprintf(cursor, "%8g %8g %8g ", to_mm(r->x),
                                to_mm(r->y), to_mm(r->z));
        } else {
                goupil_coordinates_geographic * r =
                    (goupil_coordinates_geographic *)frame->position;
                goupil_sampler_coordinates_as(recorder->module.sampler,
                    GOUPIL_COORDINATES_GEOGRAPHIC, r, r);

                const double dlat =  r->latitude -
                    settings->GetEarthLocationLatitude() / CLHEP::deg;
                const double dlon =  r->longitude -
                    settings->GetEarthLocationLongitude() / CLHEP::deg;
                cursor += sprintf(cursor, "%8.1e %8.1e %8.3f ", dlat, dlon,
                                  r->altitude);
        }

        G4double dE = to_MeV(frame->kinetic_energy - history->kinetic_energy); 
        G4double ds = to_mm(frame->distance - history->distance);
        const char * medium_name;
        if (frame->medium == NULL)
                medium_name = "OutOfWorld";
        else
                medium_name = goupil_medium_get_name(frame->medium);

        cursor += sprintf(cursor, "%9g %8g %8g %9g %11s %8s",
                          to_MeV(frame->kinetic_energy), dE, ds,
                          to_mm(frame->distance), medium_name,
                          event_to_str(frame->event));
        *history->stream << line << G4endl;
}


static void recorder_copy_state(struct goupil_recorder * recorder,
    struct goupil_recorder_frame * frame, IGoupilState * state,
    bool geographic)
{
        goupil_coordinates_cartesian * r = frame->position;
        if (geographic) {
                GOUPIL_COORDINATES_GEOGRAPHIC_CREATE(rgeo, 0, 0, 0);
                goupil_sampler_coordinates_as(recorder->module.sampler,
                    GOUPIL_COORDINATES_GEOGRAPHIC, r, rgeo);
                const double dz = rgeo->altitude;
                goupil_sampler_coordinates_as(recorder->module.sampler,
                    GOUPIL_COORDINATES_GEODETIC, rgeo, rgeo);

                state->SetLatitude(rgeo->latitude * CLHEP::deg);
                state->SetLongitude(rgeo->longitude * CLHEP::deg);
                state->SetAltitude(rgeo->altitude * CLHEP::m);
                state->SetTopographyLevel((rgeo->altitude - dz) * CLHEP::m);
        }

        goupil_sampler_coordinates_transform(
            recorder->module.sampler, &gGlobals.g4frame, r, r);
        state->SetPosition(
            G4ThreeVector(r->x, r->y, r->z) * CLHEP::m);

        goupil_coordinates_cartesian * u = frame->direction;
        goupil_sampler_coordinates_transform(
            recorder->module.sampler, &gGlobals.g4frame, u, u);
        state->SetMomentumDirection(
            G4ThreeVector(u->x, u->y, u->z));

        state->SetKineticEnergy(
            frame->kinetic_energy * CLHEP::GeV);
}


static void recorder_record(struct goupil_recorder * recorder,
    struct goupil_recorder_frame * frame)
{
        SteppingHistory * history = (SteppingHistory *)recorder->module.data;
        if (frame->event < 0) {
                recorder_copy_state(recorder, frame, history->ancester, false);

                /* TODO: print a summary */
                return;
        }

        IGoupilUserAction * userAction =
            (IGoupilUserAction *)goupil_sampler_data_get(recorder->module.sampler);
        if ((!history->verbose) && (!userAction)) {
                return;
        } else if (frame->event & GOUPIL_EVENT_START) {
                /* Initialise the history */
                history->index = 0;
                history->geometry = NULL;
                history->geometry++; /* Initialise to an unreachable address */
                history->kinetic_energy = frame->kinetic_energy;
                history->distance = frame->distance;
        }

        if (history->verbose) {
                recorder_dump(recorder, frame);
        }

        if (userAction) {
                IGoupilState state;
                recorder_copy_state(recorder, frame, &state, true);

                G4VPhysicalVolume * physical;
                double field[6] = { 0., 0., 0., 0., 0., 0. };
                if (frame->medium) {
                        if (frame->geometry != NULL) {
                                struct medium_data * data =
                                    (struct medium_data *)goupil_medium_data(frame->medium);
                                physical = data->physical;

                                const G4ThreeVector & position = state.GetPosition();
                                const double r[3] = { position.x() / CLHEP::m,
                                    position.y() / CLHEP::m, position.z() / CLHEP::m };
                                const G4ThreeVector & direction = state.GetMomentumDirection();
                                const double u[3] = { direction.x(), direction.y(), direction.z() };
                                geometry_get_magnet(
                                    frame->geometry, frame->medium, r, u, field + 3);
                        } else {
                                physical = history->outer[frame->medium];
                                if (physical == NULL) {
                                        /* Create a volume wrapper */
                                        static G4LogicalVolume logical(0, 0, "(nil)");
                                        physical = new G4PVPlacement(0, G4ThreeVector(0, 0, 0), &logical,
                                            goupil_medium_get_name(frame->medium), 0, false, 0);
                                        history->outer[frame->medium] = physical;
                                }
                        }
                } else {
                        physical = NULL;
                }

                userAction->SteppingAction(history->index, frame->event,
                    physical, &state, field);
        }

        /* Update the history */
        history->index++;
        history->geometry = frame->geometry;
        history->kinetic_energy = frame->kinetic_energy;
        history->distance = frame->distance;
}

static void clean_recorder(struct goupil_sampler_module * module)
{
        SteppingHistory * history = (SteppingHistory *)module->data;
        if (history->stream != &G4cout) {
                std::ofstream * ofs =
                    static_cast<std::ofstream *>(history->stream);
                ofs->close();
                delete history->stream;
        }
        delete (SteppingHistory *)history;
}

static int recorder_initialise_sampler(struct goupil_sampler_module * module)
{
        SteppingHistory * history = new SteppingHistory;
        history->verbose = IGoupilSettings::GetVerboseLevel();
        const G4String& file = IGoupilSettings::GetVerboseFile();
        /* TODO: handle multiple threads, e.g. rename the file */
        if (file.size()) {
                std::ofstream * ofs = new std::ofstream;
                ofs->open(file);
                history->stream = ofs;
        } else {
                history->stream = &G4cout;
        }
        module->data = history;
        module->clean = &clean_recorder;

        struct goupil_recorder * recorder = (struct goupil_recorder *)module;
        recorder->record = &recorder_record;

        return EXIT_SUCCESS;
}

static const char * goupil_signature_recorder(void)
{
        return "recorder";
}

static int recorder_register(void)
{
        return goupil_recorder_register(&goupil_signature_recorder, NULL,
            &recorder_initialise_sampler);
}

static G4String GetGoupilDataDir()
{
        const char * path = getenv("GOUPIL_DATA_DIR");
        if (path == NULL)
                path = GOUPIL_DATA_DIR;
        return path;
}

static void InitialiseSampler(
    struct goupil_sampler ** sampler, IGoupilUserAction * userAction, const std::string& samplerMode)
{
        if (*sampler != NULL)
                return;

        /* Check the configuration */
        IGoupilSettings * settings = IGoupilSettings::GetInstance();
        if (!settings->CheckEarthLocation()) {
                G4ExceptionDescription description;
                description << G4endl 
                            << "==> Earth location has not been set <=="
                            << G4endl;
                G4Exception("IGoupilSource", "Invalid configuration",
                    FatalException, description);
        }

        /* Initialise GOUPIL : TODO: put a barrier here */
        if (gGlobals.references == 0) {
                /* Initialise the GOUPIL library */
                G4String data_dir = GetGoupilDataDir();
                goupil_error_handler_set(&handle_goupil);
                G4String mdf_path = data_dir + "/materials/materials.xml";
                G4String dedx_path = data_dir + "/materials/dedx";
                G4String cfg_path = data_dir + "/.g4goupil.cfg";
                char * argv[] = {
                    (char *)mdf_path.c_str(),
                    (char *)dedx_path.c_str(),
                    (char *)cfg_path.c_str() };
                goupil_initialise(
                    NULL, NULL, sizeof(argv) / sizeof(*argv), argv);

                /* Configure GOUPIL to use the Geant4 pseudo-random engine */
                random_register();

                /* Configure the atmosphere model */
                goupil_eole_register();

                /* Configure the geomagnetic field */
                G4String geomagnet_path = data_dir + "/geomagnet/" +
                    settings->GetGeomagnetModel() + ".COF";
                goupil_gull_register(geomagnet_path.c_str(),
                    settings->GetGeomagnetDateDay(),
                    settings->GetGeomagnetDateMonth(),
                    settings->GetGeomagnetDateYear());

                /* Configure the topography */
                goupil_turtle_register(NULL, NULL); /* TODO: multithreading */
                goupil_turtle_set_bottom(
                    settings->GetTopographyBottom() / CLHEP::m);
                const IGoupilTopography & topography =
                    settings->GetTopography();
                const G4String & reference = settings->GetTopographyReference();
                bool referenceSet = false;
                for (IGoupilTopography::const_iterator i = topography.begin();
                    i != topography.end(); i++) {
                        struct goupil_medium * layer;
                        G4String name = i->GetName();
                        goupil_turtle_layer_create(&layer, name.c_str(),
                            i->GetMaterial().c_str(),
                            i->GetDensity() / CLHEP::kg * CLHEP::m3);
                        IGoupilTopographyData data = i->GetData();
                        for (IGoupilTopographyData::const_iterator j =
                            data.begin(); j != data.end(); j++) {
                                struct goupil_turtle_data * datum;
                                const char * path = j->first.size() ?
                                    j->first.c_str() : NULL;
                                goupil_turtle_data_create(&datum, path);
                                goupil_turtle_layer_link(
                                    layer, datum, j->second / CLHEP::m);
                        }
                        if (reference == name) {
                                goupil_turtle_set_reference(layer);
                                referenceSet = true;
                        }
                }
                if (reference.empty()) {
                                goupil_turtle_set_reference(NULL);
                } else if (!referenceSet) {
                        G4ExceptionDescription description;
                        description << G4endl << "==> Unknown reference layer `"
                                    << reference << "' <==" << G4endl;
                        G4Exception("IGoupilSource",
                            "A GOUPIL configuration error occurred",
                            FatalException, description);
                }

                /* Configure the primary flux model */
                const G4String & model = settings->GetPrimaryModel();
                if (strncmp(model.c_str(), "table:", 6) == 0) {
                        std::string filename = model.substr(6);
                        if (filename[0] != '/') {
                                G4String dirname = data_dir + "/primary/";
                                filename = dirname + filename;
                        }
                        goupil_primate_register(filename.c_str());
                        const double height = *goupil_primary_get_altitude();
                        settings->SetPrimaryHeight(height * CLHEP::m);
                } else {
                        goupil_primal_register();
                        goupil_primal_model_set(model);
                        goupil_primal_altitude_set(
                            settings->GetPrimaryHeight() / CLHEP::m);
                }

                /* Configure the recorder */
                G4int verboseLevel = settings->GetVerboseLevel();
                recorder_register();
                if ((verboseLevel >= 3) || userAction)
                       goupil_recorder_period_set(1);
                else
                       goupil_recorder_period_set(0);

                /* Configure the interface with the Geant4 geometry */
                geometry_register();
                GOUPIL_COORDINATES_GEOGRAPHIC_CREATE(origin,
                    settings->GetEarthLocationLatitude() / CLHEP::deg,
                    settings->GetEarthLocationLongitude() / CLHEP::deg,
                    settings->GetEarthLocationAltitude() / CLHEP::m);

                GOUPIL_COORDINATES_EULER_CREATE(orientation,
                    settings->GetWorldOrientationAlpha() / CLHEP::deg,
                    settings->GetWorldOrientationBeta() / CLHEP::deg,
                    settings->GetWorldOrientationGamma() / CLHEP::deg);

                goupil_geometry_layer_create(&gGlobals.g4layer, 0,
                    "Geant4", &goupil_signature_geometry, NULL,
                    &geometry_initialise_sampler, origin, orientation);

                /* Configure the sampling mode */
                goupil_sampler_mode_set(samplerMode.c_str());
        }

        /* Create the GOUPIL sampler */
        goupil_sampler_create(sampler);
        goupil_sampler_data_set(*sampler, userAction);

        /* Set some remaining globals, requiring the sampler */
        if (gGlobals.references++ == 0) {
                const struct goupil_coordinates_frame * frame =
                    goupil_sampler_get_frame(*sampler, gGlobals.g4layer);
                memcpy(&gGlobals.g4frame, frame, sizeof(gGlobals.g4frame));
        }
}

/* Instanciate a new GOUPIL generator */
IGoupilSource::IGoupilSource() : G4SingleParticleSource(), fSampler(NULL),
    fSamplerMode("DETAILED"), fUserAction(NULL)
{
        /* Instanciate the muons, in case the Physics didn't do so */
        G4MuonPlus::Definition();
        G4MuonMinus::Definition();

        /* Initialise the particle source */
        G4SingleParticleSource::SetParticleDefinition(
            G4MuonMinus::Definition());
}

/* Properly clean a GOUPIL generator */
IGoupilSource::~IGoupilSource()
{
        goupil_sampler_destroy(&fSampler);
        if (--gGlobals.references == 0) {
                goupil_geometry_layer_destroy(&gGlobals.g4layer);
                goupil_finalise();
        }
}

/* Generate a muon final state and backward sample its atmospheric flux */
void IGoupilSource::GeneratePrimaryVertex(G4Event * event)
{
        InitialiseSampler(&fSampler, fUserAction, fSamplerMode);

        IGoupilSettings * settings = IGoupilSettings::GetInstance();
        double weight;
        if (settings->GetChargeRandomisation()) {
                /* Generate the charge */
                weight = 2.;
                if (G4UniformRand() < 0.5)
                        G4SingleParticleSource::SetParticleDefinition(
                            G4MuonMinus::Definition());
                else
                        G4SingleParticleSource::SetParticleDefinition(
                            G4MuonPlus::Definition());
        } else {
                weight = 1.;
        }

        /* Generate the final state(s) */
        G4SingleParticleSource::GeneratePrimaryVertex(event);

        /* Loop over the generated state(s) */
        const G4int iVertex = event->GetNumberOfPrimaryVertex() - 1;
        G4PrimaryVertex * vertex = event->GetPrimaryVertex(iVertex);
        G4ThreeVector r = vertex->GetPosition() / CLHEP::m;
        GOUPIL_COORDINATES_CARTESIAN_CREATE_POINT(
            position, &gGlobals.g4frame, r[0], r[1], r[2]);

        const G4int nprim = vertex->GetNumberOfParticle();
        if ((G4int)fAncester.size() < nprim)
                fAncester.resize(nprim);

        int q = (int)this->GetParticleDefinition()->GetPDGCharge();
        for (G4int i = 0; i < nprim; i++) {
                /* Forward the ancester container to the recorder */
                struct goupil_recorder * recorder =
                    goupil_sampler_get_recorder(fSampler);
                SteppingHistory * history =
                    (SteppingHistory *)recorder->module.data;
                history->ancester = &fAncester[i];

                /* Transport the final state backward */
                G4PrimaryParticle * primary = vertex->GetPrimary(i);
                const double kinetic_energy =
                    primary->GetKineticEnergy() / CLHEP::GeV;
                G4ThreeVector u = primary->GetMomentumDirection();
                GOUPIL_COORDINATES_CARTESIAN_CREATE_VECTOR(direction,
                    &gGlobals.g4frame, u[0], u[1], u[2]);

                double flux;
                goupil_sampler_run(
                    fSampler, q, kinetic_energy, position, direction, &flux);

                /* Update the MC weight */
                weight *= primary->GetWeight();
                primary->SetWeight(weight * flux);
        }
}

/*
 * Overload the base setter in order to check that the primary particle is
 * indeed a mu+ or mu-
 */
void IGoupilSource::SetParticleDefinition(G4ParticleDefinition * aDefinition)
{
        if ((aDefinition == G4MuonMinus::Definition()) ||
            (aDefinition == G4MuonPlus::Definition())) {
                G4SingleParticleSource::SetParticleDefinition(aDefinition);
        } else {
                G4Exception("IGoupilSource::SetParticleDefinition",
                    "Invalid primary", FatalException,
                    "\tOnly mu- or mu+ primaries are supported\n");
        }
}

void IGoupilSource::SetSamplerMode(const std::string& mode) {
    fSamplerMode = mode;
}