/*
    This file is part of ctak.

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful, but
    WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
    General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with ctak. If not, see <https://www.gnu.org/licenses/>.
*/

#include "negamax.h"
#include "actions.h"
#ifdef RECORD_PV
#include "pv.h"
#endif

// ===================================================================
// Variables
// ===================================================================

const float infty = 3.0;
char negamax_ptn[9];
uint8_t negamax_search_depth = 3;
static uint64_t negamax_best_action;
static nn_function evaluate;

// ===================================================================
// check_win memoisation
// ===================================================================

#define WON_CACHE_SIZE 8192

static uint64_t won_cache_key[WON_CACHE_SIZE];
static uint8_t won_cache_val[WON_CACHE_SIZE];

static inline enum WIN_TYPE check_win_cached(tak_state_p state,
                                             const uint64_t hash) {
  const uint32_t idx = hash & (WON_CACHE_SIZE - 1);
  if (won_cache_key[idx] == hash)
    return (enum WIN_TYPE)won_cache_val[idx];
  const enum WIN_TYPE w = check_win(state);
  won_cache_key[idx] = hash;
  won_cache_val[idx] = (uint8_t)w;
  return w;
}

// ===================================================================
// leaf evaluation memoisation
// ===================================================================

typedef struct eval_memo_s {
  uint64_t key;
  nn_function eval;
  float value;
} eval_memo_t;

#define EVAL_MEMO_SIZE 262144
static eval_memo_t eval_memo[EVAL_MEMO_SIZE];

static inline float evaluate_memo(tak_state_p state, const uint64_t key) {
  const uint32_t midx = key & (EVAL_MEMO_SIZE - 1);
  if (eval_memo[midx].key == key && eval_memo[midx].eval == evaluate)
    return eval_memo[midx].value;
  const float value = evaluate(state);
  eval_memo[midx].key = key;
  eval_memo[midx].eval = evaluate;
  eval_memo[midx].value = value;
  return value;
}

// ===================================================================
// Helper declarations
// ===================================================================

static float negamax(tak_state_p state, const uint8_t cur_depth,
                     const uint8_t init_depth, float alpha, float beta,
                     const float colour, const uint64_t hash
#ifdef RECORD_PV
                     , pv_leaf_t *out_leaf
#endif
);

// ===================================================================
// Exported functions
// ===================================================================

void negamax_init(const uint8_t new_board_size) {
  actions_init(new_board_size);
  zobrist_init(new_board_size);
  tt_init();

  if (new_board_size == 5) {
    evaluate = &five_nn_fn;
  } else {
    evaluate = &six_nn_fn;
  }
}

void negamax_set_evaluator(const nn_function f) { evaluate = f; }

void negamax_free(void) { zobrist_free(); }

float negamax_generate(tak_state_p state) {
  // We need to start with something outside of [-∞,∞] because those values are
  // wins
  const float safe_infty = infty + 1;
  float result = -infty;

  negamax_best_action = -1;

  // The transposition table is deliberately not cleared here: entries persist
  // across plies. This is sound because the table is keyed on (board, side to
  // move) and an entry is only reusable at a depth no greater than the one it
  // was stored at. In particular the new root (the previous root's child,
  // stored at depth D-1) can never be reused at depth D, so every turn still
  // runs a full fresh search at the root.

  // this is constant across iterative deepening
  const uint64_t root_hash = zobrist_compute(state);
  state->hash = root_hash; // anchor the incremental hash at the root
#ifdef RECORD_PV
  pv_leaf_t iteration_leaf;
#endif
  for (int d = 1; d <= negamax_search_depth; d++) {
    result = negamax(state, d, d, -safe_infty, safe_infty,
                     (state->ply & 1) ? +1.0 : -1.0, root_hash
#ifdef RECORD_PV
                     , &iteration_leaf
#endif
    );
  }

#ifdef RECORD_PV
  pv_publish_leaf(&iteration_leaf);
#endif

  action_to_ptn(negamax_best_action, negamax_ptn);

  return result;
}

// ===================================================================
// α-β negamax with iterative deepening.
// ===================================================================

static enum TT_FLAG flag;
static enum WIN_TYPE w;

