main.cpp

#include <iostream>
#include <filesystem>
#include "Particle.h"
#include "Node.h"
#include <random>
#include "rapidcsv.h"
#include <cmath>
#include <chrono>
#include <future>

#define SIM_AREA 10000 // Parsecs
//#define DEBUG // Define DEBUG to enable debug mode

/**
 * Function to generate a vector of particles with random positions, velocities, and masses
 * @param numParticles: Number of particles to generate
 * @param posMin: Minimum position value in pc (default 0)
 * @param posMax: Maximum position value in pc (default SIM_AREA/10)
 * @param velMin: Minimum velocity value in km/s (default 0)
 * @param velMax: Maximum velocity value in km/s (default 300)
 * @param massMin: Minimum mass value in solar masses (default 0.001)
 * @param massMax: Maximum mass value in solar masses (default 100)
 * @return A vector of pointers to Particle objects
 */
std::vector<Particle*> generateParticles(int numParticles, int posMin = -SIM_AREA/10, int posMax = SIM_AREA/10,
                                         int velMin = 0, int velMax = 300,
                                         double massMin = 0.001, double massMax = 100) {
    std::vector<Particle*> particles;
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<> posDis(posMin, posMax); // Parsecs
    std::uniform_real_distribution<> velDis(velMin, velMax); // km/s
    std::uniform_real_distribution<> massDis(massMin, massMax); // Solar masses - When more refined can use an initial mass function

    for (int i = 0; i < numParticles; ++i) {
        Particle* p = new Particle (
            posDis(gen), posDis(gen), posDis(gen),
            velDis(gen), velDis(gen), velDis(gen),
            0, 0, 0,
            massDis(gen));
        p->id = std::to_string(i); // Set unique ID for the particle
        particles.push_back(p);
    }
    return particles;
}

/**
 * Function to check if a string is an integer
 * @param str: String to check
 * @return true if the string is an integer, false otherwise
 */
bool isInt(const char* str) {
    // Check if the string is a number
    for (int i = 0; str[i] != '\0'; ++i) {
        if (!isdigit(str[i]) && str[i] != '-' && str[i] != '+') {
            return false;
        }
    }
    return true;
}


/** Function to check if a string is an integer
 * @param str: String to check
 * @return true if the string is an integer, false otherwise
 */
bool isInt(const std::string& str) {
    // Check if the string is a number
    for (char c : str) {
        if (!isdigit(c) && c != '-' && c != '+') {
            return false;
        }
    }
    return true;}


/** Function to check if a string is a float
 * @param str: String to check
 * @return true if the string is a float, false otherwise
 */
bool isFloat(const std::string& str) {
    // Check if the string is a float
    bool hasDecimal = false;
    for (char c : str) {
        if (!isdigit(c)) {
            if (c == '.' && !hasDecimal) {
                hasDecimal = true;
            } else if (c != '-' && c != '+') {
                return false;
            }
        }
    }
    return true;
}


/** Function to check if a string is a float
 * @param str: String to check
 * @return true if the string is a float, false otherwise
 */
bool isFloat(const char* str) {
    // Check if the string is a float
    bool hasDecimal = false;
    for (int i = 0; str[i] != '\0'; ++i) {
        if (!isdigit(str[i])) {
            if (str[i] == '.' && !hasDecimal) {
                hasDecimal = true;
            } else if (str[i] != '-' && str[i] != '+') {
                return false;
            }
        }
    }
    return true;
}


/**
 * @struct inputOptions
 * @brief Structure to hold input options for the simulation
 * @var numIterations: Number of iterations for the simulation
 * @var outputDirLocation: Directory location for output files
 * @var outputHistoryFile: Name of the history file
 * @var outputParticlesFile: Name of the particles file
 * @var generateParticles: Flag to indicate if particles should be generated
 * @var numParticles: Number of particles to generate if generateParticles is true
 * @var useInputFile: Flag to indicate if an input file is used
 * @var useGaiaData: Flag to indicate if Gaia data is used
 * @var inputFile: Location of the input file
 * @var pheta: Threshold for size to distance ratio
 * @var simArea: Simulation area in parsecs
 * @var timeDelta: Time step for the simulation in seconds
 */
