g++ -O2 sheep.cpp

sheep.cpp

#include <bits/stdc++.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/select.h>
#include <sys/stat.h>
#include <termios.h>
#include <unistd.h>
using namespace std;

struct MmapFile {
  int fd{-1};
  size_t size{0};
  const unsigned char* data{nullptr};
};

static MmapFile g_wolf_endgame_db;
static MmapFile g_sheep_endgame_db;

bool mmap_file_if_exists(const char* path, MmapFile& out) {
  int fd = open(path, O_RDONLY);
  if (fd < 0) return false;
  struct stat st{};
  if (fstat(fd, &st) != 0 || st.st_size <= 0) {
    close(fd);
    return false;
  }
  void* ptr = mmap(nullptr, static_cast<size_t>(st.st_size), PROT_READ, MAP_PRIVATE, fd, 0);
  if (ptr == MAP_FAILED) {
    close(fd);
    return false;
  }
  out.fd = fd;
  out.size = static_cast<size_t>(st.st_size);
  out.data = static_cast<const unsigned char*>(ptr);
  return true;
}

void munmap_file(MmapFile& f) {
  if (f.data) {
    munmap(const_cast<unsigned char*>(f.data), f.size);
    f.data = nullptr;
  }
  if (f.fd >= 0) {
    close(f.fd);
    f.fd = -1;
  }
  f.size = 0;
}

void init_endgame_dbs() {
  static const char* wolf_paths[] = {"Documents/wolf-chess-10w.bin"};
  static const char* sheep_paths[] = {"Documents/wolf-chess-10s.bin"};
  for (const char* path : wolf_paths) {
    if (mmap_file_if_exists(path, g_wolf_endgame_db)) break;
  }
  for (const char* path : sheep_paths) {
    if (mmap_file_if_exists(path, g_sheep_endgame_db)) break;
  }
}

void cleanup_endgame_dbs() {
  munmap_file(g_wolf_endgame_db);
  munmap_file(g_sheep_endgame_db);
}

enum class Piece : char {
  EMPTY = '.',
  WOLF = 'W',
  SHEEP = 'S',
};

pair<int, int> dirs[] = {
    {0, 1},
    {0, -1},
    {1, 0},
    {-1, 0},
};

bool valid(int i, int j) { return i >= 0 && i < 5 && j >= 0 && j < 5; }

struct Position {
  Piece p[5][5];
  bool is_wolf_turn{true};

  Position() {
    for (int i = 0; i < 5; i++) {
      for (int j = 0; j < 5; j++) { p[i][j] = Piece::EMPTY; }
    }
    for (int i = 0; i < 3; i++) {
      for (int j = 0; j < 5; j++) { p[i][j] = Piece::SHEEP; }
    }
    for (int j = 1; j < 4; j++) { p[4][j] = Piece::WOLF; }
    is_wolf_turn = true;
  }
  Position(string_view s) {
    for (int i = 0; i < 5; i++) {
      for (int j = 0; j < 5; j++) { p[i][j] = Piece::EMPTY; }
    }
    is_wolf_turn = true;
    int k = 0;
    for (char ch : s) {
      if (ch == ' ' || ch == '\n' || ch == '\r' || ch == '\t') continue;
      Piece piece = Piece::EMPTY;
      if (ch == 'W' || ch == 'w')
        piece = Piece::WOLF;
      else if (ch == 'S' || ch == 's')
        piece = Piece::SHEEP;
      else if (ch == '.' || ch == '-')
        piece = Piece::EMPTY;
      else
        continue;
      if (k >= 25) break;
      p[k / 5][k % 5] = piece;
      ++k;
    }
  }
  Position(const Position&) = default;
  Position(Position&&) = default;
  Position& operator=(const Position&) = default;
  Position& operator=(Position&&) = default;
  bool operator==(const Position&) const = default;
  friend ostream& operator<<(ostream& os, const Position& pos) {
    for (int i = 0; i < 5; i++) {
      for (int j = 0; j < 5; j++) { os << static_cast<char>(pos.p[i][j]); }
      os << '\n';
    }
    os << (pos.is_wolf_turn ? "wolf" : "sheep") << " to move";
    return os;
  }

  size_t encode() const {
    int bit{0}, wi{0}, setbit{0};
    size_t a{0};
    int w[3];
    w[0] = w[1] = w[2] = 0;
    for (int i = 0; i < 5; ++i)
      for (int j = 0; j < 5; ++j)
        (p[i][j] != Piece::EMPTY and ++setbit), a |= (p[i][j] != Piece::EMPTY) << bit++,
            p[i][j] == Piece::WOLF and (w[wi++] = setbit - 1);
    return a | (size_t)(w[0] + w[1] * (w[1] - 1) / 2 + w[2] * (w[2] - 1) * (w[2] - 2) / 6)
                   << 25;
  }

  void all_indices(auto&& f) const {
    for (int i = 0; i < 5; i++) {
      for (int j = 0; j < 5; j++) { f(i, j); }
    }
  }

  bool can_select(int i, int j) const {
    if (!valid(i, j)) return false;
    if (is_wolf_turn) return p[i][j] == Piece::WOLF;
    return p[i][j] == Piece::SHEEP;
  }

  bool legal_move(int si, int sj, int ti, int tj) const {
    if (!valid(si, sj) || !valid(ti, tj)) return false;
    if (!can_select(si, sj)) return false;
    int di = ti - si;
    int dj = tj - sj;
    int mi = (si + ti) / 2;
    int mj = (sj + tj) / 2;
    if (is_wolf_turn) {
      if (abs(di) + abs(dj) == 1) return p[ti][tj] == Piece::EMPTY;
      if ((abs(di) == 2 && dj == 0) || (abs(dj) == 2 && di == 0)) {
        return p[ti][tj] == Piece::SHEEP and p[mi][mj] == Piece::EMPTY;
      }
      return false;
    }
    return abs(di) + abs(dj) == 1 && p[ti][tj] == Piece::EMPTY;
  }

