/*
  This file is part of ct.

  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 ct. If not, see <https://www.gnu.org/licenses/>.
*/

#include "actions.h"
#include "tak.h"

#include <stdio.h>

static uint8_t al_board_size;
static uint32_t upper_bound_moves;

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

#define DANGER_MIN(a, b) (((a) < (b)) ? (a) : (b))

#define CLR_STONE NUM_MASK

static inline void inline_next_ply(tak_state_p state);

static inline void inline_prev_ply(tak_state_p state);

static inline uint8_t check_no_overflow(tak_state_p state, const uint8_t loc,
                                        const int8_t delta, const uint8_t num,
                                        const uint8_t steps,
                                        const uint8_t gaps);

// ===================================================================
// Exported method implementations
// ===================================================================

inline void actions_free(action_t *actions) { free(actions); }

// Keep track of move offsets
int8_t move_deltas[4];

void actions_init(const uint8_t board_size) {
  al_board_size = board_size;
  move_deltas[0] = +board_size;
  move_deltas[1] = -board_size;
  move_deltas[2] = -1;
  move_deltas[3] = +1;

  const uint32_t five_cumulative_partitions = 7 + 5 + 3 + 2 + 1; // 18
  const uint32_t six_cumulative_partitions =
      11 + five_cumulative_partitions; // 29
  if (board_size == 5) {
    // 26 = 3*3 + 4*3 + 1

    // which means central squares (4 dirs), sides (3 dirs), one corner (2 dirs)
    upper_bound_moves = 4 * 3 * 3 + 3 * 4 * 3 + 2 * 1;
    upper_bound_moves *= five_cumulative_partitions;
  } else {
    // 31 = 4*4 + 3*4 + 3

    // which means central squares (4 dirs), three sides (3 dirs) and an extra
    // three squares on the last side (1 dir)
    upper_bound_moves = 4 * 4 * 4 + 3 * 3 * 4 + 3 * 3;
    upper_bound_moves *= six_cumulative_partitions;
  }
  // In summary, using typedef uint32_t action_t
  //  74 * 18 = 1332 for 5x5, ~5.2 Kb
  // 109 * 29 = 3161 for 6x6, ~12.3 Kb
}

inline char action_in_list(action_t action, action_t *actions, const uint32_t num_actions) {
  uint32_t idx;
  for (idx = 0; idx < num_actions && actions[idx] != action; idx++);
  return idx < num_actions;
}

void action_move_to_front(action_t action, action_t *actions,
                          const uint32_t num_actions) {
  (void) num_actions; // yolo
  uint32_t idx;
  action_t prev = action, temp;
  for (idx = 0; actions[idx] != action; idx++) {
    temp = actions[idx];
    actions[idx] = prev;
    prev = temp;
  }
  actions[idx] = prev;
}