struct {
    int numIterations = 1000; // Default number of iterations
    std::string outputDirLocation = "~/Documents/AstroSimOut/"; // Default output file location
    std::string outputHistoryFile = "history"; // Default history file name
    std::string outputParticlesFile = "particles"; // Default particles file name
    bool generateParticles = false;
    int numParticles = 100; // Default number of particles to generate
    bool useInputFile = false; // Flag to indicate if an input file is used
    bool useGaiaData = false; // Flag to indicate if Gaia data is used
    std::string inputFile = ""; // Default input file location, empty means no input file
    float pheta = 0.5; // Default threshold for size to distance ratio
    long double simArea = SIM_AREA; // Default simulation area in parsecs
    double timeDelta = 86400; // Default time step for the simulation in seconds (1 day)
} inputOptions;


/**
 * Function to handle input arguments and set the inputOptions structure
 * @param argc: Number of arguments
 * @param argv: Array of argument strings
 */
void handleInput(int argc, char* argv[]) {
    // Handle input arguments
    if (argc < 3) {
        std::cerr << "Usage: " << argv[0] << " <numIterations> <outputFileLocation> "
                                             "[-g <numParticles>] [-i <inputDir> <isGaiaData(y/n)>] "
                                             "[-p <phetaValue(0-1)>] [-t <timeDelta(secs)>] "
                                             "[-s <simArea(Pc)>] " << std::endl;
        exit(EXIT_FAILURE);
    }
    if (!isInt(argv[1])) {
        std::cerr << "1st argument: numIterations - Input must be an integer" << std::endl;
        exit(EXIT_FAILURE);
    }
    inputOptions.numIterations = std::stoi(argv[1]);

    if (!std::filesystem::is_directory(argv[2])) {
        std::cerr << "2nd argument: inputDir - Input must be a directory" << std::endl;
        exit(EXIT_FAILURE);
    }
    inputOptions.outputDirLocation = argv[2];
    for (int i = 3; i < argc; ++i) {
        if (std::string(argv[i]) == "-g") {
            inputOptions.generateParticles = true;
            if (i + 1 > argc) {
                std::cerr << "Error: -g option requires a number of particles." << std::endl;
                exit(EXIT_FAILURE);
            }
            inputOptions.numParticles = std::stoi(argv[++i]);
        }
        if (std::string(argv[i]) == "-i") {
            inputOptions.useInputFile = true;
            if (i + 1 > argc) {
                std::cerr << "Error: -i option requires an input file." << std::endl;
                exit(EXIT_FAILURE);
            }
            inputOptions.inputFile = argv[++i];
            if (!std::filesystem::exists(inputOptions.inputFile)) {
                std::cerr << "Error: Input file does not exist." << std::endl;
                exit(EXIT_FAILURE);
            }
            if (i + 2 < argc && std::string(argv[i + 1]) == "y") {
                inputOptions.useGaiaData = true;
                ++i; // Skip the next argument
            } else {
                inputOptions.useGaiaData = false;
            }
        }
        if (std::string(argv[i]) == "-p") {
            if (i + 1 > argc || !isFloat(argv[i + 1])) {
                std::cerr << "Error: -p option requires a float value between 0 and 1." << std::endl;
                exit(EXIT_FAILURE);
            }
            inputOptions.pheta = std::stof(argv[++i]);
            if (inputOptions.pheta < 0 || inputOptions.pheta > 1) {
                std::cerr << "Error: -p value must be between 0 and 1." << std::endl;
                exit(EXIT_FAILURE);
            }
        }
        if (std::string(argv[i]) == "-t") {
            if (i + 1 > argc || !isFloat(argv[i + 1])) {
                std::cerr << "Error: -t option requires a float value for time step in seconds." << std::endl;
                exit(EXIT_FAILURE);
            }
            inputOptions.timeDelta = std::stof(argv[++i]);
        }
        if (std::string(argv[i]) == "-s") {
            if (i + 1 > argc || !isInt(argv[i + 1])) {
                std::cerr << "Error: -s option requires a int value for simulation area in parsecs." << std::endl;
                exit(EXIT_FAILURE);
            }
            inputOptions.simArea = std::stof(argv[++i]);
        }
    }
    if (inputOptions.generateParticles && inputOptions.useInputFile) {
        std::cerr << "Error: Cannot use -g and -i options together." << std::endl;
        exit(EXIT_FAILURE);
    }
    std::cout << "Outputting to: " << inputOptions.outputDirLocation << std::endl;
    std::cout << "Number of iterations: " << inputOptions.numIterations << std::endl;
    if (inputOptions.generateParticles) {
        std::cout << "Generating " << inputOptions.numParticles << " particles." << std::endl;
    }
    if (inputOptions.useInputFile) {
        std::cout << "Using input file: " << inputOptions.inputFile << std::endl;
    }
}


