#include <locale.h>
#include <stdio.h>
#include <string.h>
#include <time.h>
#include <stdlib.h>
#include <math.h>

#include "config.h"
#include "pieces.h"
#include "structs.h"

#define MAX(x, y) (((x) > (y)) ? (x) : (y))
#define MIN(x, y) (((x) < (y)) ? (x) : (y))

int TRANSPOSITION_TABLE_CARDINALITY = 0;

void generate_zobrist_seed(long* seed) {
  for (int i = 0; i < MAX_ZOBRIST_SEEDS; i++) {
    seed[i] = rand();
  }
}

int zobrist_hash(long* seed, int* board, int color_to_move) {
  long hash = 1;

  // Last feature means white to move
  if (color_to_move == WHITE) hash ^= seed[MAX_ZOBRIST_SEEDS - 1];

  for (int rank = 7; rank >= 0; rank--) {
    for (int file = 0; file < 8; file++) {
      int position = rank * 16 + file;
      int piece = board[position];

      // Unique index of this piece on this position
      int zobrist_index = piece * 128 + position;
      hash ^= seed[zobrist_index];
    }
  }
  return hash;
}

long hash_after_move(Move move, int* board, long* zobrist_seed, long hash) {
  int piece = board[move.origin];
  int target_piece = board[move.destination];

  hash ^= zobrist_seed[MAX_ZOBRIST_SEEDS - 1]; // Flip color
  hash ^= zobrist_seed[EMPTY * 128 + move.origin]; // Add empty to origin
  hash ^= zobrist_seed[piece * 128 + move.origin]; // Remove piece from origin
  hash ^= zobrist_seed[piece * 128 + move.destination]; // Add piece to destination
  hash ^= zobrist_seed[target_piece * 128 + move.destination]; // Remove target piece from destination

  return hash;
}


void print_board(int board[128]) {
  printf("\n");
  for (int rank = 7; rank >= 0; rank--) {
    printf("%i|", rank + 1);
    for (int file = 0; file < 8; file++) {
      int piece = board[rank * 16 + file];
      printf("%s ", pieces[piece]);
    }
    printf("\n");
  }
  printf("  a b c d e f g h\n");
}

void parse_FEN(char* FEN, int board[128]) {
  // TODO: parse everything, not position only
  int rank = 7;
  int file = 0;
  size_t lenght = strlen(FEN);

  for (int k = 0; k < lenght; k++) {
    if (FEN[k] > '0' && FEN[k] <= '8') {
      int last_file = file + FEN[k] - '0';
      for (file; file < last_file; file++) {
        board[rank * 16 + file] = EMPTY;
      }
    } else {
      switch (FEN[k]) {
        case 'r':
          board[rank * 16 + file] = BLACK | ROOK;
          break;
        case 'R':
          board[rank * 16 + file] = ROOK;
          break;
        case 'n':
          board[rank * 16 + file] = BLACK | KNIGHT;
          break;
        case 'N':
          board[rank * 16 + file] = KNIGHT;
          break;
        case 'b':
          board[rank * 16 + file] = BLACK | BISHOP;
          break;
        case 'B':
          board[rank * 16 + file] = BISHOP;
          break;
        case 'q':
          board[rank * 16 + file] = BLACK | QUEEN;
          break;
        case 'Q':
          board[rank * 16 + file] = QUEEN;
          break;
        case 'k':
          board[rank * 16 + file] = BLACK | KING;
          break;
        case 'K':
          board[rank * 16 + file] = KING;
          break;
        case 'p':
          board[rank * 16 + file] = BLACK | PAWN;
          break;
        case 'P':
          board[rank * 16 + file] = PAWN;
          break;
        case '/':
          rank--;
          file = -1; // so that it becomes 0
          break;
        case ' ':
          return;
        default:
          board[rank * 16 + file] = KNIGHT;
      }
      file++;
    }
  }
}

int notation_to_index(int file, int rank) {
  return (rank - 1) * 16 + (file - 'a');
}

void index_to_notation(int index, char notation[2]) {
  notation[0] = (index & FILE_MASK) + 'a';
  notation[1] = (index >> 4) + '1';
}

