main.cpp

//
//  Florian Probst
//  E-Mail: probstf@informatik.uni-freiburg.de / derbalvald@gmail.com
//

#pragma once

#include <corecrt_math_defines.h>
#include <dinput.h>
#include <vulkan/vulkan.h>
#include <GLFW/glfw3.h>
#include <iostream>
#include <string>
#include <filesystem>
#include <fstream>
#include <memory>
#include <thread>
#include <chrono>

#include "grid.h"
#include "sph_shared.h"
#include "vulkan_renderer.h"
#include "pause.h"
#include "iisph.h"
#include "wcsph.h"
#include "timing.h"

#define USE_MPI

#ifdef USE_MPI
#include <mpi.h>
#endif


void debug_particle(Particles &particles, const int index)
{
    // print info about particle at index
    std::cout << "Particle with index: " << index << ", pos: " << particles.positions.col(index)
                                                  << ", vel: " << particles.velocities.col(index)
                                                  << ", acc: " << particles.accelerations.col(index)
                                                  << ", dense: " << particles.densities(index)
                                                  << ", press: " << particles.pressures(index) << std::endl;
}

void debug_particles(Particles &particles, const std::vector<int> &indices)
{
    for (const auto index: indices)
    {
        debug_particle(particles, index);
    }
}

void write_xyz_file(std::string filename, const Particles& particles)
{
    std::ofstream file;
    file.open(filename);

    file << particles.positions.cols() << "\n";
    file << "Generated by SPHSolver\n";

    for (int i = 0; i < particles.positions.cols(); i++)
    {
        // also write accelerations

        if (particles.is_boundary(i))
        {
            file << "C " << particles.positions(0, i) << " " << particles.positions(1, i) << " " << particles.positions(2, i) << " " << particles.velocities(0, i) << " " << particles.velocities(1, i) << " " << particles.velocities(2, i) << " " << particles.accelerations(0, i) << " " << particles.accelerations(1, i) << " " << particles.accelerations(2, i) << "\n";
        }
        else
        {
            file << "H " << particles.positions(0, i) << " " << particles.positions(1, i) << " " << particles.positions(2, i) << " " << particles.velocities(0, i) << " " << particles.velocities(1, i) << " " << particles.velocities(2, i) << " " << particles.accelerations(0, i) << " " << particles.accelerations(1, i) << " " << particles.accelerations(2, i) << "\n";
        }
    }

    file.close();
}

void load_xyz_file(std::string filename, Particles& particles)
{
    std::ifstream file(filename);
    if (!file.is_open())
    {
        std::cerr << "Error: Could not open file " << filename << std::endl;
        return;
    }

    int num_particles;
    file >> num_particles;
    std::string line;
    std::getline(file, line); // consume the rest of the first line
    std::getline(file, line); // consume the comment line

    if (num_particles != particles.positions.cols())
    {
        std::cerr << "Error: Number of particles in file (" << num_particles << ") does not match Particles object (" << particles.positions.cols() << ")" << std::endl;
        return;
    }

    for (int i = 0; i < num_particles; i++)
    {
        char type;
        double px, py, pz, vx, vy, vz, ax, ay, az;
        if (!(file >> type >> px >> py >> pz >> vx >> vy >> vz >> ax >> ay >> az)) break;

        particles.positions.col(i) << px, py, pz;
        particles.velocities.col(i) << vx, vy, vz;
        particles.accelerations.col(i) << ax, ay, az;
        if (type == 'C')
        {
            particles.is_boundary[i] = true;
        }
        else
        {
            particles.is_boundary[i] = false;
        }
    }

    // Recalculate counts
    particles.num_boundary_particles = 0;
    particles.num_fluid_particles = 0;
    for (int i = 0; i < num_particles; i++)
    {
        if (particles.is_boundary[i])
            particles.num_boundary_particles++;
        else
            particles.num_fluid_particles++;
    }

    file.close();
}