/**
 * Function to read particles from a CSV file and convert them to Particle objects
 * @param filePath: Path to the CSV file containing particle data
 * @return A vector of Particle pointers
 */
std::vector<Particle*> getParticlesFromFile(const std::string& filePath) {;
    // Read particles from a CSV file using rapidcsv
    rapidcsv::Document doc(filePath, rapidcsv::LabelParams(0, -1));
    std::vector<Particle*> particles;
    for (size_t i = 0; i < doc.GetRowCount(); ++i) {
        Particle* p;
        p->id = doc.GetCell<std::string>("id", i); // Unique identifier for the particle
        p->posX = doc.GetCell<long double>("posX", i);
        p->posY = doc.GetCell<long double>("posY", i);
        p->posZ = doc.GetCell<long double>("posZ", i);
        p->velX = doc.GetCell<long double>("velX", i);
        p->velY = doc.GetCell<long double>("velY", i);
        p->velZ = doc.GetCell<long double>("velZ", i);
        p->accelX = doc.GetCell<double>("accelX", i);
        p->accelY = doc.GetCell<double>("accelY", i);
        p->accelZ = doc.GetCell<double>("accelZ", i);
        p->mass = doc.GetCell<double>("mass", i);
        particles.push_back(p);
    }
    return particles;
}


/**
 * Function to read particles from a Gaia data file and convert them to Particle objects
 * @param filePath: Path to the Gaia data file (CSV format)
 * @return A vector of Particle objects
 */
std::vector<Particle> getParticlesFromGaiaData(const std::string& filePath) {
    // Read particles from a Gaia data file (CSV format)
    rapidcsv::Document doc(filePath, rapidcsv::LabelParams(0, -1));
    std::vector<Particle> particles;
    for (size_t i = 0; i < doc.GetRowCount(); ++i) {
        Particle p;
        double k = 4.74047; // Conversion factor: proper motion (mas/yr) and distance (pc) to velocity (km/s)
        auto ra = doc.GetCell<double>("ra", i); // Right Ascension (deg)
        ra = ra * M_PI / 180.0; // Convert RA to radians
        auto dec = doc.GetCell<double>("dec", i); // Declination (deg)
        dec = dec * M_PI / 180.0; // Convert Dec to radians
        auto parallax = doc.GetCell<double>("parallax", i); // Parallax (in milliarcseconds)
        parallax = 1000/parallax; // Convert parallax to parsecs
        auto raProperMotion = doc.GetCell<double>("pmra", i); // Proper motion in RA (mas/yr)
        auto decProperMotion = doc.GetCell<double>("pmdec", i); // Proper motion Dec (mas/yr)
        auto radVel = doc.GetCell<double>("radial_velocity", i); // Radial velocity (km/s)
        if (parallax == 0) {
            std::cerr << "Warning: Parallax is zero for particle " << i << ", skipping." << std::endl;
            continue; // Skip particles with zero parallax
        }
        double velRa = (k * raProperMotion / 1000) * parallax; // Convert proper motion in RA to velocity in km/s
        double velDec = (k * decProperMotion / 1000) * parallax; // Convert proper motion in Dec to velocity in km/s

        // Equatorial to Galactic rotation matrix (IAU 1958)
        double rot[3][3] = {
            {-0.0548755604, -0.8734370902, -0.4838350155},
            { 0.4941094279, -0.4448296300,  0.7469822445},
            {-0.8676661490, -0.1980763734,  0.4559837762}
        };

        // Convert RA, Dec, and Parallax to Cartesian coordinates
        p.id = doc.GetCell<string>("source_id", i); // Unique identifier from Gaia

        double equposX = parallax * cos(dec) * cos(ra);
        double equposY = (parallax * cos(dec) * sin(ra));
        double equposZ = (parallax * sin(dec));
        p.posX = rot[0][0]*equposX + rot[0][1]*equposY + rot[0][2]*equposZ;
        p.posY = rot[1][0]*equposX + rot[1][1]*equposY + rot[1][2]*equposZ;
        p.posZ = rot[2][0]*equposX + rot[2][1]*equposY + rot[2][2]*equposZ;

        double equVelX = radVel * cos(dec) * cos(ra) + velRa * -sin(ra) + velDec * (-cos(ra) * sin(dec));
        double equVelY = radVel * cos(dec) * sin(ra) + velRa * cos(ra) + velDec * (-sin(ra) * sin(dec));
        double equVelZ = radVel * sin(dec) + velDec * cos(dec);

        // Rotate to Galactic velocity
        p.velX = rot[0][0]*equVelX + rot[0][1]*equVelY + rot[0][2]*equVelZ;
        p.velY = rot[1][0]*equVelX + rot[1][1]*equVelY + rot[1][2]*equVelZ;
        p.velZ = rot[2][0]*equVelX + rot[2][1]*equVelY + rot[2][2]*equVelZ;
        p.accelX = 0;
        p.accelY = 0;
        p.accelZ = 0;
        p.mass = 1; // 1 for now TODO: Add initial mass function for distributing masses of stars
        particles.push_back(p);
    }
    return particles;
}


