// This file is part of necsim project which is released under MIT license.
// See file **LICENSE.txt** or visit https://opensource.org/licenses/MIT) for full license details.

/**
 * @file Tree.h
 * @brief  Contains the main simulation object for spatially implicit coalescence simulations.
 *
 * Provides the basis for spatially explicit versions in SpatialTree, and protracted speciation versions in
 * ProtractedTree and ProtractedSpatialTree.
 *
 * @copyright <a href="https://opensource.org/licenses/MIT"> MIT Licence.</a>
 */
#ifndef TREE_H
#define TREE_H
//#ifndef sql_ram
//#define sql_ram
//#endif

#include <sqlite3.h>
#include <string>

#ifdef CXX14_SUPPORT
#include "memory.h"
#else

#include <memory>

#endif

#include "TreeNode.h"
#include "Matrix.h"
#include "SimParameters.h"
#include "RNGController.h"
#include "DataPoint.h"
#include "Community.h"
#include "file_system.h"
#include "custom_exceptions.h"
#include "Step.h"
#include "SQLiteHandler.h"

using namespace random_numbers;
namespace necsim
{
    /**
     * @brief Main simulation class for performing a non-spatial neutral simulation and generating the phylogenetic tree of
     * the individuals.
     */
    class Tree
    {
    protected:
        // storing the coalescence tree itself
        shared_ptr<vector<TreeNode>> data;
        // a reference for the last written point in data.
        unsigned long enddata{};
        // Stores the command line current_metacommunity_parameters and parses the required information.
        shared_ptr<SimParameters> sim_parameters;
        // random number generator
        shared_ptr<RNGController> NR;
        // Storing the speciation rates for later reference.
        vector<long double> speciation_rates;
        // flag for having set the simulation seed.
        bool seeded{};
        // random seed
        long long seed{};
        // for file naming - good to know which task in a series is being executed here
        long long task{};
        // The map file containing the times that we want to expand the model and record all lineages again.
        // If this is null, uses_temporal_sampling will be false and the vector will be empty.
        string times_file;
        vector<double> reference_times;
        // Set to true if we are recording at times other than the present day.
        bool uses_temporal_sampling{};
        // The time variables (for timing the simulation in real time)
        time_t start{}, sim_start{}, sim_end{}, now{}, sim_finish{}, out_finish{};
        time_t time_taken{};
        // Active lineages stored as a row of datapoints
        vector<DataPoint> active;
        // Stores the point of the end of the active vector. 0 is reserved as null
        unsigned long endactive{};
        // the maximum size of endactive
        unsigned long startendactive{};
        // the maximum simulated number of individuals in the present day.
        unsigned long maxsimsize{};
        // for create the link to the speciationcounter object which handles everything.
        Community community;
        // This might need to be updated for simulations that have large changes in maximum population size over time.
        // number of simulation num_steps
        long steps{};
        // Maximum time to run for (in seconds)
        unsigned long maxtime{};
        // number of generations passed,
        double generation{};
        // The number of individuals per cell
        double deme{};
        // The proportion of individuals to sample
        double deme_sample{};
        // the speciation rate
        long double spec{};
        // Path to output directory
        string out_directory;
        // sqlite3 object that stores all the data
        shared_ptr<SQLiteHandler> database;
        // only set to true if the simulation has finished, otherwise will be false.
        bool sim_complete{};
        // set to true when variables are imported
        bool has_imported_vars{};
        // If sql database is written first to memory, then need another object to contain the in-memory database.
#ifdef sql_ram
        SQLiteHandler outdatabase;
#endif
        // Create the step object that will be retained for the whole simulation.
        // Does not need saving on simulation pause.
        Step this_step;
        string sql_output_database;
        // If true, means the command-line imports were under the (deprecated) fullmode.
        bool bFullMode{};
        // If true, the simulation is to be resumed.
        bool bResume{};
        // If true, a config file contains the simulation variables.
        bool bConfig{};
        // If true, simulation can be resumed.
        bool has_paused{}, has_imported_pause{};
        // Should always be false in the base class
        bool bIsProtracted{};
        // variable for storing the paused sim location if files have been moved during paused/resumed simulations!
        string pause_sim_directory;
        // Set to true to use the gillespie method - this is currently only supported for spatial simulations using a
        // dispersal map and point speciation (i.e. the method is unsupported for non-spatial simulations,
        // spatial simulations not using a dispersal map and those that use protracted speciation).
        bool using_gillespie{};