void print_move(Move move) {
  char move_in_notation[] = "xy XY";
  index_to_notation(move.origin, move_in_notation);
  index_to_notation(move.destination, move_in_notation + 3);
  printf("%s\n", move_in_notation);
};

Move input_move() {
  Move move;
  char file;
  int rank;
  char file_destination;
  int rank_destination;

  scanf(" %c%i %c%i", &file, &rank, &file_destination, &rank_destination);

  move.origin = notation_to_index(file, rank);
  move.destination = notation_to_index(file_destination, rank_destination);
  return move;
};

int list_moves_with_destination(Move* moves, int* board, int color, int destination);

int validate_move(Move move, int color, int* board) {
  // Null move means player is checkmated
  if (move.origin == -100 && move.destination == -100) return INFINITY;

  if ((move.origin & 0x88) || (move.destination & 0x88)) return -1;

  int piece = board[move.origin];
  int* legal_move;

  if (piece % 2 != color) return -3;

  switch (piece & NO_COLOR) {
    case KNIGHT:
      legal_move = knightMoves;
      break;
    case BISHOP:
      legal_move = bishopMoves;
      break;
    case ROOK:
      legal_move = rookMoves;
      break;
    case QUEEN:
      legal_move = queenMoves;
      break;
    case KING:
      legal_move = kingMoves;
      break;
    case PAWN:
      if (piece & BLACK) {
        if (board[move.destination] != EMPTY) legal_move = blackPawnAttackMoves;
        else legal_move = blackPawnMoves;
      } else {
        if (board[move.destination] != EMPTY) legal_move = pawnAttackMoves;
        else legal_move = pawnMoves;
      }
      break;
    default:
      return -2;
  }

  while(*legal_move) {
    for (int square = move.origin + *legal_move; !(square & 0x88); square += *legal_move) {
      int target_piece = board[square];

      if (square == move.destination) {
        if (target_piece != EMPTY && target_piece % 2 == color) return -4;
        return target_piece;
      }

      if (target_piece != EMPTY) break;
      if (
          (piece & NO_COLOR) == PAWN &&
          !(
            (move.origin >> 4 == 6 || move.origin >> 4 == 1) && // Pawn is new
            abs(move.origin - square) == 16 // And this is only first move
          )
      ) break;
      if ((piece & NO_COLOR) == KNIGHT || (piece & NO_COLOR) == KING) break;
    }
    legal_move++;
  }

  // Handle castling
  if ((piece & NO_COLOR) == KING) {
    if (piece % 2 == WHITE && move.origin == 4) {
      if (move.destination == 2) {
        if (list_moves_with_destination(NULL, board, color ^ 1, move.origin)) return -5;
        for (int i = 1; i < 4; i++) {
          if (board[i] != EMPTY || list_moves_with_destination(NULL, board, color ^ 1, i)) return -5;
        }
        return 0;
      } else if (move.destination == 6) {
        if (list_moves_with_destination(NULL, board, color ^ 1, move.origin)) return -5;
        for (int i = 5; i < 7; i++) {
          if (board[i] != EMPTY || list_moves_with_destination(NULL, board, color ^ 1, i)) return -5;
        }
        return 0;
      }
    } else if (piece % 2 == BLACK && move.origin == 116) {
      if (move.destination == 114) {
        if (list_moves_with_destination(NULL, board, color ^ 1, move.origin)) return -5;
        for (int i = 114; i < 116; i++) {
          if (board[i] != EMPTY || list_moves_with_destination(NULL, board, color ^ 1, i)) return -5;
        }
        return 0;
      } else if (move.destination == 118) {
        if (list_moves_with_destination(NULL, board, color ^ 1, move.origin)) return -5;
        for (int i = 117; i < 119; i++) {
          if (board[i] != EMPTY || list_moves_with_destination(NULL, board, color ^ 1, i)) return -5;
        }
        return 0;
      }
    }
  }

  return -6;
}

char* VALIDATION_ERRORS[] = {
  "trying to access off-board square",
  "no piece on origin square",
  "opponent piece on origin square",
  "trying to capture your own piece",
  "move is not allowed",
  "castle is not allowed",
};