  bool apply_move(int si, int sj, int ti, int tj) {
    if (!legal_move(si, sj, ti, tj)) return false;
    Piece piece = p[si][sj];
    p[si][sj] = Piece::EMPTY;
    p[ti][tj] = piece;
    is_wolf_turn = !is_wolf_turn;
    return true;
  }

  void all_moves(auto&& f) const {
    all_indices([&](int i, int j) {
      if (!can_select(i, j)) return;
      for (auto [di, dj] : dirs) {
        int ni = i + di;
        int nj = j + dj;
        auto new_pos = *this;
        if (new_pos.apply_move(i, j, ni, nj)) f(new_pos);
      }
      if (is_wolf_turn) {
        for (auto [di, dj] : dirs) {
          int ni = i + 2 * di;
          int nj = j + 2 * dj;
          auto new_pos = *this;
          if (new_pos.apply_move(i, j, ni, nj)) f(new_pos);
        }
      }
    });
  }

  bool wolf_won() const {
    int sum = 0;
    all_indices([&](int i, int j) {
      if (p[i][j] == Piece::SHEEP) { sum++; }
    });
    return sum <= 3;
  }

  bool sheep_won() const {
    if (is_wolf_turn) {
      int sum = 0;
      all_moves([&](const Position&) { ++sum; });
      return sum == 0;
    } else {
      return false;
    }
  }
};

struct RawModeGuard {
  termios old_tio{};
  bool active{false};
  RawModeGuard() {
    if (tcgetattr(STDIN_FILENO, &old_tio) == 0) {
      termios raw = old_tio;
      raw.c_lflag &= static_cast<unsigned long>(~(ICANON | ECHO));
      raw.c_cc[VMIN] = 1;
      raw.c_cc[VTIME] = 0;
      if (tcsetattr(STDIN_FILENO, TCSANOW, &raw) == 0) active = true;
    }
  }
  ~RawModeGuard() {
    if (active) tcsetattr(STDIN_FILENO, TCSANOW, &old_tio);
  }
};

struct Move {
  int si, sj, ti, tj;
};

struct SearchNode {
  Position pos;
  vector<pair<Move, unique_ptr<SearchNode>>> children;
  vector<Move> ordered_moves;
  optional<Move> best_move;
};

struct SearchResult {
  optional<Move> move;
  int completed_depth{0};
  int best_score{0};
  bool timed_out{false};
};

bool same_move(const Move& a, const Move& b) {
  return a.si == b.si && a.sj == b.sj && a.ti == b.ti && a.tj == b.tj;
}

struct PositionHash {
  size_t operator()(const Position& pos) const {
    uint64_t h = 1469598103934665603ULL;
    for (int i = 0; i < 5; i++) {
      for (int j = 0; j < 5; j++) {
        h ^= static_cast<uint64_t>(static_cast<unsigned char>(pos.p[i][j]));
        h *= 1099511628211ULL;
      }
    }
    h ^= static_cast<uint64_t>(pos.is_wolf_turn);
    h *= 1099511628211ULL;
    return static_cast<size_t>(h);
  }
};

struct TTEntry {
  int depth;
  int score;
  int bound;
  optional<Move> best_move;
};

constexpr int INF_SCORE = 1'000'000'000;
constexpr int MATE_SCORE_BASE = INF_SCORE - 100000;

int score_to_tt(int score, int ply_from_root) {
  if (score >= MATE_SCORE_BASE) return score + ply_from_root;
  if (score <= -MATE_SCORE_BASE) return score - ply_from_root;
  return score;
}

int score_from_tt(int score, int ply_from_root) {
  if (score >= MATE_SCORE_BASE) return score - ply_from_root;
  if (score <= -MATE_SCORE_BASE) return score + ply_from_root;
  return score;
}

struct EdgeKey {
  Position from;
  Position to;
  bool operator==(const EdgeKey&) const = default;
};

struct EdgeKeyHash {
  size_t operator()(const EdgeKey& e) const {
    size_t a = PositionHash{}(e.from);
    size_t b = PositionHash{}(e.to);
    return a ^ (b + 0x9e3779b97f4a7c15ULL + (a << 6) + (a >> 2));
  }
};

uint64_t edge_hash64(const EdgeKey& e) {
  uint64_t a = static_cast<uint64_t>(PositionHash{}(e.from));
  uint64_t b = static_cast<uint64_t>(PositionHash{}(e.to));
  uint64_t x = a ^ (b + 0x9e3779b97f4a7c15ULL + (a << 6) + (a >> 2));
  x ^= (x >> 30);
  x *= 0xbf58476d1ce4e5b9ULL;
  x ^= (x >> 27);
  x *= 0x94d049bb133111ebULL;
  x ^= (x >> 31);
  return x;
}

int repeat_loss_eval(bool mover_is_wolf, int ply_from_root_after_move) {
  if (mover_is_wolf) return -INF_SCORE + ply_from_root_after_move;
  return INF_SCORE - ply_from_root_after_move;
}

struct SearchKey {
  Position pos;
  uint64_t history_sig;
  bool operator==(const SearchKey&) const = default;
};

