#include "prison.h"
#include <bits/stdc++.h>
using namespace std;
#define bv(x, i) (((x) & (1 << (i))) ? 1 : 0)
#define trit(x, i) (((x) / static_cast<int>(std::pow(3, (i)))) % 3)
using vi = vector<int>;
using vvi = vector<vi>;
struct State {
int bitID, id;
// if id = 0, then current bag is A
// if id = 1, then current bag is B, and A[bit] = 0
// if id = 2, then current bag is B, and A[bit] = 1
};
const int BITS = 8;
const int BASE = 3;
const int STATES = 30;
vector<State> states;
int stateID[BITS][BASE + 1];
void build_state_table() {
for (int i = BITS - 1; i >= 0; i--) {
for (auto state : {0, 2, 1, 3}) {
stateID[i][state] = (int) states.size();
states.push_back(State {i, state});
}
}
}
int transition(State current, int bagValue) {
State newState = current;
if (current.bitID == 0 && current.id == 0) {
if (trit(bagValue, current.bitID) == 0) {
return -1;
} else if (trit(bagValue, current.bitID) == 2) {
return -2;
}
}
if (current.id == 0) {
newState.id = 1 + trit(bagValue, current.bitID);
newState.bitID = current.bitID;
} else {
if (current.id - 1 != trit(bagValue, current.bitID)) {
return trit(bagValue, current.bitID) > (current.id - 1) ? -1 : -2;
} else {
newState.id = 0;
newState.bitID = current.bitID;
newState.bitID--;
if (newState.bitID < 0) return -1; // shouldnt happen
}
}
return stateID[newState.bitID][newState.id];
}
vi used(32, 0);
int test_strategy(int a, int b, vvi &strategy) {
int cnum = 0;
for (int i = 0; i < 500; i++) {
used[cnum]++;
int choose = strategy[cnum][0];
if (strategy[cnum][choose ? b : a] < 0) {
return -strategy[cnum][choose ? b : a];
} else {
cnum = strategy[cnum][choose ? b : a];
}
}
return 0;
}
vvi devise_strategy(int N) {
vvi s(STATES, vi(N+1, 0));
build_state_table();
for (int i = 0; i < STATES; i++) {
State cur = states[i];
s[i][0] = cur.id == 0 ? 0 : 1;
for (int j = 1; j <= N; j++) {
s[i][j] = transition(cur, j);
// if (s[i][j] >= 0) used[s[i][j]]++;
}
}
// for (int i = 1; i <= 5000; i++) {
// for (int j = 1; j <= 5000; j++) {
// if (i == j) continue;
// int ans = test_strategy(i, j, s);
// if (ans != (i < j ? 1 : 2)) {
// cout << "Fails on: " << i << ", " << j << "\n";
// break;
// }
// }
// }
// vi bits(STATES); iota(bits.begin(), bits.end(), 0);
// sort(bits.begin(), bits.end(), [&](const auto &a, const auto &b) {return used[a] < used[b];});
// for (int i = 0; i < STATES; i++) {
// cout << "Used state " << bits[i] << " " << used[bits[i]] << " times\n";
// }
return s;
}