/**
 * Function to write particles to a CSV file
 * @param particles: Vector of Particle pointers to write to the file
 * @param filePath: Path to the output CSV file
 */
void writeParticlesToFile(const std::vector<Particle*>& particles, const std::string& filePath) {
    // Check if file exists
    std::ifstream infile(filePath);
    bool fileExists = infile.good();
    infile.close();

    if (!fileExists) {
        // Create file and write headers
        std::ofstream outfile(filePath);
        outfile << "id,posX,posY,posZ,velX,velY,velZ,accelX,accelY,accelZ,mass\n";
        outfile.close();
    }

    rapidcsv::Document doc(filePath, rapidcsv::LabelParams(0, -1));
    for (auto i = 0; i < particles.size(); ++i) {
        doc.SetCell("id", i, particles[i]->id);
        doc.SetCell("posX", i, particles[i]->posX);
        doc.SetCell("posY", i, particles[i]->posY);
        doc.SetCell("posZ", i, particles[i]->posZ);
        doc.SetCell("velX", i, particles[i]->velX);
        doc.SetCell("velY", i, particles[i]->velY);
        doc.SetCell("velZ", i, particles[i]->velZ);
        doc.SetCell("accelX", i, particles[i]->accelX);
        doc.SetCell("accelY", i, particles[i]->accelY);
        doc.SetCell("accelZ", i, particles[i]->accelZ);
        doc.SetCell("mass", i, particles[i]->mass);
    }
    doc.Save();
}


/**
 * Function to append the current state of particles to a history file
 * @param particleHistoryMap: Map containing time as key and vector of Particle pointers as value
 * @param doc: Document object for the history file
 * @return true if successful, false otherwise
 */
bool appendToHistoryFile(const std::map<double, std::vector<Particle*>> &particleHistoryMap, rapidcsv::Document &doc) {
    // Append the current state of particles to a history file
    try {
        unsigned long rowCount;
        for (const auto& [time, particles] : particleHistoryMap) {
            for (const auto& particle : particles) {
                rowCount = doc.GetRowCount();
                doc.SetCell("time", rowCount, time);
                doc.SetCell("id", rowCount, particle->id);
                doc.SetCell("posX", rowCount, particle->posX);
                doc.SetCell("posY", rowCount, particle->posY);
                doc.SetCell("posZ", rowCount, particle->posZ);
                doc.SetCell("velX", rowCount, particle->velX);
                doc.SetCell("velY", rowCount, particle->velY);
                doc.SetCell("velZ", rowCount, particle->velZ);
                doc.SetCell("accelX", rowCount, particle->accelX);
                doc.SetCell("accelY", rowCount, particle->accelY);
                doc.SetCell("accelZ", rowCount, particle->accelZ);
                doc.SetCell("mass", rowCount, particle->mass);
            }
        }
        doc.Save();
        std::cout << "Appended " << particleHistoryMap.size() << " entries to history file." << std::endl;
        return true;
    }
    catch (const std::exception& e) {
        std::cerr << "Error appending to history file: " << e.what() << std::endl;
        return false;
    }
}