struct SearchKeyHash {
  size_t operator()(const SearchKey& k) const {
    size_t a = PositionHash{}(k.pos);
    size_t b = std::hash<uint64_t>{}(k.history_sig);
    return a ^ (b + 0x9e3779b97f4a7c15ULL + (a << 6) + (a >> 2));
  }
};

vector<Move> legal_moves(const Position& pos) {
  vector<Move> moves;
  moves.reserve(64);
  pos.all_indices([&](int i, int j) {
    if (!pos.can_select(i, j)) return;
    for (auto [di, dj] : dirs) {
      int ni = i + di;
      int nj = j + dj;
      if (pos.legal_move(i, j, ni, nj)) moves.push_back({i, j, ni, nj});
    }
    if (pos.is_wolf_turn) {
      for (auto [di, dj] : dirs) {
        int ni = i + 2 * di;
        int nj = j + 2 * dj;
        if (pos.legal_move(i, j, ni, nj)) moves.push_back({i, j, ni, nj});
      }
    }
  });
  return moves;
}

int sheep_count(const Position& pos) {
  int cnt = 0;
  pos.all_indices([&](int i, int j) {
    if (pos.p[i][j] == Piece::SHEEP) ++cnt;
  });
  return cnt;
}

int move_count_for(const Position& pos, bool wolf_turn) {
  int cnt = 0;
  for (int i = 0; i < 5; i++) {
    for (int j = 0; j < 5; j++) {
      if (wolf_turn) {
        if (pos.p[i][j] != Piece::WOLF) continue;
      } else {
        if (pos.p[i][j] != Piece::SHEEP) continue;
      }
      for (auto [di, dj] : dirs) {
        int ni = i + di;
        int nj = j + dj;
        if (valid(ni, nj) && pos.p[ni][nj] == Piece::EMPTY) ++cnt;
      }
      if (wolf_turn) {
        for (auto [di, dj] : dirs) {
          int ni = i + 2 * di;
          int nj = j + 2 * dj;
          int mi = i + di;
          int mj = j + dj;
          if (!valid(ni, nj) || !valid(mi, mj)) continue;
          if (pos.p[ni][nj] == Piece::SHEEP && pos.p[mi][mj] == Piece::EMPTY) ++cnt;
        }
      }
    }
  }
  return cnt;
}

int static_eval(const Position& pos) {
  if (pos.wolf_won()) return INF_SCORE;
  if (pos.sheep_won()) return -INF_SCORE;
  size_t baddr = pos.encode();
  const auto& db = (pos.is_wolf_turn ? g_wolf_endgame_db : g_sheep_endgame_db);
  if (db.data and db.size > baddr) {
    char result = db.data[baddr >> 2] >> ((baddr & 3) * 2) & 3;
    if (result == 1) return INF_SCORE;
    if (result == 2) return -INF_SCORE;
  }
  int sheep_left = sheep_count(pos);
  int wolf_mob = move_count_for(pos, true);
  int sheep_mob = move_count_for(pos, false);
  // if (pos.is_wolf_turn) return -sheep_left * 100 + wolf_mob * 50;
  return -sheep_left * 100 + wolf_mob * 50;
}

int terminal_eval(const Position& pos, int ply_from_root) {
  if (pos.wolf_won()) return INF_SCORE - ply_from_root;
  if (pos.sheep_won()) return -INF_SCORE + ply_from_root;
  return static_eval(pos);
}

string format_score(int score) {
  if (score >= MATE_SCORE_BASE) return "#" + to_string(INF_SCORE - score);
  if (score <= -MATE_SCORE_BASE) return "#-" + to_string(INF_SCORE + score);
  return to_string(score);
}

bool terminal_for_side(bool wolf_turn, int score) {
  return wolf_turn ? score >= MATE_SCORE_BASE : score <= -MATE_SCORE_BASE;
}

bool update_cursor(char ch, int& ci, int& cj) {
  if (ch == 'h')
    cj = max(0, cj - 1);
  else if (ch == 'l')
    cj = min(4, cj + 1);
  else if (ch == 'k')
    ci = max(0, ci - 1);
  else if (ch == 'j')
    ci = min(4, ci + 1);
  else
    return false;
  return true;
}

bool apply_move_checked(Position& pos, const Move& mv) {
  return pos.apply_move(mv.si, mv.sj, mv.ti, mv.tj);
}

vector<Move> reorder_moves(const vector<Move>& ordered, vector<pair<Move, int>>& scored,
                           bool wolf_turn) {
  sort(scored.begin(), scored.end(), [&](const auto& a, const auto& b) {
    return wolf_turn ? a.second > b.second : a.second < b.second;
  });
  vector<Move> next;
  next.reserve(ordered.size());
  for (const auto& [mv, _] : scored) next.push_back(mv);
  for (const auto& mv : ordered) {
    bool found = false;
    for (const auto& [smv, _] : scored) {
      if (same_move(mv, smv)) {
        found = true;
        break;
      }
    }
    if (!found) next.push_back(mv);
  }
  return next;
}

void apply_hint_order(vector<Move>& moves, const vector<Move>& hint) {
  if (hint.empty()) return;
  vector<Move> next;
  next.reserve(moves.size());
  for (const auto& hm : hint) {
    for (const auto& mv : moves) {
      if (same_move(hm, mv)) {
        next.push_back(mv);
        break;
      }
    }
  }
  for (const auto& mv : moves) {
    bool found = false;
    for (const auto& nm : next) {
      if (same_move(mv, nm)) {
        found = true;
        break;
      }
    }
    if (!found) next.push_back(mv);
  }
  moves = move(next);
}