    public:
        Tree() : data(make_shared<vector<TreeNode >>()), enddata(0), sim_parameters(make_shared<SimParameters>()),
                 NR(make_shared<RNGController>()), speciation_rates(), seeded(false), seed(-1), task(-1),
                 times_file("null"), reference_times(), uses_temporal_sampling(false), start(0), sim_start(0),
                 sim_end(0), now(0), sim_finish(0), out_finish(0), time_taken(0), active(), endactive(0),
                 startendactive(0), maxsimsize(0), community(data), steps(0), maxtime(0), generation(0.0), deme(0.0),
                 deme_sample(0.0), spec(0.0), out_directory(""), database(make_shared<SQLiteHandler>()),
                 sim_complete(false), has_imported_vars(false),
#ifdef sql_ram
                 outdatabase(),
#endif //sql_ram
                 this_step(), sql_output_database("null"), bFullMode(false), bResume(false), bConfig(true),
                 has_paused(false), has_imported_pause(false), bIsProtracted(false), pause_sim_directory("null"),
                 using_gillespie(false)
        {

        }

        Tree(Tree &&other) noexcept;

        Tree(const Tree &other) noexcept;

        virtual ~Tree()
        {
            if(database != nullptr)
            {
                database->close();
            }
#ifdef sql_ram
            outdatabase.close();
#endif
        }

        Tree &operator=(const Tree &other);

        Tree &operator=(Tree &&other);

        /**
         * @brief Import the simulation variables from the command line structure.
         *
         * This function parses the simulation variables, imports them from the config file,
         * checks that the input files exist and checks for any paused simulations. The flags are then set correctly,
         * meaning that setup() and runSimulation() can be run immediately afterwards.

         * @param configfile the path to the config file containing parameters to parse.
         */
        void importSimulationVariables(string configfile);

        /**
         * @brief Import the simulation variables from a ConfigOption.
         *
         * This function parses the simulation variables, imports them (from either the command line or a config file),
         * checks that the input files exist and checks for any paused simulations. The flags are then set correctly,
         * meaning that setup() and runSim() can be run immediately afterwards.

         * @param config the set of config parameters to import
         */
        void importSimulationVariables(ConfigParser config);

        /**
         * @brief Import the simulation variables from a string containing the config file data
         *
         * This function parses the simulation variables, imports them (from either the command line or a config file),
         * checks that the input files exist and checks for any paused simulations. The flags are then set correctly,
         * meaning that setup() and runSim() can be run immediately afterwards.

         * @param config_string the set of config parameters to import as a string
         */
        void importSimulationVariablesFromString(string config_string);

        /**
         * @brief Runs the basic file existence checks.
         * Checks for paused simulations and file existence.
         */
        virtual void runFileChecks();

        /**
         * @brief Resets all the simulation variables.
         */
        void wipeSimulationVariables();

        /**
         * @brief Sets up the simulation parameters from the one provided.
         *
         * Intended for usage with metacommunity application. No output directory is expected.
         * @param sim_parameters_in the simulation parameters to set up the simulation with
         */
        void internalSetup(shared_ptr<SimParameters> sim_parameters_in);

        /**
         * @brief Asserts that the output directory is not null and exists. If it doesn't exist, it attempts to create it.
         * @throws Fatal_Exception if the output directory creation fails
         * @return true if output creates successfully
         */
        bool checkOutputDirectory();

        /**
         * @brief Checks for existing paused simulations to resume from
         * Sets bPaused if there are.
         *
         * This version uses the default values read from the config file.
         */
        void checkSims();