// We bias place over move by prepending place actions and appending
// move actions to the generated list
action_t *actions_generate(tak_state_p state, uint32_t *num_actions) {
  /*
   * The check for whether it's a black piece to be played is actually
   *     black = (ply < 2) ? (ply==1) : (ply & 1),
   * but material will always be sufficient in ply < 2 so we might as
   * well save on the conditional.
   */

  const uint8_t material =
                    (state->ply & 1) ? state->black_count : state->white_count,
                flat = material & 0x7F,
                cap = ((state->ply >= 2) && (material & 0x80)),
                standing = ((state->ply >= 2) && flat);

  // Count placement options
  uint32_t placements = 0;
  const uint32_t how_many = (flat ? 1 : 0) + (cap ? 1 : 0) + (standing ? 1 : 0);
  for (int row = 0; row < state->board_size; row++) {
    for (int col = 0; col < state->board_size; col++) {
      const int l = THE_COORDS(al_board_size, col, row);
      if (COUNT_AT(state, l) == 0)
        placements += how_many;
    }
  }

  // TODO: trap errno
  action_t *actions =
      malloc(sizeof(action_t) * (placements + upper_bound_moves));
  uint32_t total = 0, move_idx = placements, place_idx = placements;

  // Step across the board
  for (int row = 0; row < state->board_size; row++) {
    for (int col = 0; col < state->board_size; col++) {
      // We'll need these at various points: the location of this
      // square and the maximum number of stones we could pick up
      const int loc = THE_COORDS(al_board_size, col, row);
      const uint8_t count = DANGER_MIN(COUNT_AT(state, loc), al_board_size);

      // Only try moves after CPS and if the colour is correct
      if (count) {
        if (state->ply >= 2 &&
            ((state->colours[loc] & 1) == state->current_colour)) {

          // Pre-compute end-stops and crushes
          uint8_t end_stops[4], crushes[4] = {0, 0, 0, 0};

          // These are upper bounds, not counting walls and such.
          // UP DOWN LEFT RIGHT
          end_stops[0] = DANGER_MIN(al_board_size - row - 1, count);
          end_stops[1] = DANGER_MIN(row, count);
          end_stops[2] = DANGER_MIN(col, count);
          end_stops[3] = DANGER_MIN(al_board_size - col - 1, count);

          // Now we check for caps and walls
          const uint8_t cap_top = STONE_AT(state, loc) == STONE_CAPSTONE;
          for (int d = 0; d < 4; d++) {
            const int delta = move_deltas[d];
            const int stop = end_stops[d];
            end_stops[d] = 0;
            for (int k = 1; k <= stop; k++) {
              const enum STONE_VARIANT stone = STONE_AT(state, loc + k * delta);
              if (stone == STONE_STANDING) {
                if (cap_top) {
                  crushes[d] = k;
                  end_stops[d]++;
                }
                break;
              } else if (stone == STONE_CAPSTONE) {
                break;
              }
              end_stops[d]++;
            }
          }
          /*
           * For each direction, generate all possible ordered integer
           * partitions of 1 ≤ num ≤ count whose number of summands is in the
           * range 1 ≤ # summands ≤ min(end_stops[dir], num) -- we write
           * summands as steps.
           *
           * We exploit the `gaps' bijection here and elsewhere between ordered
           * {integer partitions of n with s summands} and {binary strings of
           * length n-1 with s-1 set bits}.
           */
          for (enum MOVE_DIRECTION dir = M_UP; dir <= M_RIGHT; dir++) {
            const int8_t delta = move_deltas[dir];
            for (uint8_t num = 1; num <= count; num++) {
              for (uint8_t steps = 1; steps <= end_stops[dir] && steps <= num;
                   steps++) {
                uint8_t gaps = ((1 << (al_board_size - 2)) - 1) >>
                               (al_board_size - steps - 1);
                // For 5x5 this gives 0b0000[0XXX] where steps-1 of those X's
                // are 1s (starting with LSB) because 4-1=3 and 5-1=4
                do {
                  /*
                   * We skip the partition if it calls for multiple stones at
                   * the end with a crush, or if it would cause any stack to
                   * grow beyond height 15.
                   */
                  const uint8_t no_overflow =
                      check_no_overflow(state, loc, delta, num, steps, gaps);
                  const uint8_t last_drop_check =
                      (num > 1) ? (gaps & (1 << (num - 2))) : 1;
                  const uint8_t can_and_must_crush =
                      crushes[dir] == steps && last_drop_check;
                  const uint8_t crush_check =
                      can_and_must_crush || crushes[dir] != steps;
                  if (no_overflow && crush_check) {
                    // Store the move
                    actions[move_idx++] =
                        A_BUILD(A_MOVE, loc, (can_and_must_crush << 7) | gaps,
                                (dir << 4) | num);
                    total++;
                  }
                  /*
                   * With thanks to
                   * https://graphics.stanford.edu/~seander/bithacks.html#NextBitPermutation
                   * we have the following magic to generate the next
                   * permutation of steps-many set bits
                   */
                  uint8_t t = (gaps | (gaps - 1));
                  gaps =
                      (t + 1) | (((~t & -~t) - 1) >> (__builtin_ctz(gaps) + 1));
                } while (gaps && (gaps + 1 <= (1 << (num - 1))));
              }
            }
          }
        }
      } // end of if (count) { ... }
      else if (material) {
        // Empty square, generate placements
        if (flat) {
          actions[--place_idx] = A_BUILD(A_PLACE, loc, STONE_FLAT, 0);
          total++;
          if (standing) {
            actions[--place_idx] = A_BUILD(A_PLACE, loc, STONE_STANDING, 0);
            total++;
          }
        }
        if (cap) {
          actions[--place_idx] = A_BUILD(A_PLACE, loc, STONE_CAPSTONE, 0);
          total++;
        }
      }
    }
  }
  *num_actions = total;
  return actions;
}