SearchNode* get_or_create_child(SearchNode* node, const Move& mv, const Position& child_pos) {
  if (!node) return nullptr;
  for (auto& [cmv, ch] : node->children) {
    if (same_move(cmv, mv)) {
      if (!ch) {
        ch = make_unique<SearchNode>();
        ch->pos = child_pos;
      } else {
        ch->pos = child_pos;
      }
      return ch.get();
    }
  }
  auto ch = make_unique<SearchNode>();
  ch->pos = child_pos;
  auto* ptr = ch.get();
  node->children.push_back({mv, move(ch)});
  return ptr;
}

int search_score(const Position& pos, int depth, int alpha, int beta,
                 unordered_map<SearchKey, TTEntry, SearchKeyHash>& tt, SearchNode* node,
                 const function<bool()>& should_stop, bool& stopped, int ply_from_root = 0) {
  if (should_stop()) {
    stopped = true;
    return static_eval(pos);
  }
  if (depth == 0 || pos.wolf_won() || pos.sheep_won())
    return terminal_eval(pos, ply_from_root);
  const int alpha_orig = alpha;
  const int beta_orig = beta;
  optional<Move> tt_move;
  SearchKey key{pos, 0};
  auto it = tt.find(key);
  if (it != tt.end()) {
    int tt_score = score_from_tt(it->second.score, ply_from_root);
    if (it->second.depth >= depth) {
      if (it->second.bound == 0) return tt_score;
      if (it->second.bound == 1)
        alpha = max(alpha, tt_score);
      else
        beta = min(beta, tt_score);
      if (alpha >= beta) return tt_score;
    }
    tt_move = it->second.best_move;
  }
  auto moves = legal_moves(pos);
  if (moves.empty()) {
    int v = static_eval(pos);
    tt[key] = TTEntry{depth, v, 0, nullopt};
    return v;
  }
  if (node) apply_hint_order(moves, node->ordered_moves);
  vector<Move> ordered;
  ordered.reserve(moves.size());
  optional<Move> first_hint =
      tt_move.has_value() ? tt_move : (node ? node->best_move : nullopt);
  if (first_hint.has_value()) {
    for (const auto& mv : moves) {
      if (same_move(mv, *first_hint)) {
        ordered.push_back(mv);
        break;
      }
    }
  }
  for (const auto& mv : moves) {
    if (first_hint.has_value() && same_move(mv, *first_hint)) continue;
    if (abs(mv.ti - mv.si) == 2 || abs(mv.tj - mv.sj) == 2) ordered.push_back(mv);
  }
  for (const auto& mv : moves) {
    if (first_hint.has_value() && same_move(mv, *first_hint)) continue;
    if (abs(mv.ti - mv.si) != 2 && abs(mv.tj - mv.sj) != 2) ordered.push_back(mv);
  }
  const int INF = 1'000'000'000;
  int best = pos.is_wolf_turn ? -INF : INF;
  optional<Move> best_move;
  for (const auto& mv : ordered) {
    if (should_stop()) {
      stopped = true;
      return static_eval(pos);
    }
    Position np = pos;
    np.apply_move(mv.si, mv.sj, mv.ti, mv.tj);
    SearchNode* child = get_or_create_child(node, mv, np);
    int v = search_score(
        np, depth - 1, alpha, beta, tt, child, should_stop, stopped, ply_from_root + 1);
    if (stopped) return static_eval(pos);
    if (pos.is_wolf_turn) {
      if (v > best) {
        best = v;
        best_move = mv;
      }
      alpha = max(alpha, best);
      if (alpha >= beta) break;
    } else {
      if (v < best) {
        best = v;
        best_move = mv;
      }
      beta = min(beta, best);
      if (alpha >= beta) break;
    }
  }
  int bound = 0;
  if (best <= alpha_orig)
    bound = 2;
  else if (best >= beta_orig)
    bound = 1;
  tt[key] = TTEntry{depth, score_to_tt(best, ply_from_root), bound, best_move};
  if (node) {
    vector<pair<Move, int>> rescored;
    rescored.reserve(ordered.size());
    for (const auto& mv : ordered) {
      Position np = pos;
      if (!np.apply_move(mv.si, mv.sj, mv.ti, mv.tj)) continue;
      SearchKey ck{np, 0};
      auto ct = tt.find(ck);
      if (ct == tt.end()) continue;
      rescored.push_back({mv, score_from_tt(ct->second.score, ply_from_root + 1)});
    }
    node->ordered_moves = reorder_moves(ordered, rescored, pos.is_wolf_turn);
    node->best_move = best_move;
  }
  return best;
}

