/*
  selfplay: generate training data

  Copyright (C) 2026, tslil clingman

  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 <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#include "actions.h"
#include "negamax.h"
#include "pv.h"
#include "tak.h"
#include "xorshift64.h"

#ifndef RECORD_PV
#error "selfplay must be built with -DRECORD_PV"
#endif

#ifndef CAND_A
#error "selfplay must be built with -DCAND_A=<id>"
#endif

#ifndef CAND_B
#error "selfplay must be built with -DCAND_B=<id>"
#endif

#include CAND_A_WEIGHTS_5
#include CAND_B_WEIGHTS_5
#include CAND_A_WEIGHTS_6
#include CAND_B_WEIGHTS_6

// sorry not sorry
#define STR2U(a, b) a##_##b
#define STR4U(a, b, c, d) a##_##b##_##c##_##d
#define STR4U_EXPAND(a, b, c, d) STR4U(a, b, c, d)
#define C(variant, size, suffix) STR4U_EXPAND(cand, STR2U(CAND, variant), size, suffix)
#define NETWORK(size, called)                                           \
  static inline float STR2U(size, called)(tak_state_p state) {          \
    static float cur_board[STR2U(size, INP_NUM)] NN_ALIGNED;            \
    nn_prepare_input(cur_board, state, C(called, size, embedding));    \
    return eval_net(cur_board, STR2U(size, INP_NUM), &C(called, size, dense_weights)[0][0], \
                    C(called, size, dense_biases), C(called, size, logit_weights), C(called, size, logit_bias), \
                    STR2U(size, DENSE_NUM), state->ply & 1 ? 1.0f : -1.0f); \
  }

NETWORK(FIVE, A)
NETWORK(FIVE, B)
NETWORK(SIX, A)
NETWORK(SIX, B)

#define MAX_PLY 255
#define MAX_CELLS (6 * 6)

static tak_state_p state;

void negamax_display_progress(const uint8_t cur_depth,
                              const uint8_t init_depth,
                              const uint32_t length) {
  (void)(cur_depth);
  (void)(init_depth);
  (void)(length);
}

static int8_t win_value(const enum WIN_TYPE w) {
  switch (w) {
    case WIN_FLAT_BLACK:
    case WIN_ROAD_BLACK:
    return 1;
    case WIN_FLAT_WHITE:
    case WIN_ROAD_WHITE:
    return -1;
    default:
    return 0;
  }
}

static void emit_board(const float *head, const uint8_t *cells,
                       const uint8_t n_cells) {
  fprintf(stdout, "\"head\":[%.8f,%.8f],\"cells\":[", head[0], head[1]);
  for (uint8_t c = 0; c < n_cells; c++)
    fprintf(stdout, "%s%u", c ? "," : "", cells[c]);
  fprintf(stdout, "]");
}

static int play_game(const uint8_t min_ply, const nn_function fn_1,
                     const nn_function fn_2, const uint32_t cand_id_1,
                     const uint32_t cand_id_2,
                     char record[][16], uint8_t replay) {
  float root_head[MAX_PLY][2];
  uint8_t root_cells[MAX_PLY][MAX_CELLS];
  float row_head[MAX_PLY][2];
  uint8_t row_cells[MAX_PLY][MAX_CELLS];
  uint8_t row_term[MAX_PLY];
  float row_terminal_value[MAX_PLY];
  float row_parity_sign[MAX_PLY];
  uint32_t row_cand_id[MAX_PLY];
  int8_t row_mover_sign[MAX_PLY];
  uint32_t n_rows;

  reset_state(state, state->board_size);
  n_rows = 0;

  char ptn[16];
  for (uint32_t ply = 0; ply + negamax_search_depth < MAX_PLY; ply++) {
    if (ply < min_ply) {
      if (replay == 0) {
        uint32_t n_actions;
        action_t *actions = actions_generate(state, &n_actions);
        action_to_ptn(actions[rand32() % n_actions], ptn);
        strcpy(record[ply],ptn);
      } else {
        strcpy(ptn, record[ply]);
      }
    } else {
      const nn_function fn = (state->current_colour & 1) ? fn_1 : fn_2;
      negamax_set_evaluator(fn);
      negamax_generate(state);

      strcpy(ptn, negamax_ptn);

      const pv_leaf_t *leaf = pv_get_leaf();
      memcpy(row_head[n_rows], leaf->head, sizeof(row_head[n_rows]));
      memcpy(row_cells[n_rows], leaf->cells, sizeof(row_cells[n_rows]));
      row_term[n_rows] = leaf->terminal;
      row_terminal_value[n_rows] = leaf->terminal_value;
      row_parity_sign[n_rows] = leaf->parity_sign;
      row_cand_id[n_rows] = (state->current_colour&1) ? cand_id_1 : cand_id_2;
      row_mover_sign[n_rows] = (state->current_colour&1) ? 1 : -1;

      // Net input of the root state s_t, before the move is played
      nn_prepare_head(root_head[n_rows], state);
      nn_cell_idxs(root_cells[n_rows], state);
      n_rows++;
    }


    const enum ACT_RESULT r = do_ptn(state, ptn);
    if (r == GAME_END)
      break;
    if (r != ACT_OK) {
      fprintf(stderr, "selfplay: unexpected result %d for ptn %s\n", r,
              ptn);
      return 0;
    }
  }

  if (n_rows == 0)
    return 0;

  const uint8_t n_cells = state->board_size * state->board_size;

  // final outcome of the played game, from black's perspective
  const int8_t R = win_value(state->won);

  fprintf(stdout, "{\"result\":%d,\"cand_A\":%u,\"cand_B\":%u,\"rows\":[", R,
          CAND_A, CAND_B);
  for (uint32_t k = 0; k < n_rows; k++) {
    fprintf(stdout, "%s{\"cand\":%u,\"root\":{\"parity\":%d,", k ? "," : "",
            row_cand_id[k], row_mover_sign[k]);
    emit_board(root_head[k], root_cells[k], n_cells);
    fprintf(stdout, "},\"leaf\":{\"parity\":%.0f,", row_parity_sign[k]);
    emit_board(row_head[k], row_cells[k], n_cells);
    fprintf(stdout, ",\"terminal\":%u,\"terminal_value\":%.1f}}",
            row_term[k], row_terminal_value[k]);
  }
  fprintf(stdout, "]}\n");

  return (int)n_rows;
}

const char *usage =
  "selfplay, one game of training data between candidates.\n\n\
Usage: selfplay <size> <depth> <min_ply>\n\n\
Plays a single game of <size>x<size>, with <min_ply> moves made randomly. Only the deterministic plies are written to stdout, one JSON record per game, from ply <min_ply> onwards:\n\n\
\t{\"result\": R, \"cand_A\": i, \"cand_B\": j, \"rows\": [{\"cand\": id, \"root\": {\"parity\", \"head\", \"cells\"}, \"leaf\": {\"parity\", \"head\", \"cells\", \"terminal\", \"terminal_value\"}}]}\n\n\
R is the outcome of the game from black's perspective; cand_A and cand_B are the ids of the two candidates that played the game. Each row records the net input of the root state s_t and the PV leaf selected at that ply; the leaf value Net(pv) is taken from the trained net. nn_train.py runs many of these in parallel to generate training data.\n\n\
Copyright (C) 2026, tslil clingman\n\n\
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.";

int main(int argc, char **argv) {
  if (argc != 4) {
    puts(usage);
    return EXIT_FAILURE;
  }

  const uint8_t size = (uint8_t)atoi(argv[1]);
  negamax_search_depth = (uint8_t)atoi(argv[2]);
  const uint8_t min_ply = (uint8_t)atoi(argv[3]);

  if (size != 5 && size != 6) {
    fprintf(stderr, "selfplay: size must be 5 or 6\n");
    return EXIT_FAILURE;
  }
  if (negamax_search_depth > NEGAMAX_MAX_DEPTH) {
    fprintf(stderr, "selfplay: depth must be <= %d\n", NEGAMAX_MAX_DEPTH);
    return EXIT_FAILURE;
  }
  if (min_ply >= MAX_PLY) {
    fprintf(stderr, "selfplay: min_ply must be < %d\n", MAX_PLY);
    return EXIT_FAILURE;
  }

  // Processes launched in the same second must not share a seed
  xors = ((uint64_t)time(NULL) ^ ((uint64_t)getpid() << 32)) *
             0x9E3779B97F4A7C15ULL +
         0xC0FFEE;

  state = new_tak_state(size);
  negamax_init(size);

  char record[min_ply][16];

  play_game(min_ply,
            size == 5 ? FIVE_A : SIX_A, size == 5 ? FIVE_B : SIX_B,
            CAND_A,
            CAND_B,
            record, 0);
  play_game(min_ply,
            size == 5 ? FIVE_B : SIX_B, size == 5 ? FIVE_A : SIX_A,
            CAND_B,
            CAND_A,
            record, CAND_A == CAND_B ? 0: 1);


  negamax_free();
  free_tak_state(state);

  return EXIT_SUCCESS;
}