static float negamax(tak_state_p state, const uint8_t cur_depth,
                     const uint8_t init_depth, float alpha, float beta,
                     const float colour, const uint64_t hash
#ifdef RECORD_PV
                     , pv_leaf_t *out_leaf
#endif
) {

  // the scratch buffer is per-level so that recursion does not clobber
  const uint32_t level = init_depth - cur_depth;

  tt_entry_t *entry = tt_seek(hash, evaluate);

  // The root never early-returns from the table. Within one game the new
  // root was only ever stored at depth D-1, so an early return could not
  // fire anyway; across games the table persists across reset_state and
  // a stale full-depth entry would early-return and leave
  // negamax_best_action unset.
  if (level > 0 && entry != NULL && entry->depth >= cur_depth) {
    if (entry->flag == TT_EXACT) {
#ifdef RECORD_PV
      *out_leaf = entry->leaf;
#endif
      return entry->value;
    } else if (entry->flag == TT_LOWERBOUND && entry->value > alpha) {
      alpha = entry->value;
    } else if (entry->flag == TT_UPPERBOUND && entry->value < beta) {
      beta = entry->value;
    }
    if (alpha >= beta) {
#ifdef RECORD_PV
      *out_leaf = entry->leaf;
#endif
      return entry->value;
    }
  }

  action_t *actions;
  uint32_t num_actions;
  if ((actions = actions_generate_at(state, &num_actions, level)) == NULL)
    return alpha; // should never happen!

  // How did we land up with impossible actions for this hash?
  if (entry != NULL && action_in_list(entry->action, actions, num_actions)) {
    action_move_to_front(entry->action, actions, num_actions);
  }

  // negamax_best_action may not be in this position's move list: if the root
  // search early-returned from the transposition table (the table persists
  // across plies), it still holds the *previous* ply's best action, which is a
  // different side's move. action_move_to_front is a no-op for absent actions,
  // so this is safe either way.
  if (init_depth > 1 && cur_depth == init_depth) {
    action_move_to_front(negamax_best_action, actions, num_actions);
  }

  const float orig_alpha = alpha;

  action_t best_action = actions[0];
  float best_value = -infty;
#ifdef RECORD_PV
  pv_leaf_t candidate_leaf, best_leaf;
#endif

  for (uint32_t idx = 0; idx<num_actions; idx++) {
    negamax_display_progress(cur_depth, init_depth, num_actions);

    action_take(state, actions[idx]);

    const uint64_t child_hash = state->hash;

    // Compute the value of the node and retain the leaf that produced it.
    float node_value;
    if (state->ply >= 2 * state->board_size - 2 &&
        (w = check_win_cached(state, child_hash)) < 0xFF) {
      if (w == WIN_DRAGON)
        w = (colour > 0) ? WIN_ROAD_BLACK : WIN_ROAD_WHITE;
      node_value = -colour * infty;
      // Check win if far enough into the game
      if (w == WIN_ROAD_BLACK || w == WIN_FLAT_BLACK) {
        node_value = colour * infty;
      } else if (w == WIN_DRAW) {
        node_value = 0;
      }
#ifdef RECORD_PV
      pv_capture_leaf(&candidate_leaf, state, w);
#endif
    } else if (cur_depth > 1) {
      // If nobody won, or too early and not leaf, recurse
      node_value = -negamax(state, cur_depth - 1, init_depth, -beta, -alpha,
                            -colour, child_hash
#ifdef RECORD_PV
                            , &candidate_leaf
#endif
      );
    } else {
      node_value = colour * evaluate_memo(state, child_hash);
#ifdef RECORD_PV
      pv_capture_leaf(&candidate_leaf, state, (enum WIN_TYPE)0xFF);
#endif
    }
    action_undo(state, actions[idx]);

    if (idx == 0 || node_value > best_value) {
      best_value = node_value;
      best_action = actions[idx];
#ifdef RECORD_PV
      best_leaf = candidate_leaf;
#endif
    }

    if (best_value > alpha)
      alpha = best_value;
    if (alpha >= beta)
      break;
  }

  if (cur_depth == init_depth)
    negamax_best_action = best_action;

#ifdef RECORD_PV
  *out_leaf = best_leaf;
#endif

  flag = TT_EXACT;
  if (best_value >= beta)
    flag = TT_LOWERBOUND;
  else if (best_value <= orig_alpha)
    flag = TT_UPPERBOUND;

  entry = tt_store(hash, evaluate, flag, cur_depth, best_value, best_action);
#ifdef RECORD_PV
  if (entry != NULL)
    entry->leaf = best_leaf;
#endif

  return best_value;
}