SearchResult choose_ai_move(const Position& root, int max_depth,
                            unordered_map<SearchKey, TTEntry, SearchKeyHash>& tt,
                            SearchNode* root_node,
                            const function<void(int, int, optional<Move>)>& on_depth) {
  auto moves = legal_moves(root);
  if (moves.empty()) return {};
  static mt19937 rng(
      static_cast<unsigned int>(chrono::steady_clock::now().time_since_epoch().count()));
  auto deadline = chrono::steady_clock::now() + chrono::milliseconds(3000);
  vector<Move> ordered = moves;
  if (root_node) apply_hint_order(ordered, root_node->ordered_moves);
  optional<Move> best = ordered[0];
  int completed_depth = 0;
  int completed_score = 0;
  bool timed_out_any = false;
  if (tt.empty()) tt.reserve(200000);
  for (int d = 1; d <= max_depth; d++) {
    bool timed_out = false;
    int best_val = root.is_wolf_turn ? -INF_SCORE : INF_SCORE;
    vector<Move> best_moves;
    vector<pair<Move, int>> scored;
    for (const auto& mv : ordered) {
      if (chrono::steady_clock::now() >= deadline) {
        timed_out = true;
        break;
      }
      Position np = root;
      np.apply_move(mv.si, mv.sj, mv.ti, mv.tj);
      auto should_stop = [&]() { return chrono::steady_clock::now() >= deadline; };
      SearchNode* child = get_or_create_child(root_node, mv, np);
      int v =
          search_score(np, d - 1, -INF_SCORE, INF_SCORE, tt, child, should_stop, timed_out, 1);
      if (timed_out) break;
      scored.push_back({mv, v});
      if (root.is_wolf_turn) {
        if (v > best_val) {
          best_val = v;
          best_moves.clear();
          best_moves.push_back(mv);
        } else if (v == best_val) {
          best_moves.push_back(mv);
        }
      } else {
        if (v < best_val) {
          best_val = v;
          best_moves.clear();
          best_moves.push_back(mv);
        } else if (v == best_val) {
          best_moves.push_back(mv);
        }
      }
      if (terminal_for_side(root.is_wolf_turn, best_val)) break;
    }
    if (timed_out) {
      timed_out_any = true;
      break;
    }
    if (!best_moves.empty()) {
      uniform_int_distribution<int> dist(0, static_cast<int>(best_moves.size()) - 1);
      best = best_moves[dist(rng)];
    }
    completed_depth = d;
    completed_score = best_val;
    if (on_depth) on_depth(completed_depth, completed_score, best);
    if (terminal_for_side(root.is_wolf_turn, completed_score)) break;
    ordered = reorder_moves(ordered, scored, root.is_wolf_turn);
    if (root_node) {
      root_node->ordered_moves = ordered;
      root_node->best_move = best;
    }
  }
  return {best, completed_depth, completed_score, timed_out_any};
}

void analyze_live(const Position& root, unordered_map<SearchKey, TTEntry, SearchKeyHash>& tt,
                  SearchNode* root_node, const function<bool()>& should_stop,
                  const function<void(int, int, optional<Move>)>& on_depth) {
  auto moves = legal_moves(root);
  if (moves.empty()) return;
  vector<Move> ordered = moves;
  if (root_node) apply_hint_order(ordered, root_node->ordered_moves);
  if (tt.empty()) tt.reserve(200000);
  for (int d = 1;; d++) {
    if (should_stop()) return;
    bool stopped = false;
    int best_val = root.is_wolf_turn ? -INF_SCORE : INF_SCORE;
    vector<pair<Move, int>> scored;
    for (const auto& mv : ordered) {
      if (should_stop()) return;
      Position np = root;
      np.apply_move(mv.si, mv.sj, mv.ti, mv.tj);
      SearchNode* child = get_or_create_child(root_node, mv, np);
      int v =
          search_score(np, d - 1, -INF_SCORE, INF_SCORE, tt, child, should_stop, stopped, 1);
      if (stopped) return;
      scored.push_back({mv, v});
      if (root.is_wolf_turn) {
        best_val = max(best_val, v);
      } else {
        best_val = min(best_val, v);
      }
      if (terminal_for_side(root.is_wolf_turn, best_val)) break;
    }
    if (!scored.empty()) on_depth(d, best_val, nullopt);
    if (terminal_for_side(root.is_wolf_turn, best_val)) return;
    ordered = reorder_moves(ordered, scored, root.is_wolf_turn);
    if (root_node) root_node->ordered_moves = ordered;
  }
}

struct AiState {
  atomic<unsigned long long> current_job{0};
  atomic<unsigned long long> running_job{0};
  atomic<int> running_task{0};
  mutex m;
  bool has_result{false};
  int result_task{0};
  bool has_progress{false};
  int progress_depth{0};
  int progress_score{0};
  optional<Move> progress_move;
  SearchResult result;
  mutex tt_m;
  unordered_map<SearchKey, TTEntry, SearchKeyHash> shared_tt;
  mutex tree_m;
  unique_ptr<SearchNode> root;
};

unique_ptr<SearchNode> make_root_node(const Position& pos) {
  auto n = make_unique<SearchNode>();
  n->pos = pos;
  return n;
}

void ensure_ai_root(AiState& ai, const Position& pos) {
  lock_guard<mutex> lk(ai.tree_m);
  if (!ai.root || !(ai.root->pos == pos)) ai.root = make_root_node(pos);
}

void reroot_ai_tree(AiState& ai, const Position& before, const Move& mv,
                    const Position& after) {
  lock_guard<mutex> lk(ai.tree_m);
  if (!ai.root || !(ai.root->pos == before)) {
    ai.root = make_root_node(after);
    return;
  }
  for (auto& [cmv, ch] : ai.root->children) {
    if (same_move(cmv, mv) && ch && ch->pos == after) {
      ai.root = move(ch);
      return;
    }
  }
  ai.root = make_root_node(after);
}

void invalidate_ai(AiState& ai) {
  ai.current_job.fetch_add(1);
  ai.running_job.store(0);
  ai.running_task.store(0);
  lock_guard<mutex> lk(ai.m);
  ai.has_result = false;
  ai.result_task = 0;
  ai.has_progress = false;
  ai.progress_depth = 0;
  ai.progress_score = 0;
  ai.progress_move.reset();
  ai.result = {};
}

struct TaskResult {
  int task;
  SearchResult res;
};

struct ProgressInfo {
  int depth;
  int score;
  optional<Move> move;
};

optional<TaskResult> take_ai_result(AiState& ai) {
  lock_guard<mutex> lk(ai.m);
  if (!ai.has_result) return nullopt;
  ai.has_result = false;
  return TaskResult{ai.result_task, ai.result};
}