int make_move(Move move, int* board) {
  int piece = board[move.origin];
  int target_piece = board[move.destination];

  board[move.destination] = piece;
  board[move.origin] = EMPTY;

  if (target_piece != EMPTY) {
    // printf("Captured %s\n", pieces[target_piece]);
  }

  if ((piece & NO_COLOR) == PAWN) {
    int rank = move.destination / 16;
    if (rank == 0) {
      board[move.destination] = QUEEN | BLACK;
      return target_piece | PROMOTION_HAPPENED;
    } else if (rank == 7) {
      board[move.destination] = QUEEN | WHITE;
      return target_piece | PROMOTION_HAPPENED;
    }
  }

  if ((piece & NO_COLOR) == KING && abs(move.destination - move.origin) == 2) {
    // CASTLE!
    if (move.destination == 2) {
      int rook = board[0];
      board[0] = EMPTY;
      board[3] = rook;
    } else if (move.destination == 6) {
      int rook = board[7];
      board[7] = EMPTY;
      board[5] = rook;
    } else if (move.destination == 114) {
      int rook = board[112];
      board[112] = EMPTY;
      board[115] = rook;
    } else if (move.destination == 118) {
      int rook = board[119];
      board[119] = EMPTY;
      board[117] = rook;
    }
  }

  return target_piece;
}

void unmake_move(Move move, int captured_piece, int* board) {
  int piece = board[move.destination];

  if (captured_piece & PROMOTION_HAPPENED) {
    board[move.destination] = captured_piece ^ PROMOTION_HAPPENED;
    if (piece % 2 == WHITE) {
      board[move.origin] = PAWN | WHITE;
    } else board[move.origin] = PAWN | BLACK;
    return;
  }

  board[move.origin] = piece;
  board[move.destination] = captured_piece;

  if ((piece & NO_COLOR) == KING && abs(move.destination - move.origin) == 2) {
    // CASTLE!
    if (move.destination == 2) {
      int rook = board[3];
      board[3] = EMPTY;
      board[0] = rook;
    } else if (move.destination == 6) {
      int rook = board[5];
      board[5] = EMPTY;
      board[7] = rook;
    } else if (move.destination == 114) {
      int rook = board[115];
      board[115] = EMPTY;
      board[112] = rook;
    } else if (move.destination == 118) {
      int rook = board[117];
      board[117] = EMPTY;
      board[119] = rook;
    }
  }
}


int evaluate_pawn(int piece, int rank) {
  int color = piece % 2;
  if (color == 0) {
    switch (rank) {
      case 6:
        return 5;
      case 5:
        return 3;
      default:
        return 1;
    }
  } else {
    switch (rank) {
      case 1:
        return 5;
      case 2:
        return 3;
      default:
        return 1;
    }
  }
}

/*
 * Compute material advantage of a player over an opponent
 * Returns number in range (0, 30)
 */
int compute_material_advantage(int* board, int color) {
  int counter = 0;
  for (int rank = 7; rank >= 0; rank--) {
    for (int file = 0; file < 8; file++) {
      int piece = board[rank * 16 + file];
      if (piece != EMPTY) {
        int sign = (piece % 2 == color) ? 1 : -1;
        switch (piece & NO_COLOR) {
          case KING:
            counter += INFINITY * sign;
            break;
          case PAWN:
           counter += 10 * evaluate_pawn(piece, rank) * sign;
           break;
          case KNIGHT:
            counter += 30 * sign;
            break;
          case BISHOP:
            counter += 30 * sign;
            break;
          case ROOK:
            counter += 45 * sign;
            break;
          case QUEEN:
            counter += 90 * sign;
            break;
          default:
            counter += 0 * sign;
            break;
        }
      }
    }
  }
  return counter;
}

int list_moves_with_destination(Move* moves, int* board, int color, int destination) {
  int moves_count = 0;

  for (int rank = 7; rank >= 0; rank--) {
    for (int file = 0; file < 8; file++) {
      int index = rank * 16 + file;
      int piece = board[index];
      if (piece != EMPTY && (piece % 2) == color) {
        Move move = { index, destination, 0 };
        if (validate_move(move, color, board) >= 0) {
          if (moves) moves[moves_count] = move;
          moves_count++;
        }
      }
    }
  }

  return moves_count;
}

