/*
  pptdb, generate neural network training data from a playtak.com
  database dump

  Copyright (C) 2021, 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 this program.  If not, see <https://www.gnu.org/licenses/>.
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <tak.h>

int generate;
uint64_t heights[16];
FILE *training_fh = NULL;
float max_flats;
uint8_t outcome_black;

tak_state_p state;

static void write_input(const int dx, const int dy, const uint8_t swap) {
  // Two numbers for flats remaining
  fprintf(training_fh, "%d,%.8f,%.8f,", state->ply & 1 ? 1 : -1,
          (float)(state->white_count & 127) / max_flats,
          (float)(state->black_count & 127) / max_flats);

  // Write the board layers
  float val;
  int col, row;
  row = (dy > 0) ? -1 : state->board_size;
  for (int i = 0; i < state->board_size; i++) {
    row += dy;
    col = (dx > 0) ? -1 : state->board_size;
    for (int j = 0; j < state->board_size; j++) {
      col += dx;
      const uint8_t k = (swap) ? THE_COORDS(state->board_size, row, col)
                               : THE_COORDS(state->board_size, col, row);
      val = 0;
      if (COUNT_AT(state, k) > 0) {
        // Top layer of stacks is handled differently to indicate
        // stone type
        if (STONE_AT(state, k) == STONE_STANDING) {
          val = (state->colours[k] & 1) ? +0.25 : -0.25;
        } else if (STONE_AT(state, k) == STONE_CAPSTONE) {
          val = (state->colours[k] & 1) ? +1.00 : -1.00;
        } else {
          val = (state->colours[k] & 1) ? +0.50 : -0.50;
        }
      }
      fprintf(training_fh, "%.2f,", val);
    }
  }
  fprintf(training_fh, "%d,%d\n", outcome_black ? 1 : 0, outcome_black ? 0 : 1);
}

// Warning: performs _no_ checks on input whatsoever
static enum ACT_RESULT parse_line(const char *pt, const ssize_t read) {
  ssize_t idx;
  enum ACT_RESULT r;
  int total_plies = 0;

  for (idx = 0; idx < read; idx++) {
    if (pt[idx] == ',')
      total_plies++;
  }
  for (idx = 0;;) {
    if (pt[idx] == 'P') {
      // P [A-F][1-6] [CF]?,
      idx += 2;
      enum STONE_VARIANT stone;
      const uint8_t col = pt[idx] - 'A', row = pt[idx + 1] - '1';

      if (idx + 3 < read) {
        switch (pt[idx + 3]) {
        case 'W': {
          stone = STONE_STANDING;
          break;
        }
        case 'C': {
          stone = STONE_CAPSTONE;
          break;
        }
        default: {
          stone = STONE_FLAT;
          break;
        }
        }
      } else {
        stone = STONE_FLAT;
      }

      r = try_place(state, THE_COORDS(state->board_size, col, row),
                    state->current_colour, stone);
      if (r != ACT_OK)
        return r;
    } else if (pt[idx] == 'M') {
      // M [A-F][1-6] [A-F][1-6]( [1-6])+,
      idx += 2;
      uint8_t drops[state->board_size];
      const uint8_t s_col = pt[idx] - 'A', s_row = pt[idx + 1] - '1',
                    d_col = pt[idx + 3] - 'A', d_row = pt[idx + 4] - '1';
      idx += 4;

      enum MOVE_DIRECTION dir = M_RIGHT;
      if (s_col < d_col)
        dir = M_RIGHT;
      else if (s_col > d_col)
        dir = M_LEFT;
      else if (s_row < d_row)
        dir = M_UP;
      else if (s_row > d_row)
        dir = M_DOWN;

      uint8_t steps = 0;
      do {
        idx += 2;
        drops[steps++] = pt[idx] - '0';
      } while (idx + 2 < read && pt[idx + 1] != ',');

      r = try_move(state, THE_COORDS(state->board_size, s_col, s_row), dir,
                   steps, drops);

      if (r != ACT_OK)
        return r;

      if (generate == 0) {
        // Measure height of stacks exceeding 1
        for (int k = 0; k < state->board_size * state->board_size; k++) {
          if (COUNT_AT(state, k) > 1)
            heights[COUNT_AT(state, k)] += 1;
        }
      }
    }
    // Generate training data, not too early in the game and not at
    // the end, under all eight symmetries of the board
    if (generate && state->ply > 7 && state->ply < total_plies &&
        state->ply + 10 >= total_plies) {
      write_input(+1, +1, 1);
      write_input(+1, +1, 0);
      write_input(+1, -1, 1);
      write_input(+1, -1, 0);
      write_input(-1, +1, 1);
      write_input(-1, +1, 0);
      write_input(-1, -1, 1);
      write_input(-1, -1, 0);
    }
    // Parse next action
    while (idx < read && pt[idx++] != ',')
      ;
    if (idx >= read)
      return ACT_OK;
    next_ply(state);
  }
  return ACT_OK;
}

const char *license =
    "pptdb, generate neural network training data from a playtak.com database dump\n\
\n\
Copyright (C) 2021, tslil clingman\n\
\n\
This program comes with ABSOLUTELY NO WARRANTY; and is made available under the terms of the GNU GPL v3 license. This is free software, and you are welcome to redistribute it under certain conditions; see COPYING for details.\n";

int main(int argc, char **argv) {
  (void)(argc);

  enum ACT_RESULT r;
  enum WIN_TYPE win;
  uint32_t games = 0, overflow = 0, illegal = 0;
  uint32_t road_wins = 0, flat_wins = 0, road_turns = 0, flat_turns = 0,
           white_wins = 0, black_wins = 0;

  for (int k = 0; k < 16; k++)
    heights[k] = 0;

  size_t len = 0;
  ssize_t read = 0;
  FILE *playtak_fh = NULL;
  char *line = NULL, td_fn[65];

  const uint8_t size = argv[1][0] - '0';

  playtak_fh = fopen(argv[2], "r");
  if (playtak_fh == NULL)
    exit(EXIT_FAILURE);

  if (argc > 3 && (!strncmp("generate", argv[3], 8))) {
    generate = 1;
    max_flats = (size == 5) ? 21.0 : 30.0;
    snprintf(td_fn, 64, "data/parsed-%d.csv", size);
    training_fh = fopen(td_fn, "w");
    if (training_fh == NULL)
      exit(EXIT_FAILURE);
  } else
    generate = 0;

  state = new_tak_state(size);
  while ((read = getline(&line, &len, playtak_fh)) != -1) {
    // Reset everything
    reset_state(state, size);
    // Store the outcome of this game. Black win = 1
    outcome_black = (line[read - 4] == '0');
    // Parse the line
    r = parse_line(line, read - 4);
    // Adjust counts if we're not generating training data
    if (generate == 0) {
      if (r == ACT_ILLEGAL) {
        illegal++;
        printf("Illegal:\n%s", line);
      } else if (r == ACT_OVERFLOW) {
        printf("Overflow:\n%s", line);
        overflow++;
      } else {
        win = check_win(state);
        if (win == WIN_FLAT_BLACK || win == WIN_FLAT_WHITE || win == WIN_DRAW) {
          flat_wins++;
          flat_turns += state->ply / 2 + 1;
        } else {
          road_wins++;
          road_turns += state->ply / 2 + 1;
        }
        if (win == WIN_FLAT_BLACK || win == WIN_ROAD_BLACK)
          black_wins++;
        else if (win == WIN_FLAT_WHITE || win == WIN_ROAD_WHITE)
          white_wins++;
      }
    }
    games++;
  }

  fclose(playtak_fh);
  if (generate)
    fclose(training_fh);
  if (line)
    free(line);

  if (illegal || overflow)
    putchar('\n');
  printf("Read %d games\n", games);

  if (generate == 0) {
    printf("Illegals:  %d\nOverflows: %d\n\
Black wins: %.3f%%\n\
Road wins: %d\nFlat wins: %d\n\
Average turns to road win: %.3f\n\
Average turns to flat win: %.3f\n",
           illegal, overflow,
           (double)black_wins / (double)(black_wins + white_wins) * 100,
           road_wins, flat_wins, (double)(road_turns) / (double)(road_wins),
           (double)(flat_turns) / (double)(flat_wins));
    for (int k = 2; k < 16; k++) {
      printf("Height %2d: %7ld\n", k, heights[k]);
    }
  }

  exit(EXIT_SUCCESS);
}