        /**
         * @brief Checks for existing paused simulations to resume from
         *
         * Sets bPaused if there are.
         *
         * This version uses the values supplied to the function directly
         *
         * @param output_dir the output directory to check for
         * @param seed the seed for paused sims
         * @param task the task for paused sims
         */
        void checkSims(string output_dir, long seed, long task);

        /**
         * @brief Move the parameters from the sim_parameters object to their relevant parameters.
         */
        virtual void setParameters();

        /**
         * @brief Sets the protracted variables
         * @param speciation_gen_min the minimum number of generations to have passed before speciation is allowed
         * @param speciation_gen_max the maximum number of generations a lineage can exist for before it is speciated.
         */
        virtual void setProtractedVariables(double speciation_gen_min, double speciation_gen_max);

        /**
         * @brief Gets the has_paused variable for resuming sims.
         * @return if the simulation has paused
         */
        bool hasPaused();

        /**
         * @brief Gets the map autocorrel times
         */
        vector<double> getTemporalSampling();

        /**
        * @brief Getter for the simulation seed.
        * @return Returns the seed
        */
        virtual long long getSeed();

        /**
        * @brief Gets the job type for the simulation.
        * This is a reference number for the jobs.
        * @return Returns the job type
        */
        virtual long long getJobType();

        /**
         * @brief Sets the simulation seed for the random number generator.
         *
         * This function should only be called once.
         *
         * The seed is set within the NR  object. This will be fixed for the simulation and is only performed once.
         *
         * @param seed_in the desired seed to set for the simulation
         */
        void setSeed(long long seed_in);

        /**
         * @brief Gets the generation timer for the simulation
         * @return the number of generations that have passed
         */
        double getGeneration() const;

        /**
         * @brief Gets the initial number of individuals
         * @return the number of initial individuals to simulate
         */
        virtual unsigned long getInitialCount();

        /**
         * @brief Sets the sizes of grid, active and data, based on the number of individuals counted from the samplemask.
         * @return a count of the number of individuals that exist in the simulation
         */
        unsigned long setObjectSizes();

        /**
         * @brief The setup function for generating the simulation objects.
         *
         * The simulation parameters are set from comargs using setParameters().
         * Generates and fills the active and grid objects as well as importing all the maps from the supplied files.
         */
        virtual void setup();

        /**
         * @brief Sets the starting values for required parameters.
         */
        void setInitialValues();

        /**
         * @brief Sets the variables at the start of a simulation for temporary data.
         *
         * This is not the main set-up routine, which creates the permanent data structures including maps,
         * the coalescence tree and active lineage listings.
         */
        void setSimStartVariables();

        /**
         * @brief Prints the statement for the setup initiation.
         *
         * This is stored in a separate function so that it can be called in isolation by child classes.
         */
        void printSetup();

        /**
         * @brief Sets the temporal sampling points from the time config file.
         */
        void setTimes();

        /**
         * @brief Determines the speciation rates to apply and then applies them to the coalescence tree post-simulation.
         *
         * Detects speciation rates from the config files supplied.
         */
        void determineSpeciationRates();

        /**
         * @brief Adds the speciation rates to those to be applied.
         */
        void addSpeciationRates(vector<long double> spec_rates_in);

        /**
         * @brief Assigns the objects sizes in memory and fills with the starting lineages.
         */
        void generateObjects();

        /**
         * @brief Fills the active and data objects with the starting lineages.
         * @param initial_count the number of individuals expected to exist
         */
        virtual unsigned long fillObjects(const unsigned long &initial_count);

        /**
         * @brief Runs a single loop of the simulation.
         */
        void runSingleLoop();

        /**
         * @brief Run the entire simulation given the start conditions already defined by setup().
         *
         * Setup is assumed to have been run already. This function is the main function containing the main loop of the
         * simulation.
         * At the end of the simulation, returns true if the simulation is complete, false otherwise.
         * @return true if the simulation has completed
         */
        bool runSimulation();

        /**
         * @brief Runs the simulation using the standard coalescence algorithm.
         * @return true if the simulation has completed
         */
        bool runSimulationNoGillespie();