void init_block(Particles& particles, int& index, Eigen::Vector3i size, Eigen::Vector3d origin, double h, bool is_boundary)
{
    for (int x = 0; x < size.x(); x++)
    {
        for (int y = 0; y < size.y(); y++)
        {
            for (int z = 0; z < size.z(); z++)
            {
                if (index >= particles.positions.cols()) return;

                particles.positions.col(index) << origin.x() + x * h, origin.y() + y * h, origin.z() + z * h;
                particles.velocities.col(index).setZero();
                particles.accelerations.col(index).setZero();
                if (is_boundary)
                {
                    particles.mark_as_boundary(index);
                }
                else
                {
                    particles.is_boundary[index] = false;
                }
                index++;
            }
        }
    }
}

void init_container(Particles& particles, int& index, Eigen::Vector3i size, Eigen::Vector3d origin, double h)
{
    // Floor
    init_block(particles, index, Eigen::Vector3i(size.x(), 1, size.z()), origin, h, true);
    // Wall x=0
    init_block(particles, index, Eigen::Vector3i(1, size.y(), size.z()), origin + Eigen::Vector3d(0, h, 0), h, true);
    // Wall x=max
    init_block(particles, index, Eigen::Vector3i(1, size.y(), size.z()), origin + Eigen::Vector3d((size.x() - 1) * h, h, 0), h, true);
    // Wall z=0
    init_block(particles, index, Eigen::Vector3i(size.x() - 2, size.y(), 1), origin + Eigen::Vector3d(h, h, 0), h, true);
    // Wall z=max
    init_block(particles, index, Eigen::Vector3i(size.x() - 2, size.y(), 1), origin + Eigen::Vector3d(h, h, (size.z() - 1) * h), h, true);
    // Roof
    init_block(particles, index, Eigen::Vector3i(size.x() - 2, 1, size.z() - 2), origin + Eigen::Vector3d(h, size.y() * h, h), h, true);
}