void action_take(tak_state_p state, const action_t action) {
  const int8_t loc = A_GET_LOC(action);
  if (A_GET_TYPE(action) == A_PLACE) {
    const uint8_t black = (state->current_colour == C_BLACK);
    switch (A_GET_DATA0(action)) {
    case STONE_FLAT: {
      if (black)
        state->black_count--;
      else
        state->white_count--;
      state->colours[loc] = state->current_colour;
      state->celldat[loc] = NUM_INC | STONE_FLAT;
      break;
    }
    case STONE_STANDING: {
      if (black)
        state->black_count--;
      else
        state->white_count--;
      state->colours[loc] = state->current_colour;
      state->celldat[loc] = NUM_INC | STONE_STANDING;
      break;
    }
    default: {
      if (black)
        state->black_count &= 0x7F;
      else
        state->white_count &= 0x7F;
      state->colours[loc] = state->current_colour;
      state->celldat[loc] = NUM_INC | STONE_CAPSTONE;
      break;
    }
    }
  } else {
    /*
     * See the discussion around line 160 for an explanation of the encoding.
     * Here we are not interested in whether we crushed, it will work out by
     * anyway because we overwrite the top stone type. See (*) later for when we
     * do need to know.
     */
    const uint8_t gaps = A_GET_DATA0(action) & 0x7F,
                  num = A_GET_DATA1(action) & 0x0F, // unpack
        dir = A_GET_DATA1(action) >> 4;
    int8_t delta = move_deltas[dir];

    // Use the Kernighan method to count the set bits
    int8_t steps = 1;
    for (uint8_t _gaps = gaps; _gaps; steps++)
      _gaps &= _gaps - 1;

    // Move top stone type to destination
    state->celldat[loc + steps * delta] &= CLR_STONE; // necessary for crushing
    state->celldat[loc + steps * delta] |= STONE_AT(state, loc);
    state->celldat[loc] &= CLR_STONE;
    state->celldat[loc] |= STONE_FLAT; // should be optimised out

    uint8_t total = 1, gap_bit = 1 << (num - 2); // it's not important what
                                                 // negative shifts do here, we
                                                 // don't use gap_bit if num < 2

    // move stuff starting at destination
    for (uint8_t d = 1; d < num; d++, total++, gap_bit >>= 1) {
      // We took a step, move everything over so far
      if (gaps & gap_bit) {
        state->colours[loc + steps * delta] <<= total;
        state->colours[loc + steps * delta] |=
            state->colours[loc] & ((1 << total) - 1);
        state->colours[loc] >>= total;
        state->celldat[loc + steps * delta] += total * NUM_INC;
        state->celldat[loc] -= total * NUM_INC;
        // Reset for next step
        total = 0;
        steps--;
      }
    }
    // Move what remains (steps == 1 here always, so we simplify)
    state->colours[loc + delta] <<= total;
    state->colours[loc + delta] |= state->colours[loc] & ((1 << total) - 1);
    state->colours[loc] >>= total;
    state->celldat[loc + delta] += total * NUM_INC;
    state->celldat[loc] -= total * NUM_INC;
  }
  // Next ply
  inline_next_ply(state);
}

