#include "iom.h"
#include <iostream>
#include <omp.h>

const int dir[8][2] = {{1,1},{1,0},{1,-1},{0,-1},{0,1},{-1,-1},{-1,0},{-1,1}};
inline int simulateCell(std::vector<std::vector<int>>& oldWorld, std::vector<std::vector<int>>& newWorld, int x, int y) {
    int worldWidth = oldWorld.size();
    int worldHeight = oldWorld[0].size();
    int count[10] = {0};
    // Total of cells surrounding
    // #pragma omp sections
    {for(auto d: dir) {
        int nx = d[0] + x;
        int ny = d[1] + y;

        if (nx < 0) {
            nx += worldWidth;
        } else if (nx >= worldWidth) {
            nx -= worldWidth;
        }
        if (ny < 0) {
            ny += worldHeight;
        } else if (ny >= worldHeight) {
            ny -= worldHeight;
        }

        // #pragma omp critical
        count[oldWorld[nx][ny]]++;
    }}
    int cur =  oldWorld[x][y];
    if(cur == 0) {
        if (count[0] == 8) return 0;
        // new birth
        for(int i = 9; i > 0; i--) if(count[i] ==3) {
            newWorld[x][y] = i;
            return 0;
        }
        return 0;
    }

    // hostile
    for(int i = 1; i <= 9; ++i) if(i != cur && count[i] != 0) {
        newWorld[x][y] = 0;
        return 1;
    }

    // under/over population
    if(count[cur] >= 4 || count[cur] < 2) {
        newWorld[x][y] = 0;
        return 0;
    }

    // keep same
    newWorld[x][y] = cur;
    return 0;
}

inline int invade(std::vector<std::vector<int>>& curWorld, std::vector<std::vector<int>>& invPlan, std::vector<std::vector<int>>& oldWorld) {
    int sum = 0;
    int i, j;
    // #pragma omp parallel for shared(curWorld, invPlan, oldWorld, sum) private(i,j)
    for(i = 0; i < (int)curWorld.size(); ++i)
        for(j = 0; j < (int)curWorld[0].size(); ++j) {
            if (invPlan[i][j] != 0) {
                if((curWorld[i][j] != 0 && oldWorld[i][j] != 0)) {
                    // #pragma omp critical
                        sum++;
                }
            }
            curWorld[i][j] = invPlan[i][j];
        }
    return sum;
}

int iom(int nThreads, int nGenerations, std::vector<std::vector<int>>& startWorld, int nRows, int nCols, int nInvasions, std::vector<int> invasionTimes, std::vector<std::vector<std::vector<int>>> invasionPlans) {

    // omp_set_num_threads(nThreads);

    int totalCount = 0;
    int curInvasionIdx = 0;


    for(int i = 1; i <= nGenerations; ++i ) {
        std::vector<std::vector<int>> newWorld(nRows,std::vector<int>(nCols));
        int x,y;
        // #pragma omp parallel for private(x,y) shared(startWorld, newWorld)
        for(x = 0; x < nRows; ++x) for(y = 0; y < nCols; ++y)  {
            int res = simulateCell(startWorld,newWorld,x,y);
            // #pragma omp critical
            totalCount += res;
        }

        if(nInvasions > 0 && invasionTimes[curInvasionIdx] == i){
            totalCount += invade(newWorld,invasionPlans[curInvasionIdx], startWorld);
            curInvasionIdx++;
            nInvasions--;
        }

        std::swap(newWorld,startWorld); // 1 way to swap
        // #pragma omp parallel for private(x, y)
        // for(int x = 0; x < nRows; ++x) for(int y = 0; y < nCols; ++y) startWorld[x][y] = newWorld[x][y];
    }
    return totalCount;
}