int main(int argc, char *argv[])
{
    int rank = 0, size = 1;

#ifdef USE_MPI
    MPI_Init(&argc, &argv);

    // Retrieve process info
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
#endif

    // Default parameters
    std::string method = "iisph";
    std::string search_method = "ugm";
    std::string load_path = "";
    int max_time_steps = 10000;
    int save_interval = 1;
    double dt = 0.01;
    double h = 1.0;
    double stiffness = 50000.0;
    double rest_density = 1000.0;
    double nu = 0.1;
    double gamma = 7.0;
    double omega = 0.5;
    int lbvh_update_freq = 5;

    bool save_only_latest = false;
    const std::string latest_particles_filename = "particles.xyz";

    bool no_visualization = false;

    // Fluid block dimensions
    Eigen::Vector3i fluid_size(10, 10, 10);
    // Boundary box dimensions
    Eigen::Vector3i boundary_size(14, 14, 14);

    // Argument parsing
    for (int i = 1; i < argc; i++)
    {
        std::string arg = argv[i];
        if (arg == "--method" && i + 1 < argc)
        {
            method = argv[++i];
        }
        else if (arg == "--search" && i + 1 < argc)
        {
            search_method = argv[++i];
        }
        else if (arg == "--load" && i + 1 < argc)
        {
            load_path = argv[++i];
        }
        else if (arg == "--steps" && i + 1 < argc)
        {
            max_time_steps = std::stoi(argv[++i]);
        }
        else if (arg == "--save_interval" && i + 1 < argc)
        {
            save_interval = std::stoi(argv[++i]);
        }
        else if (arg == "--dt" && i + 1 < argc)
        {
            dt = std::stod(argv[++i]);
        }
        else if (arg == "--h" && i + 1 < argc)
        {
            h = std::stod(argv[++i]);
        }
        else if (arg == "--stiffness" && i + 1 < argc)
        {
            stiffness = std::stod(argv[++i]);
        }
        else if (arg == "--rho0" && i + 1 < argc)
        {
            rest_density = std::stod(argv[++i]);
        }
        else if (arg == "--viscosity" && i + 1 < argc)
        {
            nu = std::stod(argv[++i]);
        }
        else if (arg == "--gamma" && i + 1 < argc)
        {
            gamma = std::stod(argv[++i]);
        }
        else if (arg == "--omega" && i + 1 < argc)
        {
            omega = std::stod(argv[++i]);
        }
        else if (arg == "--lbvh_update_freq" && i + 1 < argc)
        {
            lbvh_update_freq = std::stoi(argv[++i]);
        }
        else if (arg == "--save_latest")
        {
            // If set, always overwrite the same file so only the latest snapshot exists
            save_only_latest = true;
        }
        else if (arg == "--no_visualization" || arg == "--no_visualisation")
        {
            // Run simulation without creating the graphical visualisation/UI
            no_visualization = true;
        }
        else if (arg == "--fluid_size" && i + 3 < argc)
        {
            fluid_size.x() = std::stoi(argv[++i]);
            fluid_size.y() = std::stoi(argv[++i]);
            fluid_size.z() = std::stoi(argv[++i]);
        }
        else if (arg == "--boundary_size" && i + 3 < argc)
        {
            boundary_size.x() = std::stoi(argv[++i]);
            boundary_size.y() = std::stoi(argv[++i]);
            boundary_size.z() = std::stoi(argv[++i]);
        }
        else if (arg == "--help" || arg == "-h")
        {
            if (rank == 0)
            {
                std::cout << "Usage: " << argv[0] << " [options]\n"
                          << "Options:\n"
                          << "  --method <method>       Simulation method (iisph or wcsph)\n"
                          << "  --search <method>       Neighborhood search method (shi, sh, ug, ugm, sg, sgi, lbvh, naive)\n"
                          << "  --load <file>           Path to xyz file to load scenario from\n"
                          << "  --steps <n>             Maximum number of time steps\n"
                          << "  --save_interval <n>     Save xyz file every n steps\n"
                          << "  --dt <dt>               Time step size\n"
                          << "  --h <h>                 Smoothing length\n"
                          << "  --stiffness <k>         Stiffness parameter for WCSPH\n"
                          << "  --rho0 <rho>            Rest density\n"
                          << "  --viscosity <nu>        Viscosity parameter\n"
                          << "  --gamma <gamma>         Gamma for WCSPH\n"
                          << "  --omega <omega>         Omega for IISPH\n"
                          << "  --lbvh_update_freq <n>  How often to rebuild LBVH (default: 1)\n"
                          << "  --fluid_size <x> <y> <z>  Size of fluid block\n"
                          << "  --boundary_size <x> <y> <z> Size of boundary container\n"
                          << "  --save_latest            Overwrite a single file (particles.xyz) each save so only the latest snapshot is present\n"
                          << "  --no_visualization       Run without the graphical visualisation (headless)\n";
            }
#ifdef USE_MPI
            MPI_Finalize();
#endif
            return 0;
        }
    }

    if (rank == 0)
    {
        std::cout << "Hello I am rank " << rank << " of " << size << "\n";
        std::cout << "Method: " << method << "\n";
        std::cout << "Max steps: " << max_time_steps << "\n";
        std::cout << "DT: " << dt << "\n";

		std::string test = "";

		if (search_method == "lbvh")
		{
			test = "_lbvh_freq_" + std::to_string(lbvh_update_freq);
		}

        // Generate log filename
        std::string log_filename = "sim_" + method + "_" + search_method + "_dt" + std::to_string(dt) + "_h" +
                                   std::to_string(h) + "_size" + std::to_string(fluid_size.x()) +
                                   std::to_string(fluid_size.y()) + std::to_string(fluid_size.z()) + test + ".log";
        Logger::getInstance().setLogFile(log_filename);
    }

    int num_all_particles = 0;
    if (!load_path.empty())
    {
        std::ifstream file(load_path);
        if (file.is_open())
        {
            file >> num_all_particles;
            file.close();
        }
        else
        {
            std::cerr << "Error: Could not open file " << load_path << std::endl;
            return 1;
        }
    }
    else
    {
        int num_fluid = fluid_size.x() * fluid_size.y() * fluid_size.z();
        int num_boundary = 1 * boundary_size.x() * 1 * boundary_size.z() +
                           2 * (1 * boundary_size.y() * boundary_size.z()) +
                           2 * ((boundary_size.x() - 2) * boundary_size.y() * 1) +
                           1 * ((boundary_size.x() - 2) * 1 * (boundary_size.z() - 2));
        num_all_particles = num_fluid + num_boundary;
    }

    double particle_mass = rest_density * (h * h * h);
    Particles particles(3, num_all_particles, particle_mass, h, rest_density);
    particles.lbvh_update_freq = lbvh_update_freq;

    if (!load_path.empty())
    {
        load_xyz_file(load_path, particles);
    }
    else
    {
        int index = 0;
        init_block(particles, index, fluid_size, Eigen::Vector3d(-0.0, -0.0, -0.0), h, false);
        init_container(particles, index, boundary_size, Eigen::Vector3d(-2.0, -2.0, -2.0), h);
    }

    // Initialise neighbors for boundary mass calculation



    if (search_method == "shi")
    {
        SpatialHashGrid_Ihmsen shi(particles.h * 2, particles.get_num_particles());
        particles.find_neighbors_shi(shi);
    }
    else if (search_method == "sh")
    {
        SpatialHashGrid sh(particles.h * 2);
        particles.find_neighbors_sh(sh);
    }
    else if (search_method == "ug")
    {
        UniformGrid ug;
        particles.find_neighbors_ug(ug);
    }
    else if (search_method == "ugm")
    {
        UniformGrid_Morton ugm;
        particles.find_neighbors_ugm(ugm);
    }
    else if (search_method == "sg")
    {
        SortedHashGrid sg(particles.h * 2);
        particles.find_neighbors_sg(sg);
    }
    else if (search_method == "sgi")
    {
        SortedHashGrid_Ihmsen sgi(particles.h * 2, particles.get_num_particles());
        particles.find_neighbors_sgi(sgi);
    }
    else if (search_method == "lbvh")
    {
        particles.find_neighbors_lbvh(particles.h * 2);
    }
    else if (search_method == "naive")
    {
        particles.find_neighbors_naive();
    }
    else
    {
        SpatialHashGrid_Ihmsen shi(particles.h * 2, particles.get_num_particles());
        particles.find_neighbors_shi(shi);
    }
    compute_boundary_masses(particles);

    // Create the visualisation application only if requested
    std::unique_ptr<SPHVisualisationApplication> application = nullptr;
    if (!no_visualization)
    {
        application.reset(new SPHVisualisationApplication(particles));
        application->init();
    }

    int i = 0;
    while (i < max_time_steps)
    {
        // If visualisation is active, poll window events and check for close
        if (!no_visualization && application)
        {
            if (glfwWindowShouldClose(application->window)) { break; }
            glfwPollEvents();
            application->processInput(application->window);
        }
        else
        {
            // headless: no-op for event polling
        }

        // If visualization is enabled and the user requested pause, skip advancing the simulation
        if (!no_visualization && application && g_paused.load())
        {
            // Still update the visualiser so camera moves/inputs are reflected
            application->updateShaderStorageBuffers(particles);
            application->drawFrame();

            // Sleep briefly to avoid busy-looping while paused
            std::this_thread::sleep_for(std::chrono::milliseconds(10));

            // Do not increment i — simulation time is paused
            continue;
        }

        // Run simulation step (headless or not paused)
        if (rank == 0) {
            // std::cout << "Step " << i << "\n";
            Logger::getInstance().logStep(i);
        }

        if (method == "iisph")
        {
            simulation_step_iisph(particles, dt, omega, gamma, rest_density, nu, search_method);
        }
        else
        {
            simulation_step(particles, dt, stiffness, gamma, rest_density, nu, search_method);
        }

        // Write the output file if we want to.
        if (i % save_interval == 0)
        {
            if (save_only_latest)
            {
                write_xyz_file(latest_particles_filename, particles);
            }
            else
            {
                write_xyz_file("particles_" + std::to_string(i) + ".xyz", particles);
            }
        }

        if (!no_visualization && application)
        {
            application->updateShaderStorageBuffers(particles);
            application->drawFrame();
        }

        // Advance simulation time step only after we've performed the step
        i++;
    }

    if (!no_visualization && application)
    {
        vkDeviceWaitIdle(application->device);
        application->cleanup();
    }

    std::cout << "Goodbye from rank " << rank << "\n";

#ifdef USE_MPI
    MPI_Finalize();
#endif

    return 0;
}