void action_undo(tak_state_p state, const action_t action) {
  // Previous ply
  inline_prev_ply(state);

  const int8_t loc = A_GET_LOC(action);
  if (A_GET_TYPE(action) == A_PLACE) {
    const uint8_t black = (state->current_colour == C_BLACK);
    state->celldat[loc] = 0;
    if (A_GET_DATA0(action) == STONE_CAPSTONE) {
      if (black)
        state->black_count |= 0x80;
      else
        state->white_count |= 0x80;
    } else {
      if (black)
        state->black_count++;
      else
        state->white_count++;
    }
  } else {
    // See action_take for comments, this is the time reversal, but
    // there is one caveat -- undoing a crush! (*)
    const uint8_t gaps = A_GET_DATA0(action) & 0x7F,
                  crush = A_GET_DATA0(action) & 0x80,
                  num = A_GET_DATA1(action) & 0x0F,
                  dir = A_GET_DATA1(action) >> 4;
    const int8_t delta = move_deltas[dir];

    int8_t steps = 1;
    uint8_t gap_bit = 1, total = 1;
    for (int8_t d = 1; d < num; d++, total++, gap_bit <<= 1) {
      if (gaps & gap_bit) {
        state->colours[loc] <<= total;
        state->colours[loc] |=
            state->colours[loc + steps * delta] & ((1 << total) - 1);
        state->colours[loc + steps * delta] >>= total;
        state->celldat[loc] += total * NUM_INC;
        state->celldat[loc + steps * delta] -= total * NUM_INC;
        total = 0;
        steps++;
      }
    }
    state->colours[loc] <<= total;
    state->colours[loc] |=
        state->colours[loc + steps * delta] & ((1 << total) - 1);
    state->colours[loc + steps * delta] >>= total;

    state->celldat[loc] += total * NUM_INC;
    // celldat[loc] &= CLR_STONE; is not necessary, as STONE_FLAT == 0
    state->celldat[loc] |= STONE_AT(state, loc + steps * delta);
    state->celldat[loc + steps * delta] -= total * NUM_INC;
    state->celldat[loc + steps * delta] &= CLR_STONE;
    if (crush) {
      state->celldat[loc + steps * delta] |= STONE_STANDING;
    } else {
      state->celldat[loc + steps * delta] |=
          STONE_FLAT; // should be optimised out
    }
  }
}

void action_to_ptn(const action_t action, char *out_ptn) {
  const int8_t loc = A_GET_LOC(action);
  if (A_GET_TYPE(action) == A_PLACE) {
    generate_place(al_board_size, loc, A_GET_DATA0(action), out_ptn);
  } else {
    const uint8_t gaps = A_GET_DATA0(action) & 0x7F,
                  num = A_GET_DATA1(action) & 0x0F, // unpack
        dir = A_GET_DATA1(action) >> 4;

    uint8_t drops[al_board_size]; // we only ever need al_board_size-1 in drops
                                  // actually, the last spot is to skip a bounds
                                  // check at (**)
    uint8_t mask = 1, steps = 0;
    // Translate to a drop sequence
    drops[0] = 1;
    mask = 1;
    for (uint8_t d = 1; d < num; d++) {
      if (gaps & mask) {
        steps++;
        drops[steps] = 1; // (**) no bounds check
      } else {
        drops[steps] += 1;
      }
      mask <<= 1;
    }
    generate_move(al_board_size, loc, dir, steps + 1, drops, out_ptn);
  }
}

// ===================================================================
// Helper method implementations
// ===================================================================

static inline void inline_next_ply(tak_state_p state) {
  state->ply++;
  if (state->ply == 2) {
    state->current_colour = C_WHITE;
  } else {
    if (state->current_colour == C_BLACK)
      state->current_colour = C_WHITE;
    else
      state->current_colour = C_BLACK;
  }
}

static inline void inline_prev_ply(tak_state_p state) {
  if (state->ply > 0)
    state->ply--;
  if (state->ply == 1) {
    state->current_colour = C_WHITE;
  } else {
    if (state->current_colour == C_BLACK)
      state->current_colour = C_WHITE;
    else
      state->current_colour = C_BLACK;
  }
}

static inline uint8_t check_no_overflow(tak_state_p state, const uint8_t loc,
                                        const int8_t delta, const uint8_t num,
                                        const uint8_t steps,
                                        const uint8_t gaps) {
  uint8_t total = 1, gap_bit = 1 << (num - 2), step = steps;
  for (uint8_t d = 1; d < num; d++, total++, gap_bit >>= 1) {
    if (gaps & gap_bit) {
      if (COUNT_AT(state, loc + step * delta) + total > 15)
        return 0;
      step--;
      total = 0;
    }
  }
  if (COUNT_AT(state, loc + delta) + total > 15)
    return 0;

  return 1;
}