#include <cassert>
#include <cmath>
#include <fstream>
#include <iostream>
#include <random>
#include <thread>
#define gpuErrchk(code) gpuAssert((code), FILE, LINE)
void gpuAssert(cudaError_t code, const char *file, int line) {
if (code != cudaSuccess) {
std::cerr << "gpuAssert failed " << file << ":" << line << ": " << cudaGetErrorString(code) << std::endl;
exit(code);
}
}

constexpr int m = 5, n = 5, cnt_monsters = 5, board_size = m * n;
//constexpr int cnt_states = pow(board_size, cnt_monsters) + .5;
static_assert(cnt_monsters <= 5);
constexpr int cnt_states = board_size * board_size * (cnt_monsters >= 3 ? board_size : 1) * (cnt_monsters >= 4 ? board_size : 1) * (cnt_monsters >= 5 ? board_size : 1);
//constexpr int cnt_outcomes = std::pow(5, cnt_monsters) + .5;
constexpr int cnt_outcomes = 5 * 5 * (cnt_monsters >= 3 ? 5 : 1) * (cnt_monsters >= 4 ? 5 : 1) * (cnt_monsters >= 5 ? 5 : 1);
//constexpr compute_t monster_move_prob = std::pow(0.2, cnt_monsters);  // monster_move_prob is 0.2 ^ cnt_monsters
constexpr float monster_move_prob = 1. / cnt_outcomes;

constexpr int antiverbosity = 10;
constexpr int k1 = 512, k2 = 128, k12 = k1 * k2;
using compute_t = float;

host device int get_state_index(const int state[cnt_monsters]) {
int index = 0;
for (int i = 0; i < cnt_monsters; ++i) {
index = index * board_size + state[i];
}
return index;
}
host device void get_state_from_index(int index, int state_out[cnt_monsters]) {
for (int i = cnt_monsters - 1; i >= 0; --i) {
state_out[i] = index % board_size;
index /= board_size;
}
}

// Function to update DP table for a range of states (single-threaded)
host device void update_dp_for_state(int state_idx, const compute_t** dp, compute_t** new_dp, int** next_move) {
const int directions[5][2] = {{-1, 0}, {1, 0}, {0, -1}, {0, 1}, {0, 0}};
int monster_outcomes[cnt_outcomes][cnt_monsters];
int initial_monster_state[cnt_monsters];
get_state_from_index(state_idx, initial_monster_state);
int total_outcomes = 1, outcome_i = 0;
for (int mi = 0; mi < cnt_monsters; ++mi) {
total_outcomes *= 5;
}
int direction_powers[cnt_monsters];
direction_powers[0] = 1;
for (int mi = 1; mi < cnt_monsters; ++mi) {
direction_powers[mi] = direction_powers[mi - 1] * 5;
}
while (outcome_i < total_outcomes) {
for (int mi = 0; mi < cnt_monsters; ++mi) {
int d = outcome_i / direction_powers[cnt_monsters - 1 - mi] % 5;
int mx = initial_monster_state[mi] / n, my = initial_monster_state[mi] % n, nx = mx + directions[d][0], ny = my + directions[d][1];
monster_outcomes[outcome_i][mi] = nx >= 0 && nx < m && ny >= 0 && ny < n ? nx * n + ny : initial_monster_state[mi];
}
++outcome_i;
}

// Inner loop over man's positions
for (int man_pos = 0; man_pos < board_size; ++man_pos) {
if (dp[man_pos][state_idx] == -1) { // If player is already dead in this state
new_dp[man_pos][state_idx] = -1;
continue;
}
int mx = man_pos / n, my = man_pos % n;
compute_t expected_scores[5];
// Initialize scores to a very low value to mark invalid moves
for(int d=0; d<5; ++d) {
expected_scores[d] = -1e18;
}
for (int d = 0; d < 5; ++d) { // Loop over man's 5 possible moves
int dx = directions[d][0], dy = directions[d][1], new_mx = mx + dx, new_my = my + dy;
if (new_mx < 0 || new_mx >= m || new_my < 0 || new_my >= n) {
expected_scores[d] = -1e18; // Mark as invalid move
continue;
}
int new_man_pos = new_mx * n + new_my;
compute_t total_score = 0;
for (size_t k = 0; k < cnt_outcomes; ++k) {
int new_monster_state_idx = get_state_index(monster_outcomes[k]);
if (dp[new_man_pos][new_monster_state_idx] != -1) {
total_score += monster_move_prob * (1 + dp[new_man_pos][new_monster_state_idx]);
}
}
expected_scores[d] = total_score;
}
int best_move = 0;
compute_t max_score = expected_scores[0];
for (int d = 1; d < 5; ++d) {
if (expected_scores[d] > max_score) {
max_score = expected_scores[d];
best_move = d;
}
}
new_dp[man_pos][state_idx] = max_score;
next_move[man_pos][state_idx] = best_move;
}
}
global void update_all_dp(const compute_t** dp, compute_t** new_dp, int** next_move) {
int id = blockIdx.x * blockDim.x + threadIdx.x;
for (int state_i = id; state_i < cnt_states; state_i += k12) {
update_dp_for_state(state_i, dp, new_dp, next_move);
}
}