int list_available_moves(Move* moves, int* board, int color) {
  int moves_count = 0;

  for (int rank = 7; rank >= 0; rank--) {
    for (int file = 0; file < 8; file++) {
      int destination = rank * 16 + file;
      Move* next_ptr = moves ? moves + moves_count : NULL;
      moves_count += list_moves_with_destination(next_ptr, board, color, destination);
    }
  }

  return moves_count;
}

// Count doubled, blocked and isolated pawns
int count_bad_pawns(int* board, int color) {
  int bad_pawns = 0;
  long occupied_files = 0b00000000;

  for (int file = 0; file < 8; file++) {
    int file_pawns = 0;
    for (int rank = 7; rank >= 0; rank--) {
      int piece = board[rank * 16 + file];
      if ((piece & NO_COLOR) == PAWN && piece % 2 == color) {
        file_pawns++;
        occupied_files |= 1 << file;

        // If pawn can't advance, it's blocked
        int advance_rank = rank + (color == WHITE ? 1 : -1);
        if (board[advance_rank * 16 + file] != EMPTY) bad_pawns++;
      }
    }

    // Doubled or tripled pawns will get punished
    bad_pawns += file_pawns - 1;
  }

  // Shift everything to the left to get some space around right-most bit
  occupied_files = occupied_files << 1;

  // Count isolated pawns based on bits in occupied files
  for (int bit = 1; bit < 9; bit++) {
    long bit_mask = (1 << (bit + 1)) | (1 << (bit - 1));
    if (!(occupied_files & bit_mask)) bad_pawns++;
  }

  return bad_pawns;
}

/*
 * Return a total board evaluation symbolizing advantage
 * for WHITE, meaining WHITE tries to maximize this
 * value and BLACK tries to minimize it.
 */
int evaluate_position(int* board, int precomputed_mobility, int mobility_color) {
  Move dummy[MAX_AVAILABLE_MOVES];

  int white_material_advantage = compute_material_advantage(board, WHITE);

  // If pre-computed mobility has not been passed, compute it
  int mobility = precomputed_mobility ? precomputed_mobility : list_available_moves(NULL, board, mobility_color);
  int opponent_mobility = list_available_moves(NULL, board, mobility_color ^ 1);

  int white_mobility_advantage = (mobility - opponent_mobility) * (mobility_color == WHITE ? 1 : -1);

  int pawn_structure_advantage = count_bad_pawns(board, BLACK) - count_bad_pawns(board, WHITE);

  return white_material_advantage + white_mobility_advantage + 5 * pawn_structure_advantage;
}

int get_piece_raw_value(int piece) {
  switch (piece & NO_COLOR) {
    case KNIGHT:
      return 3;
    case BISHOP:
      return 3;
    case ROOK:
      return 5;
    case QUEEN:
      return 9;
    case KING:
      return INFINITY;
    case PAWN:
      return 1;
    default:
      return 0;
  }
}

int compare_moves(const void* a, const void* b) {
  return ((Move*)b)->value - ((Move*)a)->value;
}

int order_moves(Move* moves, int moves_count, int* board, long hash, int* transposition_table, long* zobrist_seed) {
  int cache_hits = 0;
  for (int i = 0; i < moves_count; i++) {
    int color = board[moves[i].origin] % 2;
    // Lookup move in transpoition table
    long move_hash = hash_after_move(moves[i], board, zobrist_seed, hash);
    if (transposition_table[move_hash % TRANSPOSITION_TABLE_SIZE]) {
      moves[i].value = transposition_table[move_hash % TRANSPOSITION_TABLE_SIZE] * (color == WHITE ? 1 : -1);
      cache_hits++;
    } else {
      int origin_value = get_piece_raw_value(board[moves[i].origin]);
      int destination_value = get_piece_raw_value(board[moves[i].destination]);
      moves[i].value = 10 * destination_value - origin_value;
    }
  }

  qsort(moves, moves_count, sizeof(Move), compare_moves);

  return cache_hits; // Only needed for display
}