        /**
         * @brief Writes to the console that the simulation is beginning
         */
        void writeSimStartToConsole();

        /**
         * @brief Write the step counter to console. This function should only be called in debugging mode.
         */
        void writeStepToConsole();

        /**
         * @brief Increments the generation counter and step references
         */
        virtual void incrementGeneration();

        /**
         * @brief Chooses a random lineage from active.
         *
         * The index of the random lineage is stored in this_step, as chosen.
         * Also records the required variables for the step process, like x, y position.
         */
        void chooseRandomLineage();

        /**
         * @brief Updates the coalescence variables in the step object.
         */
        virtual void updateStepCoalescenceVariables();

        /**
         * @brief Speciation to supplied lineage.
         *
         * Also calls the removeOldPos() and switchPositions() functions for removing the lineage out of active reference.
         * @param chosen
         */
        void speciation(const unsigned long &chosen);

        /**
         * @brief Performs the actual speciation.
         * @param data_position the position in the array of TreeNodes for this lineage
         */
        virtual void speciateLineage(const unsigned long &data_position);

        /**
         * @brief Removes the old position within active.
         * @param chosen the desired active reference to remove from the grid.
         */
        virtual void removeOldPosition(const unsigned long &chosen);

        /**
         * @brief Switches the chosen position with the endactive position.
         * @param chosen the chosen lineage to switch with endactive.
         */
        virtual void switchPositions(const unsigned long &chosen);

        /**
         * @brief Calculates the next step for the simulation.
         */
        virtual void calcNextStep();

        /**
         * @brief Calculates the speciation probability from the random number, speciation rate and number of generations a
         * lineage has existed for.
         * @param random_number the generated random number from 0-1
         * @param speciation_rate the speciation rate to be applied
         * @param no_generations the number of generations a lineage has existed for
         * @return if true, speciation has occured
         */
        virtual bool calcSpeciation(const long double &random_number,
                                    const long double &speciation_rate,
                                    const unsigned long &no_generations);

        /**
        * @brief  Perform the coalescence between lineages. Once coalesced, lineages are removed from the active scope.
        * @param chosen the chosen lineage for coalescence
        * @param coalchosen the target lineage for coalscence
        */
        void coalescenceEvent(const unsigned long &chosen, unsigned long &coalchosen);

        /**
         * @brief Checks if the number of lineages should be expanded at another sample point
         */
        void checkTimeUpdate();

        /**
         * @brief Adds the required lineages at the generation time.
         * @param generation_in the generation to add lineages at
         */
        virtual void addLineages(double generation_in);

        /**
         * @brief Randomly checks if a lineage should be added, based on the proportion added.
         * @param proportion_added the proportion of lineages that should be added
         * @return true if the lineage should be added, false otherwise
         */
        bool checkProportionAdded(const double &proportion_added);

        /**
         * @brief Checks the size of the main active and data objects is large enough
         * @param req_data the required data object size
         * @param req_active the required active object size
         */
        void checkSimSize(unsigned long req_data, unsigned long req_active);

        /**
         * @brief Sets the active reference to a tip, if it isn't one already. Otherwise, creates a new tip for the new
         * generation time.
         * @param tmp_active the reference in active
         * @param generation_in the generation to set for the new lineage
         * @param data_added vector of lineages to add to data
         */
        void makeTip(const unsigned long &tmp_active, const double &generation_in, vector<TreeNode> &data_added);

        /**
         * @brief Creates a new reference in data containing the tip with a new generation counter.
         * @param i the reference in active of the lineage to make a tip
         * @param generationin the generation to make the lineage a tip at
         * @param data_added vector of lineages to add to the data object
         */
        void createNewTip(unsigned long i, double generationin, vector<TreeNode> &data_added);

        /**
         * @brief Finalises the simulation, and performs the correct tasks depending if the sim has been paused or finished.
         * @return
         */
        bool stopSimulation();

