/*
  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 <string.h>

#include "nn.h"

#include "weights_5.h"
#include "weights_6.h"

// ===================================================================
// Helpers
// ===================================================================

static float fast_tanh(const float x) {
  // tanh via a degree-11 odd polynomial, least-squares fit of tanh on [-3, 3]
  //  (max absolute error ~4e-3), saturating to +-1 outside that interval.
  if (x > 3.0f)
    return 1.0f;
  if (x < -3.0f)
    return -1.0f;
  const float x2 = x * x;
  return x * (0.997136032f + x2 * (-0.305215124f + x2 *
                                   (0.0832466371f + x2 *
                                    (-0.0139947924f + x2 *
                                     (0.00123106314f + x2 * -4.29748644e-05f)))));
}

static float fast_silu(const float x) {
  // [4,2] Pade approximant with cutoff
  if (x < -4.6f)
    return 0.0;
  if (x > 4.6f)
    return x;
  return (x * (0.5f + x * (0.25f + x * (0.05f + 0.004167f * x)))) / (1+0.1f*x);
}
// ===================================================================
// Input encoding
// ===================================================================

inline uint8_t cell_idx_at(const tak_state_p state, const uint8_t loc) {
  const uint8_t d = state->celldat[loc];
  switch (d & NUM_MASK) {
    case 0: {
      // index 0
      return 0;
    }
    default: {
      // |{flat, standing, cap} * {black, white}| = 6
      // index 1 .. 6
      const uint8_t top_colour = TOP_COLOUR_AT(state, loc);
      return (uint8_t)(1u + 3u * ((state->ply & 1) ? 1-top_colour : top_colour)  + (d & STONE_MASK));
    }
  }
}

inline void nn_cell_idxs(uint8_t *out, tak_state_p state) {
  const unsigned int n_cells = state->board_size * state->board_size;
  for (unsigned int loc = 0; loc < n_cells; loc++)
    out[loc] = cell_idx_at(state, loc);
}

inline void nn_prepare_head(float out_head[2], tak_state_p state) {
  const float inv_max_flats =
    (state->board_size == 6) ? (1.0f / 30.0f) : (1.0f / 21.0f);
  const uint8_t idx = state->ply & 1;
  out_head[idx] = (float)(state->white_count & 127) * inv_max_flats;
  out_head[1-idx] = (float)(state->black_count & 127) * inv_max_flats;
}

// The cell indices computed by the most recent nn_prepare_input, and the hash
// of the state that produced them.
static uint8_t prepared_idxs[36]; // max board is 6x6
static uint64_t prepared_hash = 0;

void nn_prepare_input(float *cur_board, tak_state_p state,
                      const float embedding[NN_N_EMB]) {
  const unsigned int n_cells = state->board_size * state->board_size;
  nn_prepare_head(cur_board, state);
  for (unsigned int loc = 0; loc < n_cells; loc++) {
    const uint8_t idx = cell_idx_at(state, loc);
    prepared_idxs[loc] = idx;
    cur_board[2 + loc] = embedding[idx];
  }
  prepared_hash = state->hash;
}

// Copy the cell indices produced by the most recent nn_prepare_input into out
// and return 1, but only if that call was made on exactly this state, otherwise
// return 0 and leave out unfilled.
int nn_last_cell_idxs(uint8_t *out, tak_state_p state) {
  const unsigned int n_cells = state->board_size * state->board_size;
  if (prepared_hash != state->hash)
    return 0;
  memcpy(out, prepared_idxs, n_cells);
  return 1;
}

// ===================================================================
// Implementation of small neural networks
// ===================================================================

float eval_net(const float *board, const unsigned int n_inp, const float *w,
               const float *biases, const float *logit_w, const float logit_b,
               const unsigned int n_dense,
              const float parity) {
  float dense1[FIVE_DENSE_NUM]; // FIVE_DENSE_NUM == SIX_DENSE_NUM
  for (unsigned int d1 = 0; d1 < n_dense; d1++)
    dense1[d1] = biases[d1];
  for (unsigned int fl = 0; fl < n_inp; fl++) {
    const float b = board[fl];
    const float *wrow = w + fl * n_dense;
    for (unsigned int d1 = 0; d1 < n_dense; d1++)
      dense1[d1] += b * wrow[d1];
  }
  float z = logit_b;
  for (unsigned int d1 = 0; d1 < n_dense; d1++) {
    const float h = dense1[d1];
    dense1[d1] = fast_silu(h);
    z += dense1[d1] * logit_w[d1];
  }
  return parity * fast_tanh(z);
}

// ===================================================================
// Interfaces to the two built-in evaluators
// ===================================================================

float five_nn_fn(tak_state_p state) {
  static float cur_board[FIVE_INP_NUM] NN_ALIGNED;

  nn_prepare_input(cur_board, state, FIVE_embedding);
  return eval_net(cur_board, FIVE_INP_NUM, &FIVE_dense_weights[0][0],
                  FIVE_dense_biases, FIVE_logit_weights, FIVE_logit_bias,
                  FIVE_DENSE_NUM, state->ply & 1 ? 1.0f : -1.0f);
}

float six_nn_fn(tak_state_p state) {
  static float cur_board[SIX_INP_NUM] NN_ALIGNED;

  nn_prepare_input(cur_board, state, SIX_embedding);
  return eval_net(cur_board, SIX_INP_NUM, &SIX_dense_weights[0][0],
                  SIX_dense_biases, SIX_logit_weights, SIX_logit_bias,
                  SIX_DENSE_NUM, state->ply & 1 ? 1.0f : -1.0f);
}