const int secInYears = 31536000; // Number of seconds in a year

/* Input goes as:
 * Arg 0: Name of the program
 * Arg 1: Number of iterations
 * Arg 2: Where to store the output
 * -g <numParticles>: Number of particles to generate
 * -i input file: Input file to read particles from
 * */
int main(int argc, char * argv[]) {
    const std::vector<std::string> csvHeaders = {
        "time", "id", "posX", "posY", "posZ",
        "velX", "velY", "velZ", "accelX", "accelY", "accelZ", "mass"
    };

#ifndef DEBUG
    handleInput(argc, argv);
#else
    inputOptions.outputDirLocation = "/home/lucy/CLionProjects/AstroSim/";
    inputOptions.numIterations = 1001;
    inputOptions.generateParticles = true;
    inputOptions.numParticles = 100;
    inputOptions.useInputFile = false;
#endif
    std::vector<Particle*> particlesAddr;
    if (inputOptions.generateParticles) {
        particlesAddr = generateParticles(inputOptions.numParticles);
    }
    else if (inputOptions.useInputFile) {
        particlesAddr = getParticlesFromFile(inputOptions.inputFile);
    }
    std::cout << inputOptions.outputDirLocation << std::endl;
    writeParticlesToFile(particlesAddr, inputOptions.outputDirLocation + "/particles" + ".csv");
    Node root(0, 0, 0, inputOptions.simArea);
    for (Particle* particleAddr : particlesAddr) {
        root.addParticle(particleAddr);
        particleAddr->printParticle();
    }
    root.makeChildren();
    root.calculateCOM();

    std::ifstream csvFile(inputOptions.outputDirLocation + inputOptions.outputHistoryFile + ".csv");
    if (!csvFile.good()) {
        // If the file does not exist, create it and write headers
        std::ofstream outFile(inputOptions.outputDirLocation + inputOptions.outputHistoryFile + ".csv");
        outFile.close();
    }
    else {
        int i = 0;
        while (true) {
            std::string path = inputOptions.outputDirLocation + inputOptions.outputHistoryFile + "(" + std::to_string(i) + ").csv";
            std::ifstream csvFile(path);
            if (!csvFile.good()) {
                // If the file does not exist, create it and write headers
                std::ofstream outFile(path);
                outFile.close();
                inputOptions.outputHistoryFile = inputOptions.outputHistoryFile + "(" + std::to_string(i) + ")";
                std::cout << "Created new history file: " << path << std::endl;
                break;
            }
            i++;
        }
    }

    rapidcsv::Document doc(inputOptions.outputDirLocation + inputOptions.outputHistoryFile + ".csv", rapidcsv::LabelParams(0, -1));
    // Add headers to the history file
    for (int i = 0; i < csvHeaders.size(); ++i) {
        doc.SetColumnName(i, csvHeaders[i]);
    }


    std::map<double, std::vector<Particle*>> particlesHistoryMap;
    std::future<bool> futureBool;
    double timeOutput = 0;
    for (int i =0; i < inputOptions.numIterations; i++) {
        root.updateParticlesRoot(inputOptions.timeDelta, inputOptions.pheta);
        timeOutput = (i ==0) ? 0 : (i * inputOptions.timeDelta) / secInYears; // Convert time to years
        particlesHistoryMap [timeOutput] = root.particles;
        if (i % 1000 == 0) {
            if (futureBool.valid()) {
                // Wait for the previous append operation to finish
                futureBool.get();
                // Remove the particles from the history map to avoid overlapping entries in the csv
            }
            std::cout << "Iteration: " << i << "\nAppending to file" << std::endl;
            // Copy the particles history map to avoid data loss
            std::map<double, std::vector<Particle*>> particlesHistoryMapCopy = particlesHistoryMap;
            particlesHistoryMap.erase(particlesHistoryMap.begin(), particlesHistoryMap.end());
            futureBool = std::async(std::launch::async, appendToHistoryFile, particlesHistoryMapCopy, std::ref(doc));
        }
        std::cout << "Iteration: " << i << std::endl;
    }
    // root.updateParticlesRoot(3e9, 0.5);
    // root.printNode();
    return 0;
}