optional<ProgressInfo> get_ai_progress(AiState& ai) {
  lock_guard<mutex> lk(ai.m);
  if (!ai.has_progress) return nullopt;
  return ProgressInfo{ai.progress_depth, ai.progress_score, ai.progress_move};
}

void start_ai_search(AiState& ai, Position pos, int depth, int task) {
  unsigned long long job = ai.current_job.fetch_add(1) + 1;
  ai.running_job.store(job);
  ai.running_task.store(task);
  {
    lock_guard<mutex> lk(ai.m);
    ai.has_result = false;
    ai.result_task = task;
    ai.has_progress = false;
    ai.progress_depth = 0;
    ai.progress_score = 0;
    ai.progress_move.reset();
    ai.result = {};
  }
  thread([&ai, pos, depth, job, task]() {
    auto should_stop = [&]() { return ai.running_job.load() != job; };
    unordered_map<SearchKey, TTEntry, SearchKeyHash> tt;
    unique_ptr<SearchNode> root;
    {
      lock_guard<mutex> lk(ai.tt_m);
      tt = move(ai.shared_tt);
    }
    {
      lock_guard<mutex> lk(ai.tree_m);
      if (!ai.root || !(ai.root->pos == pos)) ai.root = make_root_node(pos);
      root = move(ai.root);
    }
    if (tt.empty()) tt.reserve(200000);
    auto publish_progress = [&](int d, int s, optional<Move> mv) {
      lock_guard<mutex> lk(ai.m);
      if (ai.running_job.load() != job) return;
      ai.has_progress = true;
      ai.progress_depth = d;
      ai.progress_score = s;
      ai.progress_move = mv;
    };
    if (task == 1 || task == 3) {
      auto res = choose_ai_move(pos, depth, tt, root.get(), publish_progress);
      lock_guard<mutex> lk(ai.m);
      if (ai.running_job.load() != job) return;
      ai.result_task = task;
      ai.result = res;
      ai.has_result = true;
    } else {
      analyze_live(pos, tt, root.get(), should_stop, publish_progress);
    }
    {
      lock_guard<mutex> lk(ai.tt_m);
      ai.shared_tt = move(tt);
    }
    {
      lock_guard<mutex> lk(ai.tree_m);
      ai.root = move(root);
    }
    ai.running_job.store(0);
    ai.running_task.store(0);
  }).detach();
}

optional<char> read_key(int timeout_ms) {
  fd_set set;
  FD_ZERO(&set);
  FD_SET(STDIN_FILENO, &set);
  timeval tv{};
  timeval* ptv = nullptr;
  if (timeout_ms >= 0) {
    tv.tv_sec = timeout_ms / 1000;
    tv.tv_usec = (timeout_ms % 1000) * 1000;
    ptv = &tv;
  }
  int rv = select(STDIN_FILENO + 1, &set, nullptr, nullptr, ptv);
  if (rv <= 0) return nullopt;
  char ch = 0;
  if (read(STDIN_FILENO, &ch, 1) == 1) return ch;
  return nullopt;
}

void render(const Position& pos, int ci, int cj, optional<pair<int, int>> selected,
            optional<Move> last_move, optional<Move> best_move_hint, bool show_best_move,
            const string& msg, int mode, bool human_is_wolf, int last_ai_depth,
            int last_ai_score) {
  cout << "\x1b[2J\x1b[H";
  cout << "h/j/k/l move  space|enter select/move  esc unselect  1/2/3/4 mode  q quit\n";
  if (mode == 0) cout << "self play (E/e live eval, F/f best move)\n";
  if (mode == 1)
    cout << "vs ai: you are " << (human_is_wolf ? "wolf" : "sheep")
         << "  (S/s switch side + reset, E/e eval)\n";
  if (mode == 2) {
    cout << "ai vs ai watch mode (F/f best move)\n";
    cout << "last depth: " << last_ai_depth << "  eval: " << format_score(last_ai_score)
         << '\n';
  }
  if (mode == 3)
    cout << "board editor: w wolf, s sheep, space empty, c clear, tab side, S start self "
            "play, a start vs ai\n";
  cout << (pos.is_wolf_turn ? "wolf" : "sheep") << " to move\n";
  for (int i = 0; i < 5; i++) {
    for (int j = 0; j < 5; j++) {
      bool is_cursor = (i == ci && j == cj);
      bool is_sel = selected && selected->first == i && selected->second == j;
      bool is_from = last_move && last_move->si == i && last_move->sj == j;
      bool is_last = last_move && last_move->ti == i && last_move->tj == j;
      bool is_best = show_best_move && best_move_hint &&
                     ((best_move_hint->si == i && best_move_hint->sj == j) ||
                      (best_move_hint->ti == i && best_move_hint->tj == j));
      Piece pc = pos.p[i][j];
      int fg = (pc == Piece::WOLF) ? 31 : 37;
      int bg = -1;
      if (is_best) bg = 102;
      if (is_last || is_from) bg = 43;
      if (is_cursor) bg = 44;
      if (is_sel) bg = 42;
      if (is_cursor && is_sel) bg = 45;
      cout << "\x1b[" << fg;
      if (bg != -1) cout << ';' << bg;
      cout << "m " << static_cast<char>(pc) << " \x1b[0m";
    }
    cout << '\n';
  }
  string result;
  if (pos.wolf_won())
    result = "wolves win";
  else if (pos.sheep_won())
    result = "sheep win";
  cout << "result: " << result << '\n';
  cout << "message: " << msg << '\n';
  cout.flush();
}