        /**
         * @brief Applies the given speciation rate to the tree.
         * @param sr the required speciation rate
         * @param t the required time of speciation
         */
        void applySpecRate(long double sr, double t);

        /**
         * @brief Applies the given speciation rate to the tree, but does not output to a file. Instead returns a pointer
         * to the nodes object.
         * @param sr the required speciation rate
         * @param t the required time of speciation
         */
        void applySpecRateInternal(long double sr, double t);

        /**
         * @brief Gets the sorted cumulative species abundances from the contained TreeList.
         *
         * For use with metacommunity applications
         * @return row of cumulative species abundances
         */
        shared_ptr<vector<unsigned long>> getCumulativeAbundances();

        /**
         * @brief Gets the species abundances from the internal tree.
         * @param community_reference the community reference
         * @return the species abundances
         */
        shared_ptr<map<unsigned long, unsigned long>> getSpeciesAbundances(const unsigned long &community_reference);

        /**
         * @brief Gets the species abundances from the internal tree.
         * @return the species abundances
         */
        shared_ptr<vector<unsigned long>> getSpeciesAbundances();

        /**
         * @brief Sets up Community member to point to the same output database as the simulation.
         * @return the protracted speciation parameters to set up
         */
        ProtractedSpeciationParameters setupCommunity();

        /**
         * @brief Sets up the generation of the tree object.
         * @param sr the required speciation rate
         * @param t the required time of speciation
         */
        void setupCommunityCalculation(long double sr, double t);

        /**
         * @brief Overloaded version of applySpecRates for the default generation (0.0).
         * @param sr the speciation rate to apply to the tree
         */
        void applySpecRate(long double sr);

        /**
         * @brief Applies multiple speciation rates to the coalescence tree, ignoring repeated speciation rates.
         *
         * Speciation rates are read from the speciation_rates object, which should have already been calculated.
         */
        void applyMultipleRates();

        /**
         * @brief Gets the protractedness of the simulation.
         * Overridden by protracted child classes.
         * @return
         */
        virtual bool getProtracted();

        /**
         * @brief Gets the protracted variables and returns them as a single, newline separated string.
         * This method is intended to be overridden in derived classes.
         * @return string containing the protracted variables, separated by newlines.
         */
        virtual string getProtractedVariables();

        /**
         * @brief Gets the minimum number of generations a lineage must exist.
         *
         * Without overriding this function, should always return 0.0.
         * @return double the number of generations a lineage must exist
         */
        virtual double getProtractedGenerationMin();

        /**
         * @brief Gets the maximum number of generations a lineage can exist.
         *
         * Without overriding this function, should always return 0.0 (no maximum).
         *
         * @return double the number of generations a lineage must exist
         */
        virtual double getProtractedGenerationMax();

        /**
         * @brief Copy the in-memory database to file.
         *
         * This function should not be called if the database is already opened on disc, and won't do anything if it is.
         */
        void sqlOutput();

        /**
         * @brief Creates the output data objects and outputs important simulation data to a database file.
         */
        void createAndOutputData();

        /**
         * @brief Outputs the simulation data (including the coalescence tree) to an sql database.
         */
        void outputData();

        /**
         * @brief Sort and process the species list so that the useful information can be extracted from it.
         */
        void sortData();

        /**
         * @brief Writes the times to the terminal for simulation information.
         */
        void writeTimes();

        /**
         * @brief Opens a connection to the in-memory database, or the on-disk database, depending on the compilation
         * options.
         */
        void openSQLDatabase();

        /**
         * @brief Generates the SQL database file from the full simulation data. This allows for greater analysis of the
         * data after completion of the simulation
         */
        void sqlCreate();

        /**
         * @brief if the sql_output_database variable is "null", sets the output database and directory.
         *
         * Creates the output directory, if it doesn't already exist, and deletes any existing database with the
         * output name.
         */
        void setupOutputDirectory();

        /**
         * @brief Creates the SIMULATION_PARAMETERS table in the SQL database.
         */
        void sqlCreateSimulationParameters();

