Untitled

#include <omp.h>
#include <array>
#include <vector>
#include <iostream>
#include "collision.h"
#include "io.h"
#include "sim_validator.h"

inline int GetTileFromDir(int curTileIndex, std::pair<int, int> dir, int tilesPerSide) {
    return (dir.second * tilesPerSide) + curTileIndex + dir.first;
}

inline int DetermineCellType(int xCoord, int yCoord) {
    int xParity = xCoord & 1;          // Same as xCoord % 2
    int yParity = yCoord & 1;          // Same as yCoord % 2
    return xParity + (yParity << 1) + 1;
}

inline bool IsInBounds(int tileIndex, int totalTileCount) {
    return (tileIndex >= 0 && tileIndex < totalTileCount);
}

inline int GetTileFromCoords(int xCoord, int yCoord, int tilesPerSide) {
    return (yCoord * tilesPerSide) + xCoord;
}

inline void addParticle(std::vector<std::vector<int>>* tileToParticles, Particle* particle, int tilesPerSide, int tileLength, omp_lock_t* tileLocks) {
    int xCoord = std::max(0, std::min(static_cast<int>(particle->loc.x / tileLength), tilesPerSide - 1));
    int yCoord = std::max(0, std::min(static_cast<int>(particle->loc.y / tileLength), tilesPerSide - 1));
    int tileIndex = GetTileFromCoords(xCoord, yCoord, tilesPerSide);

    // Lock the tile before modifying
    omp_set_lock(&tileLocks[tileIndex]);
    (*tileToParticles)[tileIndex].push_back(particle->i);
    omp_unset_lock(&tileLocks[tileIndex]);
}

inline void clearGrid(std::vector<std::vector<int>>* tileToParticles, int totalTileCount) {
    #pragma omp parallel for
    for (int i = 0; i < totalTileCount; ++i) {
        (*tileToParticles)[i].clear();  // Dereference the pointer and access the vector
    }
}

int main(int argc, char* argv[]) {
    // Read arguments and input file
    Params params{};
    std::vector<Particle> particles;
    read_args(argc, argv, params, particles);

    // Set number of threads
    omp_set_num_threads(params.param_threads);
    static constexpr std::array<std::pair<int, int>, 4> directions = {
        {{0, 1}, {1, 0}, {1, 1}, {1, -1}}};
#if CHECK == 1
    // Initialize collision checker
    SimulationValidator validator(params.param_particles, params.square_size, params.param_radius);
    // Initialize with starting positions
    validator.initialize(particles);
    // Uncomment the line below to enable visualization (makes program much slower)
    // validator.enable_viz_output("test.out");
#endif

    // TODO: this is the part where you simulate particle behavior.

    /*
    After simulating each timestep, you must call:

    #if CHECK == 1
    validator.validate_step(particles);
    #endif
    */

    // code start here
    // Initialize per-particle locks for thread safety
    int tileLength = params.param_radius * 2;
    int tilesPerSide = params.square_size / tileLength;
    int totalTileCount = tilesPerSide * tilesPerSide;
    std::vector<std::vector<int>> tileToParticles(totalTileCount, std::vector<int>());
    std::vector<omp_lock_t> tileLocks(totalTileCount);
    for (int i = 0; i < totalTileCount; ++i) {
        omp_init_lock(&tileLocks[i]);
    }

    for (int i = 0; i < params.param_steps; i++) {
        // update all particle's position
        #pragma omp parallel for
        for (int j = 0; j < params.param_particles; j++) {
            particles[j].loc.x += particles[j].vel.x;
            particles[j].loc.y += particles[j].vel.y;
        }

        
        clearGrid(&tileToParticles, totalTileCount);
        // check for particle collisions
        #pragma omp parallel for
        for (int j = 0; j < params.param_particles; j++) {
            addParticle(&tileToParticles, &(particles[j]), tilesPerSide, tileLength, tileLocks.data());
        }

        // for every particle in the queue we check against the rest
        bool collisionsOccurred = true;
        while (collisionsOccurred) {
            collisionsOccurred = false;
            for (int color = 1; color <= 4; ++color) {
                int xStart = (color == 2 || color == 4) ? 1 : 0;
                int yStart = (color >= 3) ? 1 : 0;

                #pragma omp parallel for schedule(guided) collapse(2)
                for (int x = xStart; x < tilesPerSide; x += 2) {
                    for (int y = yStart; y < tilesPerSide; y += 2) {
                        int tileIndex = GetTileFromCoords(x, y, tilesPerSide);
                        std::vector<int> particlesInTile = tileToParticles[tileIndex];
                        for(int particleIndex : particlesInTile) {
                            Particle* particle = &(particles[particleIndex]);
                            // Check and resolve wall collision
                            if (x == 0 || y == 0 || x == tilesPerSide - 1 || y == tilesPerSide - 1) {
                                if (is_wall_collision(particle->loc, particle->vel, params.square_size, params.param_radius)) {
                                    resolve_wall_collision(particle->loc, particle->vel, params.square_size, params.param_radius);
                                    collisionsOccurred = true;
                                }
                            }
                            for(int pIndex : particlesInTile) {
                                if (pIndex == particleIndex) continue;  // Avoid double processing

                                Particle* otherParticle = &(particles[pIndex]);

                                // Check and resolve particle collision
                                if (is_particle_collision(particle->loc, particle->vel, otherParticle->loc,
                                                        otherParticle->vel, params.param_radius)) {
                                    resolve_particle_collision(particle->loc, particle->vel, otherParticle->loc,
                                                            otherParticle->vel);
                                    collisionsOccurred = true;
                                }
                            }
                            for (auto dir : directions) {
                                int nextTileIndex = GetTileFromDir(tileIndex, dir, tilesPerSide);
                                if (IsInBounds(nextTileIndex, totalTileCount)) {
                                    for (int pIndex : tileToParticles[nextTileIndex]) {
                                        Particle* otherParticle = &(particles[pIndex]);

                                        // Check and resolve particle collision
                                        if (is_particle_collision(particle->loc, particle->vel, otherParticle->loc,
                                                                otherParticle->vel, params.param_radius)) {
                                            resolve_particle_collision(particle->loc, particle->vel, otherParticle->loc,
                                                                    otherParticle->vel);
                                            collisionsOccurred = true;
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }



#if CHECK == 1
        validator.validate_step(particles);
#endif
    }

    // code end here
#if CHECK == 1
    // Check final positions
    // validator.validate_step(particles);
#endif
}