optional<int> title_screen() {
  auto s = R"(
- - - S W
- - - - W
- - W S S
- - S S S
- - S - -)";
  auto s1 = R"(
- - S - -
- S - - S
- - - S -
- - - W -
- - W - W)";
  auto se = R"(
- - - S -
- - - S W
- - S W -
- - - S -
- - S S W)";
  Position p(s);
  Position p1(s1);
  Position pe(se);
  pe.is_wolf_turn = p1.is_wolf_turn = p.is_wolf_turn = false;
  while (true) {
    cout << "\x1b[2J\x1b[H";
    cout << p.encode() << "," << p1.encode() << "," << pe.encode() << '\n';
    cout << static_eval(p) << "," << static_eval(p1) << '\n';
    cout << "wolf sheep\n";
    cout << "mmap dbs: wolf=" << (g_wolf_endgame_db.data ? "yes" : "no")
         << " sheep=" << (g_sheep_endgame_db.data ? "yes" : "no") << '\n';
    cout << "1) play vs self\n";
    cout << "2) play vs ai\n";
    cout << "3) watch ai vs ai\n";
    cout << "4) board editor\n";
    cout << "q) quit\n";
    cout.flush();
    char ch = 0;
    if (read(STDIN_FILENO, &ch, 1) != 1) return nullopt;
    if (ch == '1') return 0;
    if (ch == '2') return 1;
    if (ch == '3') return 2;
    if (ch == '4') return 3;
    if (ch == 'q') return nullopt;
  }
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  init_endgame_dbs();
  atexit(cleanup_endgame_dbs);

  RawModeGuard guard;
  auto mode = title_screen();
  if (!mode.has_value()) return 0;

  Position pos;
  int mode_id = *mode;
  bool human_is_wolf = true;
  const int ai_depth = 19;
  bool live_eval = false;
  bool show_best_move = false;
  AiState ai;
  int ci = 0, cj = 0;
  optional<pair<int, int>> selected;
  optional<Move> last_move;
  optional<Move> best_move_hint;
  int last_ai_depth = 0;
  int last_ai_score = 0;
  string msg;
  auto reset_game = [&](bool reset_side = false) {
    pos = Position();
    ci = 0;
    cj = 0;
    selected.reset();
    last_move.reset();
    last_ai_depth = 0;
    last_ai_score = 0;
    live_eval = false;
    show_best_move = false;
    best_move_hint.reset();
    invalidate_ai(ai);
    if (reset_side) human_is_wolf = !human_is_wolf;
  };
  auto apply_ai_move_now = [&](const Move& mv, int d, int s) {
    Position before = pos;
    apply_move_checked(pos, mv);
    reroot_ai_tree(ai, before, mv, pos);
    last_move = mv;
    last_ai_depth = d;
    last_ai_score = s;
    msg = "ai moved (depth " + to_string(d) + " score " + format_score(s) + ")";
  };
  auto start_from_editor = [&](int new_mode) {
    mode_id = new_mode;
    selected.reset();
    last_move.reset();
    live_eval = false;
    show_best_move = false;
    best_move_hint.reset();
    invalidate_ai(ai);
    if (new_mode == 1) human_is_wolf = pos.is_wolf_turn;
  };

  while (true) {
    bool over = pos.wolf_won() || pos.sheep_won();
    bool ai_turn = mode_id == 2 ? !over
                                : mode_id == 1 && !over &&
                                      ((pos.is_wolf_turn && !human_is_wolf) ||
                                       (!pos.is_wolf_turn && human_is_wolf));
    auto progress = get_ai_progress(ai);
    if (progress.has_value()) {
      last_ai_depth = progress->depth;
      last_ai_score = progress->score;
      if ((mode_id == 0 || mode_id == 2) && show_best_move && progress->move.has_value())
        best_move_hint = progress->move;
      if (live_eval && (mode_id == 0 || mode_id == 1)) {
        msg = "live eval depth " + to_string(progress->depth) + " score " +
              format_score(progress->score);
      }
    }
    if (ai_turn) {
      if (ai.running_task.load() == 2) invalidate_ai(ai);
      auto res = take_ai_result(ai);
      if (res.has_value()) {
        if (res->task == 1 && res->res.move.has_value()) {
          auto mv = *res->res.move;
          if (mode_id == 2 && show_best_move) best_move_hint = mv;
          apply_ai_move_now(mv, res->res.completed_depth, res->res.best_score);
        }
        selected.reset();
        continue;
      }
      if (ai.running_job.load() == 0) {
        ensure_ai_root(ai, pos);
        start_ai_search(ai, pos, ai_depth, 1);
      }
    } else {
      auto res = take_ai_result(ai);
      if (res.has_value() && res->task == 2) {
        last_ai_depth = res->res.completed_depth;
        last_ai_score = res->res.best_score;
        msg = "analysis depth " + to_string(res->res.completed_depth) + " score " +
              format_score(res->res.best_score);
      } else if (res.has_value() && res->task == 3) {
        last_ai_depth = res->res.completed_depth;
        last_ai_score = res->res.best_score;
        best_move_hint = res->res.move;
      }
      if ((live_eval || (mode_id == 0 && show_best_move)) && (mode_id == 0 || mode_id == 1) &&
          ai.running_job.load() == 0) {
        int task = (mode_id == 0 && show_best_move) ? 3 : 2;
        ensure_ai_root(ai, pos);
        start_ai_search(ai, pos, ai_depth, task);
      }
    }

    string display_msg = msg;
    if (live_eval && (mode_id == 0 || mode_id == 1)) {
      if (progress.has_value())
        display_msg = "live eval depth " + to_string(progress->depth) + " score " +
                      format_score(progress->score);
      else if (last_ai_depth > 0)
        display_msg = "live eval depth " + to_string(last_ai_depth) + " score " +
                      format_score(last_ai_score);
      else if (ai.running_job.load() != 0)
        display_msg = "live eval...";
    } else if (ai_turn && ai.running_job.load() != 0) {
      display_msg = "ai thinking...";
    }
    if (mode_id == 0 && show_best_move && display_msg.empty() && ai.running_job.load() != 0) {
      display_msg = "best move...";
    }
    render(pos,
           ci,
           cj,
           selected,
           last_move,
           best_move_hint,
           show_best_move,
           display_msg,
           mode_id,
           human_is_wolf,
           last_ai_depth,
           last_ai_score);
    msg.clear();
    auto ch_opt = read_key((ai_turn || ai.running_job.load() != 0) ? 50 : -1);
    if (!ch_opt.has_value()) continue;
    char ch = *ch_opt;
    if (ch == 'q') break;
    if (ch == '1' || ch == '2' || ch == '3' || ch == '4') {
      int new_mode = ch - '1';
      if (new_mode != mode_id) {
        bool keep_board = (new_mode == 3) ||
                          (mode_id == 3 && (new_mode == 0 || new_mode == 1)) ||
                          ((mode_id == 0 && new_mode == 1) || (mode_id == 1 && new_mode == 0));
        if (keep_board) {
          mode_id = new_mode;
          selected.reset();
          last_move.reset();
          live_eval = false;
          show_best_move = false;
          best_move_hint.reset();
          last_ai_depth = 0;
          last_ai_score = 0;
          invalidate_ai(ai);
        } else {
          mode_id = new_mode;
          reset_game();
        }
        if (mode_id == 1) human_is_wolf = pos.is_wolf_turn;
        if (mode_id == 0)
          msg = "switched to vs self";
        else if (mode_id == 1)
          msg = "switched to vs ai";
        else if (mode_id == 2)
          msg = "switched to ai vs ai";
        else
          msg = "switched to board editor";
      }
      continue;
    }
    if (mode_id == 1 && (ch == 'S' || ch == 's')) {
      reset_game(true);
      msg = "side switched and board reset";
      continue;
    }
    if ((mode_id == 0 || mode_id == 1) && (ch == 'E' || ch == 'e')) {
      live_eval = !live_eval;
      invalidate_ai(ai);
      if (live_eval) {
        int task = (mode_id == 0 && show_best_move) ? 3 : 2;
        ensure_ai_root(ai, pos);
        start_ai_search(ai, pos, ai_depth, task);
        msg = "live eval on";
      } else {
        msg = "live eval off";
      }
      continue;
    }
    if ((mode_id == 0 || mode_id == 2) && (ch == 'F' || ch == 'f')) {
      show_best_move = !show_best_move;
      best_move_hint.reset();
      if (mode_id == 0) { invalidate_ai(ai); }
      if (show_best_move && !over && mode_id == 0) {
        ensure_ai_root(ai, pos);
        start_ai_search(ai, pos, ai_depth, 3);
        msg = "best move on";
      } else {
        msg = "best move off";
      }
      continue;
    }
    if (ch == 27) selected.reset();
    if (ai_turn) {
      update_cursor(ch, ci, cj);
      continue;
    }
    if (mode_id == 3) {
      if (update_cursor(ch, ci, cj)) continue;
      if (ch == '\t') {
        pos.is_wolf_turn = !pos.is_wolf_turn;
        invalidate_ai(ai);
        best_move_hint.reset();
        last_move.reset();
        continue;
      }
      if (ch == 'w') {
        pos.p[ci][cj] = Piece::WOLF;
        invalidate_ai(ai);
        best_move_hint.reset();
        last_move.reset();
        continue;
      }
      if (ch == 's') {
        pos.p[ci][cj] = Piece::SHEEP;
        invalidate_ai(ai);
        best_move_hint.reset();
        last_move.reset();
        continue;
      }
      if (ch == ' ') {
        pos.p[ci][cj] = Piece::EMPTY;
        invalidate_ai(ai);
        best_move_hint.reset();
        last_move.reset();
        continue;
      }
      if (ch == 'c') {
        pos.all_indices([&](int i, int j) { pos.p[i][j] = Piece::EMPTY; });
        invalidate_ai(ai);
        best_move_hint.reset();
        last_move.reset();
        continue;
      }
      if (ch == 'S') {
        start_from_editor(0);
        msg = "started self play from editor";
        continue;
      }
      if (ch == 'a') {
        start_from_editor(1);
        msg = "started vs ai (you play first)";
        continue;
      }
      continue;
    }
    if (over) continue;
    if (!update_cursor(ch, ci, cj) && (ch == ' ' || ch == '\n' || ch == '\r')) {
      if (!selected) {
        if (pos.can_select(ci, cj))
          selected = {{ci, cj}};
        else
          msg = "cannot select that piece";
      } else if (selected->first == ci && selected->second == cj) {
        selected.reset();
      } else {
        auto [si, sj] = *selected;
        Move mv{si, sj, ci, cj};
        Position before = pos;
        if (apply_move_checked(pos, mv)) {
          reroot_ai_tree(ai, before, mv, pos);
          last_move = Move{si, sj, ci, cj};
          best_move_hint.reset();
          if ((mode_id == 0 || mode_id == 1) && (live_eval || show_best_move))
            invalidate_ai(ai);
          selected.reset();
        } else
          msg = "illegal move";
      }
    }
  }
  cout << "\x1b[0m\n";
}