#include "mushrooms.h"
#include <bits/stdc++.h>
using namespace std;
#ifdef LOCAL
#define dbg(x) (cout << (x))
#else
#define dbg(x)
#endif
using vi =vector<int>;
#define pb push_back
#define FOR(i, a, b) for (int i = a; i < b; i++)
// determine the identities of the first k-1 mushrooms
void determine(int n, int k, int &cur_pos, int &tl_light, set<int> &light, set<int> &dark) {
int target = (k - 1) / 2;
light.insert(0);
(use_machine(vi({0, 1})) ? dark : light).insert(1);
(use_machine(vi({0, 2})) ? dark : light).insert(2);
bool using_light = light.size() > dark.size();
vi check; for (auto &el : (using_light ? light : dark)) check.push_back(el);
cur_pos = 3;
while (!(light.size() >= target || dark.size() >= target)) {
// run with config a0b1
int ans = use_machine(vi({cur_pos, check[0], cur_pos + 1, check[1]}));
bool a_same = ans & 1, b_same = ans & 2; // answer ranges from 0-3. if ans = 2|3 then b is set. if ans = 1|3 a is set
bool a_light = a_same != using_light, b_light = b_same != using_light;
(a_light ? light : dark).insert(cur_pos);
(b_light ? light : dark).insert(cur_pos+1);
cur_pos += 2;
}
tl_light += light.size();
}
void run_sieve(int n, int &cur_pos, int &tl_light, set<int> &light, set<int> &dark) {
int filter_size = max(light.size(), dark.size());
bool using_light = light.size() > dark.size();
filter_size = min(filter_size, n - cur_pos);
vi filter(2*filter_size);
int start = cur_pos;
auto it = (using_light ? light : dark).begin();
for (int i = 0; i < filter.size(); i++) {
if (i % 2 == 0) {
filter[i] = cur_pos++;
} else {
filter[i] = *it;
++it;
}
}
int ans = use_machine(filter);
bool start_light = (using_light && ans % 2 == 0) || (!using_light && ans % 2 == 1);
(start_light ? light : dark).insert(start);
int amount = (ans + 1) / 2;
if (using_light) amount = filter_size - amount;
tl_light += amount;
}
int count_mushrooms(int n) {
if (n <= 227) {
int tl_a = 1;
for (int i = 1; i < n; i++) {
tl_a += 1 - use_machine({0, i});
}
return tl_a;
}
int cur_pos = 0, tl_light = 0;
set<int> light, dark;
determine(n, 146, cur_pos, tl_light, light, dark);
while (cur_pos < n) {
run_sieve(n, cur_pos, tl_light, light, dark);
}
return tl_light;
}