/*
 * Alpha-beta pruning
 * Value here is white material advantage over black
 * Alpha is the best value for maximizer (white)
 * Beta is the best value for minimizer (black)
 */
Move minimax_search(int* board, int color, int depth, int alpha, int beta, int* metrics, long hash, int* transposition_table, long* zobrist_seed) {
  int is_maximizer = (color == 0);

  Move best_move = { -100, -100, 0 };
  best_move.value = is_maximizer ? -100 * INFINITY / (depth + 1) : 100 * INFINITY / (depth + 1);

  // You only have available moves if your king hasn't been taken
  if (compute_material_advantage(board, color) > -INFINITY / 2) {
    Move available_moves[MAX_AVAILABLE_MOVES];
    int available_moves_count = list_available_moves(available_moves, board, color);
    int cache_hits = order_moves(available_moves, available_moves_count, board, hash, transposition_table, zobrist_seed);

    for (int i = 0; i < available_moves_count; i++) {
      *metrics += 1;
      Move move = available_moves[i];

      long move_hash = hash_after_move(move, board, zobrist_seed, hash);
      int captured_piece = make_move(move, board);

      move.value = depth < MAX_DEPTH
        ? minimax_search(board, 1 - color, depth + 1, alpha, beta, metrics, hash, transposition_table, zobrist_seed).value
        : evaluate_position(board, available_moves_count, color);

      if (!transposition_table[move_hash % TRANSPOSITION_TABLE_SIZE]) TRANSPOSITION_TABLE_CARDINALITY++;
      transposition_table[move_hash % TRANSPOSITION_TABLE_SIZE] = move.value;

      unmake_move(move, captured_piece, board);

      if (is_maximizer) {
        if (move.value > best_move.value) best_move = move;
        alpha = MAX(best_move.value, alpha);
      } else {
        if (move.value < best_move.value) best_move = move;
        beta = MIN(best_move.value, beta);
      }

      if (beta <= alpha) break;
    }
  }

  return best_move;
}

int main() {
  int board[128];
  parse_FEN(DEFAULT_FEN, board);

  long zobrist_seed[MAX_ZOBRIST_SEEDS];
  generate_zobrist_seed(zobrist_seed);
  long hash = zobrist_hash(zobrist_seed, board, WHITE);

  int transposition_table[TRANSPOSITION_TABLE_SIZE];

  Move move;

  print_board(board);

  int ply = 0;
  while (1) {
    int color = ply % 2;
    printf("##### Ply %i #####\n", ply);
    printf("Current evaluation: %.2f\n", (double)evaluate_position(board, 0, WHITE) / 10);

    if (color == PLAYER) {
      printf("Enter a move for %s:\n", COLORS[color]);
      move = input_move();
    } else {
      printf("Evaluating move for %s...\n", COLORS[color]);
      clock_t start, end;
      int metrics = 0;
      start = clock();
      move = minimax_search(board, color, 0, -INFINITY, +INFINITY, &metrics, hash, transposition_table, zobrist_seed);
      end = clock();
      printf("[%i positions analyzed in %f seconds]\n", metrics, (double)(end - start) / CLOCKS_PER_SEC);
      printf("[Transposition table cardinality: %i]\n", TRANSPOSITION_TABLE_CARDINALITY);

      char move_in_notation[] = "xy XY";
      index_to_notation(move.origin, move_in_notation);
      index_to_notation(move.destination, move_in_notation + 3);
      printf("%s: %s with evaluation %c=%.2f on Ply %i\n", COLORS[color], move_in_notation, color == WHITE ? '>' : '<', (double)move.value / 10, ply + MAX_DEPTH);
    }

    int error = validate_move(move, color, board);
    if (error < 0) {
      printf("Invalid move: %s!\n", VALIDATION_ERRORS[-1 - error]);
    } else if (error == INFINITY) {
      printf("Checkmate, %s is victorious!\n", COLORS[1 - color]);
      break;
    } else {
      hash = hash_after_move(move, board, zobrist_seed, hash);
      make_move(move, board);

      print_board(board);

      ply++;
    }
  }

	return 0;
}