// Function to compute divergence for a range of states (single-threaded)
host device int compute_divergence_for_state(int state_idx, const compute_t** old_dp, const compute_t** current_dp, compute_t &add_div, compute_t &add_sum) {
int count = 0;
for (int man_pos = 0; man_pos < board_size; ++man_pos) {
if (old_dp[man_pos][state_idx] != -1) { // Only consider reachable states
compute_t val1 = old_dp[man_pos][state_idx], val2 = current_dp[man_pos][state_idx];
assert(val2 >= val1);
add_div += val2 - val1;
add_sum += val2;
count++;
}
}
return count;
}
global void compute_all_divergences(const compute_t **old_dp, const compute_t **current_dp, compute_t *put_divs, compute_t *put_sums, int *put_counts) {
int id = blockIdx.x * blockDim.x + threadIdx.x;
compute_t div = 0, sum = 0;
int count = 0;
for (int state_i = id; state_i < cnt_states; state_i += k12) {
count += compute_divergence_for_state(state_i, old_dp, current_dp, div, sum);
}
put_divs[id] = div, put_sums[id] = sum, put_counts[id] = count;
}

int main() {
std::cout << m << "x" << n << " board with " << cnt_monsters << " monsters" << std::endl;
std::cout << "State space size: " << board_size * cnt_states / 1000000 << "M" << std::endl;
gpuErrchk(cudaDeviceSetLimit(cudaLimitStackSize, 100000));
std::clog << "Allocating memory (" << board_size * cnt_states * (2 * sizeof(compute_t) + sizeof(int)) / 1e9 << "GB)...";
// Allocate DP tables (dynamic 2D arrays)
compute_t dp, new_dp;
int next_move;
gpuErrchk(cudaMallocManaged(&dp, board_size * sizeof(compute_t)));
gpuErrchk(cudaMallocManaged(&new_dp, board_size * sizeof(compute_t)));
gpuErrchk(cudaMallocManaged(&next_move, board_size * sizeof(int)));
if (dp == NULL || new_dp == NULL || next_move == NULL) {
std::cerr << "Initial malloc failed" << std::endl;
return 1;
}
for (int i = 0; i < board_size; ++i) {
gpuErrchk(cudaMallocManaged(&dp[i], cnt_states * sizeof(compute_t)));
gpuErrchk(cudaMallocManaged(&new_dp[i], cnt_states * sizeof(compute_t)));
gpuErrchk(cudaMallocManaged(&next_move[i], cnt_states * sizeof(int)));
if (dp[i] == NULL || new_dp[i] == NULL || next_move[i] == NULL) {
std::cerr << "Row malloc failed" << std::endl;
return 1;
}
// Initialize with 0.0
memset(dp[i], 0, cnt_states * sizeof(compute_t));
memset(new_dp[i], 0, cnt_states * sizeof(compute_t));
memset(next_move[i], 0, cnt_states * sizeof(int));
}
compute_t divergences, sums;
int counts;
gpuErrchk(cudaMallocManaged(&divergences, k12 * sizeof(compute_t)));
gpuErrchk(cudaMallocManaged(&sums, k12 * sizeof(compute_t)));
gpuErrchk(cudaMallocManaged(&counts, k12 * sizeof(int)));
std::clog << "  done" << std::endl;
// Initialize states where man starts on a monster (-1 signifies death)
int monster_state_buffer[cnt_monsters];
for (int man_pos = 0; man_pos < board_size; ++man_pos) {
for (int state_idx = 0; state_idx < cnt_states; ++state_idx) {
get_state_from_index(state_idx, monster_state_buffer);
for (int i = 0; i < cnt_monsters; ++i) {
if (monster_state_buffer[i] == man_pos) {
dp[man_pos][state_idx] = -1;
new_dp[man_pos][state_idx] = -1;
break;
}
}
}
}
for (int t = 0;; ++t) {
if (t % antiverbosity == 0) {
std::cout << "Iteration " << t << std::endl;
}
// Call update function directly (single-threaded)
update_all_dp<<<k1,k2>>>((const compute_t)dp, new_dp, next_move);
gpuErrchk(cudaDeviceSynchronize());
// Call divergence function directly (single-threaded)
if (t % antiverbosity == 0) {
compute_t divergence = 0, sum = 0;
int count = 0;
compute_all_divergences<<<k1,k2>>>((const compute_t)dp, (const compute_t)new_dp, divergences, sums, counts);
gpuErrchk(cudaDeviceSynchronize());
for (int i = 0; i < k12; ++i) {
divergence += divergences[i], sum += sums[i], count += counts[i];
}
std::cout << "Divergence " << divergence << ", average survival time " << sum / count << std::endl;
if (divergence < 1e-6) break;
}
// Swap dp and new_dp pointers for next iteration
std::swap(dp, new_dp);
}
std::cout << "\nSample play:\n";
// SAMPLE PLAY → Fancy grid + write to file
// ========================================
{
std::ofstream out("sample_play");
if(!out){
std::cerr<<"Failed to open sample_play.txt for writing\n";
return 1;
}

int man_pos        = m * n - 1;
int monster_state  = 0;
int monsters[cnt_monsters], newm[cnt_monsters];
get_state_from_index(monster_state, monsters);

// RNG for monster moves
std::mt19937_64 rng{ std::random_device{}() };
std::uniform_int_distribution<int> uni(0, cnt_outcomes - 1);

// Pre-define directions
const int dir[5][2] = {{-1,0},{1,0},{0,-1},{0,1},{0,0}};

auto draw_grid = [&](int turn){
// Top border
out << "Turn " << turn << "\n";
out << "┌";
for(int j=0;j<n;j++){
out << "───";
out << (j+1<n ? "┬" : "┐");
}
out << "\n";

// Rows
for(int i=0;i<m;i++){
out << "│";
for(int j=0;j<n;j++){
int idx = i * n + j;
char c = '.';
if(idx == man_pos)             c = 'M';
else if(idx == monsters[0])    c = '1';
else if(cnt_monsters >= 2 && idx == monsters[1])    c = '2';
else if(cnt_monsters >= 3 && idx == monsters[2])    c = '3';
else if(cnt_monsters >= 4 && idx == monsters[3])    c = '4';
out << " " << c << " ";
out << "│";
}
out << "\n";

// Middle or bottom border
if(i+1 < m){
out << "├";
for(int j=0;j<n;j++){
out << "───";
out << (j+1<n ? "┼" : "┤");
}
out << "\n";
} else {
out << "└";
for(int j=0;j<n;j++){
out << "───";
out << (j+1<n ? "┴" : "┘");
}
out << "\n";
}
}
out << "\n";

};

out << "=== Sample Play Begins ===\n\n";

for(int turn=1; turn <= 10000; turn++){
if (turn % antiverbosity == 0) {
std::clog << "Sample rollout: turn " << turn << std::endl;
}
// 1) Draw board
draw_grid(turn);

// 2) Man’s best move
int mv = next_move[man_pos][monster_state];
int mx = man_pos / n, my = man_pos % n;
mx += dir[mv][0];
my += dir[mv][1];
man_pos = mx * n + my;

// 3) Sample one monster outcome and move monsters
int64_t choice_val = uni(rng); // Get a random combined movement choice

std::vector<int> newm_positions(cnt_monsters);
long long temp_choice_decoder = choice_val;

// Decode individual monster directions from the combined choice_val
// Monster 0 uses the most significant 'digit', Monster cnt_monsters-1 uses the least significant 'digit'
std::vector<int> monster_move_directions(cnt_monsters);
for (int i = cnt_monsters - 1; i >= 0; --i) {
monster_move_directions[i] = temp_choice_decoder % 5;
temp_choice_decoder /= 5;
}

// Move each monster according to its decoded direction
for (int i = 0; i < cnt_monsters; ++i) {
int x = monsters[i] / n; // Current row of monster i
int y = monsters[i] % n; // Current column of monster i
int chosen_dir_idx = monster_move_directions[i]; // Direction index for monster i

int nx = x + dir[chosen_dir_idx][0]; // New row
int ny = y + dir[chosen_dir_idx][1]; // New column

// Handle boundary conditions: if new position is out of bounds, monster stays in place
if (nx < 0 || nx >= m || ny < 0 || ny >= n) {
    nx = x; ny = y;
}
newm_positions[i] = nx * n + ny; // Store new position as a single integer

}

// Update the global monster positions and calculate the new combined monster state
long long new_combined_monster_state = 0;
for (int i = 0; i < cnt_monsters; ++i) {
monsters[i] = newm_positions[i]; // Update individual monster's position
new_combined_monster_state = new_combined_monster_state * board_size + newm_positions[i];
}
monster_state = new_combined_monster_state; // Update the overall combined state variable

// 4) Check capture
bool is_captured = false;
for (int i = 0; i < cnt_monsters; ++i) {
if (man_pos == monsters[i]) {
is_captured = true;
break; // Player captured by this monster, no need to check others
}
}

if (is_captured) {
draw_grid(turn);
out << ">>> Capture on turn " << turn << "! Game over.\n";
std::clog << ">>> Capture on turn " << turn << "! Game over.\n";
// The break; statement here implies it's part of a larger game loop,
// which is outside the scope of this code snippet.
break;
}

// optional pause in console demo
//std::this_thread::sleep_for(std::chrono::milliseconds(200));

}

out << "\n=== Sample Play Ends ===\n";
out.close();
}

//compute_t value = dp[m * n - 1][0];
//std::cout << "Expected survival time from man at (" << m - 1 << ", " << n - 1 << "), all monsters at (0, 0): " << value << std::endl;
int top_man_pos, top_monster_pos;
int state[cnt_monsters];
std::cout << "Best situation: ";
for (int man_pos = 0; man_pos < board_size; ++man_pos) {
for (int monster_pos = 0; monster_pos < cnt_states; ++monster_pos) {
if (dp[man_pos][monster_pos] > dp[top_man_pos][top_monster_pos]) {
top_man_pos = man_pos, top_monster_pos = monster_pos;
}
}
}
std::cout << "man @" << top_man_pos << " monsters ";
get_state_from_index(top_monster_pos, state);
for (int i = 0; i < cnt_monsters; ++i) {
std::cout << "@" << state[i] << " ";
}
std::cout << "value " << dp[top_man_pos][top_monster_pos] << std::endl;
std::cout << "Worst situation: ";
for (int man_pos = 0; man_pos < board_size; ++man_pos) {
for (int monster_pos = 0; monster_pos < cnt_states; ++monster_pos) {
if (dp[man_pos][monster_pos] < dp[top_man_pos][top_monster_pos] && dp[man_pos][monster_pos] != -1) {
top_man_pos = man_pos, top_monster_pos = monster_pos;
}
}
}
std::cout << "man @" << top_man_pos << " monsters ";
get_state_from_index(top_monster_pos, state);
for (int i = 0; i < cnt_monsters; ++i) {
std::cout << "@" << state[i] << " ";
}
std::cout << "value " << dp[top_man_pos][top_monster_pos] << std::endl;
return 0;
}

This code works correctly, but is too inefficient. The problem is complexity dependence on num_monsters: monsters' turn gets processed in 5^num_monsters time, which is awful
you can do in O(num_monsters) if you make them move one by one (the dp will be a bit larger)
Provide complete paste-ready code with two-space indentation
don't leave too many comments
be very careful, this is very important for my game theory research

IMPORTANT: your new dp should be +1-dimensional:
compute_t dp[man_pos][monsters_pos][whose_turn], where whose_turn may be man (0) or a monster (1 to cnt_monsters inclusive)
IMPORTANT: first think THOROUGHLY. which functions you will have, how exactly they will work. and only AFTER that, write complete, paste-ready code