        /**
         * @brief Creates a string containing the SQL insertion statement for the simulation parameters.
         * @return string containing the SQL insertion statement
         */
        virtual string simulationParametersSqlInsertion();

        /**
         * @brief Outputs the protracted variables to a string.
         *
         * This function is intended to be overridden by derived classes.
         * It is intended the output is used for writing to SQL databases.
         *
         * @return string containing a list of the protracted speciation variables.
         */
        virtual string protractedVarsToString();

        /**
         * @brief Pause the simulation and dump data from memory.
         */
        virtual void simPause();

        /**
         * @brief Checks the output folder exists and initiates the pause.
         * @return the output file stream to save objects to
         */
        shared_ptr<ofstream> initiatePause();

        /**
         * @brief Saves the main simulation variables to file.
         * @param out the output file stream to save the object to
         */
        void dumpMain(shared_ptr<ofstream> out);

        /**
         * @brief Saves the active object to file.
         * @param out the output file stream to save the object to
         */
        void dumpActive(shared_ptr<ofstream> out);

        /**
         * @brief Saves the data object to file.
         * @param out the output file stream to save the object to
         */
        void dumpData(shared_ptr<ofstream> out);

        /**
         * @brief Completes the pause routine and outputs the sql dump.
         * @param out the output stream to close up
         */
        void completePause(shared_ptr<ofstream> out);

        /**
         * @brief Sets the resume variables so that the simulation can be resumed.
         *
         * The pause directory can be the same as the output directory if it is desirable to save to the same location.

         * @param pausedir the directory containing the pause folder for resuming the simulation
         * @param outdir the directory to write simulation output to
         * @param seed the simulation seed
         * @param task the simulation job reference number
         * @param new_max_time the maximum simulation time to run for in seconds (0 keeps old simulation max time)
         */
        void setResumeParameters(string pausedir,
                                 string outdir,
                                 unsigned long seed,
                                 unsigned long task,
                                 unsigned long new_max_time);

        /**
         * @brief Sets the resume variables to the defaults
         */
        void setResumeParameters();

        shared_ptr<ifstream> openSaveFile();

        /**
         * @brief Loads the main simulation parameters from the save file into memory.
         */
        virtual void loadMainSave(shared_ptr<ifstream> in1);

        /**
         * @brief Loads the data object from the save file into memory.
         */
        void loadDataSave(shared_ptr<ifstream> in1);

        /**
         * @brief Loads the active object from the save file into memory.
         */
        void loadActiveSave(shared_ptr<ifstream> in1);

        /**
         * @brief Checks for resuming and prints to the terminal
         */
        void initiateResume();

        /**
         * @brief Resumes the simulation from a previous state.
         *
         * Reads in the parameters and objects from file and re-starts the simulation.
         */
        virtual void simResume();

        virtual void addGillespie(const double &g_threshold);

        virtual bool runSimulationGillespie();

#ifdef DEBUG

        /**
         * @brief Validates all lineages have been set up correctly. This may take considerable time for larger simulations.
         *
         * This function is only relevant in debug mode.
         */
        virtual void validateLineages();

        /**
         * @brief Validates the coalescence tree.
         */
        void validateCoalescenceTree();

        /**
         * @brief Runs the debug checks at the end of each step.
         *
         * Should only be run if compiled with the DEBUG flag.
         */
        virtual void debugEndStep();

        /**
         * @brief Runs the debug checks upon a coalescence event.
         *
         * Should only be run if compiled with the DEBUG flag.
         */
        void debugCoalescence();

        /**
         * @brief Run checks at the end of each cycle which make certain the step has been successful.
         * @param chosen the chosen lineage to check
         * @param coalchosen the lineage which is coalescing with the chosen lineage which we are also required to check
         */
        virtual void runChecks(const unsigned long &chosen, const unsigned long &coalchosen);

        /**
         * @brief Runs a mini check on the chosen lineage.
         * @param chosen the lineage to check for
         */
        void miniCheck(const unsigned long &chosen);

#endif // DEBUG
    };
}
#endif //TREE_H
