// Because we want to simulate non-neutral evolution, we have to provide a
// custom initialization callback -- slendr will use it to replace its default
// neutral genomic architecture (i.e. the initialize() {...} callback it uses
// by default for neutral simulations). Note that we can refer to slendr's
// constants SEQUENCE_LENGTH and RECOMBINATION_RATE, which will carry values
// passed through from R via slendr's slim() R function.
initialize() {
    // define some parameters of the model
    defineConstant("s", 0.1);
    defineConstant("onset_time", 15000);
    defineConstant("target_pop", "EUR");

    initializeMutationType("m1", 0.5, "f", s);

    initializeGenomicElementType("g1", m1, 1.0);
    initializeGenomicElement(g1, 0, SEQUENCE_LENGTH - 1);

    initializeMutationRate(0);
    initializeRecombinationRate(RECOMBINATION_RATE);

    defineConstant("traj_file", paste0(PATH, "trajectory.tsv"));
}

function (void) add_mutation(void) {
    // sample one target carrier of the new mutation...
    target = sample(population(target_pop).haplosomes, 1);
    // ... and add the mutation in the middle of it
    mut = target.addNewDrawnMutation(m1, position = asInteger(SEQUENCE_LENGTH / 2));

    // save the mutation for later reference
    defineGlobal("MUTATION", mut);

    write_log("adding beneficial mutation to population " + target_pop);

    // write the header of the output file
    writeFile(traj_file, "time\tfrequency");
}

tick(onset_time) late() {
    // save simulation state in case we need to restart if the mutation is lost
    // (save_state() is a built-in function provided by slendr for customization)
    save_state();

    add_mutation();
}

tick(onset_time):SIMULATION_END late() {
    // the mutation is not segregating and is not fixed either -- we must restart
    if (!MUTATION.isSegregating & !MUTATION.isFixed) {
        write_log("mutation lost -- restarting");

        // reload the simulation state from just before we added the beneficial
        // mutation above (reset_state() is another built-in slendr function)
        reset_state();

        add_mutation();
    }

    // compute the frequency of the mutation of interest
    frequency = population("EUR").haplosomes.mutationFrequenciesInHaplosomes();

    // save the current frequency to the output file
    writeFile(traj_file,
              model_time(community.tick) + "\t" +
              frequency, append = T);
}
