// Copyright 2019 DeepMind Technologies Ltd. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "open_spiel/algorithms/kriegspiel/darkboard.h"

#include "open_spiel/games/chess/chess_board.h"
#include "open_spiel/games/kriegspiel.h"

#include "open_spiel/algorithms/mcts.h"
#include "open_spiel/abseil-cpp/absl/random/discrete_distribution.h"
#include "open_spiel/abseil-cpp/absl/random/distributions.h"

namespace open_spiel {
namespace algorithms {
constexpr double kTieTolerance = 0.00001;
constexpr double kDumpingCoef = 0.9;
constexpr double kAfterCheckCapture = 0.3;
constexpr double kCapture = 0.2;

double RandomNumber(std::mt19937 rng) { return absl::Uniform(rng, 0.0, 1.0); }

SimpleKriegspielUmpireMessage GetKriegspielMsg(std::mt19937 rng, const std::vector<std::pair<double, SimpleKriegspielUmpireMessage>>& messages) {
  double r = RandomNumber(rng);
  for (const auto &[p, msg] : messages) {
    r -= p;
    if (r <= 0) {
      return msg;
    }
  }
  return messages[messages.size() - 1].second;
}

chess::Square RandomSquare(std::mt19937 rng, const std::vector<std::pair<chess::Square, double>>& sq_probs) {
  double sum = 0.0;
  for (const auto &[sq, p] : sq_probs) {
    sum += p;
  }
  double r = absl::Uniform(rng, 0.0, sum);;
  for (const auto &[sq, p] : sq_probs) {
    r -= p;
    if (r <= 0) {
      return sq;
    }
  }
  return sq_probs[sq_probs.size() - 1].first;
}


ProbabilityBoard::ProbabilityBoard(chess::Color color, Stats* stats, int t)
  : color_(color),
    stats_(stats),
    time_(t) {

  std::string FEN = kPartialFENs[chess::ColorToPlayer(color)];
  std::optional<chess::ChessBoard> board = chess::ChessBoard::BoardFromFEN(FEN);
  SPIEL_CHECK_TRUE(board);
  board_ = *board;

  chess::ChessBoard temp_board = chess::MakeDefaultBoard();
  for (int x = 0; x < 8; ++x) {
    for (int y = 0; y < 8; ++y) {
      chess::Square sq = {x, y};

      chess::Piece piece = temp_board.at(sq);
      if (piece.color == chess::OppColor(color_)) {
        if (piece.type == chess::PieceType::kKing) {
          king_probs_.Set(sq, 1.0);
        }
        else if (piece.type == chess::PieceType::kPawn) {
          pawn_probs_.Set(sq, 1.0);
        }
        else {
          piece_probs_.Set(sq, 1.0);
        }
      }
    }
  }

}

void ProbabilityBoard::Reset() {
  std::string FEN = kPartialFENs[chess::ColorToPlayer(color_)];
  std::optional<chess::ChessBoard> board = chess::ChessBoard::BoardFromFEN(FEN);
  SPIEL_CHECK_TRUE(board);
  board_ = *board;

  chess::ChessBoard temp_board = chess::MakeDefaultBoard();
  for (int x = 0; x < 8; ++x) {
    for (int y = 0; y < 8; ++y) {
      chess::Square sq = {x, y};

      chess::Piece piece = temp_board.at(sq);
      if (piece.color == chess::OppColor(color_)) {
        if (piece.type == chess::PieceType::kKing) {
          king_probs_.Set(sq, 1.0);
        }
        else if (piece.type == chess::PieceType::kPawn) {
          pawn_probs_.Set(sq, 1.0);
        }
        else {
          piece_probs_.Set(sq, 1.0);
        }
      }
    }
  }
  time_ = 0;
  captured_pawns_ = 0;
  captured_pieces_ = 0;

}

void Normalize(Mat& mat, double desired_sum) {
  if (desired_sum == 0.0) {
    for (int x = 0; x < 8; ++x) {
      for (int y = 0; y < 8; ++y) {
        chess::Square sq = {x, y};
        mat.Set(sq, 0.0);
      }
    }
  }
  else {
    double sum = 0.0;
    for (int8_t x = 0; x < 8; ++x) {
      for (int8_t y = 0; y < 8; ++y) {
        sum += mat.At({x, y});
      }
    }
    if (sum == 0) {
      int a = 5;
      for (int x = 0; x < 8; ++x) {
        for (int y = 0; y < 8; ++y) {
          chess::Square sq = {x, y};
          mat.Set(sq, 1.0/64.0);
        }
      }
    }
    else {
      double multiplier = desired_sum / sum;
      for (int x = 0; x < 8; ++x) {
        for (int y = 0; y < 8; ++y) {
          chess::Square sq = {x, y};
          double old_prob = mat.At(sq);
          double new_prob = old_prob * multiplier;
          mat.Set(sq, new_prob);
        }
      }
    }
  }

}

void ProbabilityBoard::Normalize() {
  int pawn_count = 8 - captured_pawns_;
  algorithms::Normalize(pawn_probs_, std::max(pawn_count, 0));

  int piece_count = 7 - captured_pieces_;
  algorithms::Normalize(piece_probs_, std::max(piece_count, 0));

  algorithms::Normalize(king_probs_, 1.0);
}


void ProbabilityBoard::Update(chess::Move move, SimpleKriegspielUmpireMessage msg) {
  //std::cout << move.ToString() << std::endl;
  SPIEL_CHECK_EQ(board_.ToPlay(), color_);
  if (msg.illegal) {
    // todo maybe I can do something about this
    // todo like raise the sum of the squares to one
  }
  else {
    pawn_probs_.Set(move.to, 0.0);
    piece_probs_.Set(move.to, 0.0);
    king_probs_.Set(move.to, 0.0);

    if (move.piece.type != chess::PieceType::kKnight) {
      int8_t offset_x = 0;
      int8_t offset_y = 0;

      int diff_x = move.to.x - move.from.x;
      if (diff_x > 0) offset_x = 1;
      else if (diff_x < 0) offset_x = -1;

      int diff_y = move.to.y - move.from.y;
      if (diff_y > 0) offset_y = 1;
      else if (diff_y < 0) offset_y = -1;

      chess::Offset offset_step = {offset_x, offset_y};

      for (chess::Square dest = move.from + offset_step; dest != move.to;
           dest += offset_step) {
        pawn_probs_.Set(dest, 0.0);
        piece_probs_.Set(dest, 0.0);
        king_probs_.Set(dest, 0.0);
      }
    }

    if (msg.capture_type == kriegspiel::KriegspielCaptureType::kPawn) {
      captured_pawns_++;
    }
    else if (msg.capture_type == kriegspiel::KriegspielCaptureType::kPiece) {
      captured_pieces_++;
    }

    board_.ApplyMove(move);

  }
  Normalize();
  time_++;
  //std::cout <<board_.DebugString() << std::endl;
}

void ProbabilityBoard::Update(SimpleKriegspielUmpireMessage msg) {
  SPIEL_CHECK_EQ(board_.ToPlay(), chess::OppColor(color_));
  if (msg.capture_type != kriegspiel::kNoCapture) {
    //std::cout << chess::SquareToString(msg.square) << std::endl;
    board_.set_square(msg.square, {chess::Color::kEmpty, chess::PieceType::kEmpty});

    // bring the sum of opp probabilities to 1.0
    double old_pawn = pawn_probs_.At(msg.square);
    double old_piece = piece_probs_.At(msg.square);
    double old_king = king_probs_.At(msg.square);
    double sum = old_pawn + old_piece + old_king;

    if (sum == 0.0) {
      pawn_probs_.Set(msg.square, 0.33333);
      piece_probs_.Set(msg.square, 0.33333);
      king_probs_.Set(msg.square, 0.33333);
    }
    else {
      pawn_probs_.Set(msg.square, old_pawn / sum);
      piece_probs_.Set(msg.square, old_piece / sum);
      king_probs_.Set(msg.square, old_king / sum);
    }

  }
  else {
    UpdateKingProbs();
    UpdatePieceProbs();
    UpdatePawnProbs();
  }
  Normalize();
  time_++;
  board_.SetToPlay(color_);
  //std::cout <<board_.DebugString() << std::endl;
}

std::vector<std::pair<double, SimpleKriegspielUmpireMessage>> ProbabilityBoard::MyProbabilisticMessage(chess::Move move) {

  std::vector<std::pair<double, SimpleKriegspielUmpireMessage>> ret;

  double legal_prob = Legal(move);
  SimpleKriegspielUmpireMessage illegal_msg = {false};
  ret.emplace_back(1.0 - legal_prob, illegal_msg);

  double prob_capture_pawn = CapturePawn(move);
  double prob_capture_piece = CapturePiece(move);
  double prob_check = Check(move);
  //std::cout << prob_capture_pawn << " " << prob_capture_piece << " " << prob_check << " " << std::endl;

  SimpleKriegspielUmpireMessage piece_msg =
      {false,
       kriegspiel::KriegspielCaptureType::kPiece,
       move.to,
       false};
  double prob_piece_msg = legal_prob * prob_capture_piece * (1.0 - prob_check);
  ret.emplace_back(prob_piece_msg, piece_msg);

  SimpleKriegspielUmpireMessage pawn_msg =
      {false,
       kriegspiel::KriegspielCaptureType::kPawn,
       move.to,
       false};
  double prob_pawn_msg = legal_prob * prob_capture_pawn * (1.0 - prob_check);
  ret.emplace_back(prob_pawn_msg, pawn_msg);

  SimpleKriegspielUmpireMessage piece_check_msg =
      {false,
       kriegspiel::KriegspielCaptureType::kPiece,
       move.to,
       true};
  double prob_piece_check_msg = legal_prob * prob_capture_piece * prob_check;
  ret.emplace_back(prob_piece_check_msg, piece_check_msg);

  SimpleKriegspielUmpireMessage pawn_check_msg =
      {false,
       kriegspiel::KriegspielCaptureType::kPawn,
       move.to,
       true};
  double prob_pawn_check_msg = legal_prob * prob_capture_pawn * prob_check;
  ret.emplace_back(prob_pawn_check_msg, pawn_check_msg);

  SimpleKriegspielUmpireMessage check_msg =
      {false,
       kriegspiel::KriegspielCaptureType::kNoCapture,
       chess::kInvalidSquare,
       true};
  double prob_check_msg = legal_prob * (1.0 - prob_capture_piece - prob_capture_pawn) * prob_check;
  ret.emplace_back(prob_check_msg, check_msg);

  SimpleKriegspielUmpireMessage empty_msg =
      {false,
       kriegspiel::KriegspielCaptureType::kNoCapture,
       chess::kInvalidSquare,
       false};
  double prob_empty_msg = legal_prob * (1.0 - prob_capture_piece - prob_capture_pawn) * (1.0 - prob_check);
  ret.emplace_back(prob_empty_msg, empty_msg);

  return ret;
}

chess::Square ProbabilityBoard::ChooseVictim() {
  std::vector<std::pair<chess::Square, double>> sq_probs;

  for (int8_t x = 0; x < 8; ++x) {
    for (int8_t y = 0; y < 8; ++y) {
      chess::Square sq = {x, y};

      if (board_.at(sq).color == color_ && board_.at(sq).type != chess::PieceType::kKing) {

        sq_probs.emplace_back(sq, Control(sq));

      }
    }
  }
  if (sq_probs.empty()) {
    return chess::kInvalidSquare;
  }

  chess::Square ret = RandomSquare(rng_, sq_probs);
  return ret;
}

std::vector<std::pair<double, SimpleKriegspielUmpireMessage>> ProbabilityBoard::OppProbabilisticMessage(const SimpleKriegspielUmpireMessage& my_msg, const chess::Square& checking_square) {
  std::vector<std::pair<double, SimpleKriegspielUmpireMessage>> ret;

  double prob_recapture = 0.0;
  if (my_msg.capture_type != kriegspiel::KriegspielCaptureType::kNoCapture && board_.at(my_msg.square).type != chess::PieceType::kKing) {
    prob_recapture = Control(my_msg.square);
    if (prob_recapture > 0.0) {
      SimpleKriegspielUmpireMessage recapture_msg =
          {false,
           kriegspiel::KriegspielCaptureType::kPiece,
           my_msg.square};
      ret.emplace_back(prob_recapture, recapture_msg);
    }
  }

  for (const auto &[p, m] : ret) {
    if (m.capture_type != kriegspiel::KriegspielCaptureType::kNoCapture && board_.at(m.square).type == chess::PieceType::kKing) {
      int a= 5;
    }
  }
  double prob_after_check_capture = 0.0;
  if (checking_square != chess::kInvalidSquare && board_.at(checking_square).type != chess::PieceType::kKing) {
    prob_after_check_capture = my_msg.check && (!(prob_recapture > 0.0)) ? kAfterCheckCapture : 0.0;
    if (prob_after_check_capture > 0.0) {
      SimpleKriegspielUmpireMessage after_check_capture_msg =
          {false,
           kriegspiel::KriegspielCaptureType::kPiece,
           checking_square};
      ret.emplace_back(prob_recapture, after_check_capture_msg);
    }
  }

  chess::Square victim_sq = ChooseVictim();

  for (const auto &[p, m] : ret) {
    if (m.capture_type != kriegspiel::KriegspielCaptureType::kNoCapture && board_.at(m.square).type == chess::PieceType::kKing) {
      int a= 5;
    }
  }
  double prob_capture = 0.0;
  if (victim_sq != chess::kInvalidSquare) {
    prob_capture = kCapture * (1.0 - prob_recapture - prob_after_check_capture);
    SimpleKriegspielUmpireMessage capture_msg =
        {false,
         kriegspiel::KriegspielCaptureType::kPiece,
         victim_sq};
    ret.emplace_back(prob_capture, capture_msg);
  }

  for (const auto &[p, m] : ret) {
    if (m.capture_type != kriegspiel::KriegspielCaptureType::kNoCapture && board_.at(m.square).type == chess::PieceType::kKing) {
      int a= 5;
    }
  }
  double prob_empty = 1 - prob_recapture - prob_after_check_capture - prob_capture;
  SimpleKriegspielUmpireMessage empty_msg =
      {false,
       kriegspiel::KriegspielCaptureType::kNoCapture,
       chess::kInvalidSquare};
  ret.emplace_back(prob_empty, empty_msg);

  for (const auto &[p, m] : ret) {
    if (m.capture_type != kriegspiel::KriegspielCaptureType::kNoCapture && board_.at(m.square).type == chess::PieceType::kKing) {
      int a= 5;
    }
  }

  // check is basically Control(king_square) and adds no value to the simulation
  return ret;
}

double ProbabilityBoard::Control(const chess::Square& sq) {
  double king_control = 0;
  for (int8_t i = -1; i <= 1; ++i) {
    for (int8_t j = -1; j <= 1; ++j) {
      chess::Square k_sq = sq + chess::Offset{i, j};
      if (board_.InBoardArea(k_sq)) {
        king_control += king_probs_.At(k_sq);
      }
    }
  }

  double pawn_control = 0;
  int8_t opp_y_direction = color_ == chess::Color::kWhite ? -1 : 1;
  if (board_.InBoardArea(sq + chess::Offset{-1, opp_y_direction}))
    pawn_control += pawn_probs_.At(sq + chess::Offset{-1, opp_y_direction});
  if (board_.InBoardArea(sq + chess::Offset{1, opp_y_direction}))
    pawn_control += pawn_probs_.At(sq + chess::Offset{1, opp_y_direction});


  double piece_control = 0;
  chess::Offset offsets[] =
      {{0, 1}, {1, 0}, {0, -1}, {-1, 0},
       {1, 1}, {-1, -1}, {1, -1}, {-1, 1}};
  for (const chess::Offset& offset : offsets) {

    chess::Square p_sq = sq;
    double c2 = 1.0;
    while (board_.InBoardArea(p_sq)) {
      double addition = piece_probs_.At(p_sq);
      piece_control += c2 * addition;

      c2 *= (1.0 - addition);
      p_sq += offset;
    }
    
  }
  piece_control *= 3.0 / 7.0;

  return king_control + pawn_control + piece_control;
}

double ProbabilityBoard::Legal(const chess::Move& move) {
 double prob_legal = 1.0;
 // straight pawn advance needs the target square empty
 if (move.piece.type == chess::PieceType::kPawn &&
      move.to.x - move.from.x == 0) {
   prob_legal *= 1.0 - king_probs_.At(move.to) - piece_probs_.At(move.to) - pawn_probs_.At(move.to);
 }
 else if (move.piece.type != chess::PieceType::kKnight){
   int8_t offset_x = 0;
   int8_t offset_y = 0;

   int diff_x = move.to.x - move.from.x;
   if (diff_x > 0) offset_x = 1;
   else if (diff_x < 0) offset_x = -1;

   int diff_y = move.to.y - move.from.y;
   if (diff_y > 0) offset_y = 1;
   else if (diff_y < 0) offset_y = -1;

   chess::Offset offset_step = {offset_x, offset_y};

   for (chess::Square dest = move.from + offset_step; dest != move.to;
        dest += offset_step) {

     prob_legal *= 1.0 - king_probs_.At(move.to) - piece_probs_.At(move.to) - pawn_probs_.At(move.to);
   }

 }

 double prob_pin = 0.0;
 if (move.piece.type == chess::PieceType::kKing) {
   prob_pin = Control(move.to);
 }


 return prob_legal - prob_pin;
}

double ProbabilityBoard::CapturePiece(const chess::Move& move) {
  return piece_probs_.At(move.to);
}

double ProbabilityBoard::CapturePawn(const chess::Move& move) {
  if (!(pawn_probs_.At(move.to) >= 0)) {
    double b = pawn_probs_.At(move.to);
    int a = 5;
  }
  return pawn_probs_.At(move.to);
}

double ProbabilityBoard::Check(const chess::Move& move) {
  chess::ChessBoard board_copy = board_;
  board_copy.ApplyMove(move);

  std::vector<chess::Square> attacked_squares;
  board_copy.GeneratePseudoLegalMoves([&attacked_squares] (const chess::Move& move_inner) {
    attacked_squares.push_back(move_inner.to);
    return true;
    }, move.to, move.piece.color);

  double sum = 0.0;
  for (const chess::Square& sq : attacked_squares) {
    double curr_prob = king_probs_.At(sq);
    if (move.piece.type == chess::PieceType::kKnight) {
      sum += curr_prob;
    }
    else {
      int dist = std::max(std::abs(move.to.x - sq.x), std::abs(move.to.y - sq.y));
      sum += std::pow(kDumpingCoef, dist - 1) * curr_prob;
    }
  }

  return sum;
}

void ProbabilityBoard::UpdateKingProbs() {
  int pawn_count = 8 - captured_pawns_;
  int piece_count = 7 - captured_pieces_;
  double prob_king_moves = 1.0 / (pawn_count + piece_count + 1.0);

  Mat prev_king_probs = king_probs_;
  for (int8_t i = 0; i < 8; ++i) {
    for (int8_t j = 0; j < 8; ++j) {
      chess::Square sq = {i, j};

      double sum = 0.0;
      for (int x = -1; x <= 1; ++x) {
        for (int y = -1; y <= 1; ++y) {
          chess::Square loc_sq = {i + x, j + y};
          if (!board_.InBoardArea(loc_sq) || loc_sq == sq) {
            continue;
          }
          sum += prev_king_probs.At(loc_sq);
        }
      }
      double target_prob = prev_king_probs.At(sq);
      double new_prob = (1 - prob_king_moves) * target_prob
                        + prob_king_moves * sum;
      king_probs_.Set(sq, new_prob);
    }
  }
}

void ProbabilityBoard::UpdatePieceProbs() {
  int pawn_count = 8 - captured_pawns_;
  int piece_count = 7 - captured_pieces_;
  double prob_piece_moves = piece_count / (pawn_count + piece_count + 1.0);

  Mat prev_piece_probs = piece_probs_;

  // scan files
  for (int8_t x = 0; x < 8; ++x) {
    chess::Offset offset = {0, 1};
    chess::Square sq = {x, 0};
    chess::Square seq_start_sq = chess::kInvalidSquare;
    int k = 0;
    double sum = 0.0;
    while (board_.InBoardArea(sq)) {
      bool sq_empty = board_.at(sq).type == chess::PieceType::kEmpty;
      if (seq_start_sq == chess::kInvalidSquare && sq_empty) {
        seq_start_sq = sq;
        k = 1;
      }
      if (sq_empty) {
        sum += prev_piece_probs.At(sq);
        k++;
      }
      else {
        seq_start_sq = chess::kInvalidSquare;
        k = 0;
        sum = 0.0;
      }
      if (seq_start_sq != chess::kInvalidSquare &&
           (!board_.InBoardArea(sq + offset) || board_.at(sq + offset).type != chess::PieceType::kEmpty)) {
        if (sq.y - seq_start_sq.y > 1) {
          // sq is the end of the sequence
          // seq_start_sq is the beginning of the sequence

          chess::Square loc_sq = seq_start_sq;
          while(loc_sq != sq + offset) {

            double c1 = 3.0 / 7.0;
            double c2 = 1.0 / (k - 1);

            double target_prob = prev_piece_probs.At(loc_sq);
            double new_prob = (1 - prob_piece_moves) * target_prob
                              + c1 * prob_piece_moves * c2 * sum;
            piece_probs_.Set(loc_sq, new_prob);

            loc_sq += offset;
          }


        }

        seq_start_sq = chess::kInvalidSquare;
      }
      sq += offset;
    }

  }


  // scan ranks
  for (int8_t y = 0; y < 8; ++y) {
    chess::Offset offset = {1, 0};
    chess::Square sq = {0, y};
    chess::Square seq_start_sq = chess::kInvalidSquare;
    int k = 0;
    double sum = 0.0;
    while (board_.InBoardArea(sq)) {
      bool sq_empty = board_.at(sq).type == chess::PieceType::kEmpty;
      if (seq_start_sq == chess::kInvalidSquare && sq_empty) {
        seq_start_sq = sq;
        k = 1;
      }
      if (sq_empty) {
        sum += prev_piece_probs.At(sq);
        k++;
      }
      else {
        seq_start_sq = chess::kInvalidSquare;
        k = 0;
        sum = 0.0;
      }
      if (seq_start_sq != chess::kInvalidSquare &&
          (!board_.InBoardArea(sq + offset) || board_.at(sq + offset).type != chess::PieceType::kEmpty)) {
        if (sq.y - seq_start_sq.y > 1) {
          // sq is the end of the sequence
          // seq_start_sq is the beginning of the sequence

          chess::Square loc_sq = seq_start_sq;
          while(loc_sq != sq + offset) {

            double c1 = 3.0 / 7.0;
            double c2 = 1.0 / (k - 1);

            double target_prob = prev_piece_probs.At(loc_sq);
            double new_prob = (1 - prob_piece_moves) * target_prob
                              + c1 * prob_piece_moves * c2 * sum;
            piece_probs_.Set(loc_sq, new_prob);

            loc_sq += offset;
          }


        }

        seq_start_sq = chess::kInvalidSquare;
      }
      sq += offset;
    }

  }

}

void ProbabilityBoard::UpdatePawnProbs() {
  int pawn_count = 8 - captured_pawns_;
  int piece_count = 7 - captured_pieces_;
  double prob_pawn_moves = pawn_count / (pawn_count + piece_count + 1.0);

  int8_t y_direction = color_ == chess::Color::kWhite ? -1 : 1;
  chess::Offset offset = {0, y_direction};

  Mat prev_pawn_probs = pawn_probs_;
  for (int8_t x = 0; x < 8; ++x) {
    for (int8_t y = 1; y < 7; ++y) {
      chess::Square sq = {x, y};
      chess::Square target_sq = sq + offset;

      double new_prob = prob_pawn_moves * prev_pawn_probs.At(sq) + (1 - prob_pawn_moves) * prev_pawn_probs.At(target_sq);
      pawn_probs_.Set(target_sq, new_prob);
    }
  }
}

DarkboardBot::DarkboardBot(int seed, int move_time,
                           std::function<double(const ProbabilityBoard&, std::optional<SimpleKriegspielUmpireMessage>, chess::Square)> evaluator,
                           double uct_c,
                           ISMCTSFinalPolicyType final_policy_type,
                           chess::Color color,
                           Stats* stats)
    : rng_(seed),
      move_time_(move_time),
      evaluator_(std::move(evaluator)),
      uct_c_(uct_c),
      final_policy_type_(final_policy_type),
      prob_board_(color, stats),
      color_(color),
      player_(chess::ColorToPlayer(color)) {

  my_turn_node_pool_.reserve(10000);
  opp_turn_node_pool_.reserve(10000);
  my_msg_turn_node_pool_.reserve(10000);
}

void DarkboardBot::Reset() {
  my_turn_node_pool_.clear();
  opp_turn_node_pool_.clear();
  my_msg_turn_node_pool_.clear();
}

void DarkboardBot::Restart() {
  prob_board_.Reset();
  Reset();
}

Action DarkboardBot::Step(const State& state) {
  std::chrono::time_point<std::chrono::system_clock> step_start = std::chrono::high_resolution_clock::now();
  std::chrono::time_point<std::chrono::system_clock> step_end = step_start + std::chrono::milliseconds(move_time_);

  legal_actions = state.LegalActions();

  SetCurrentPlayer(state);
  if (!current_partial_board) {
    current_partial_board = chess::ChessBoard::BoardFromFEN(kPartialFENs[*player_]);
  }

  if (last_move_ && last_my_msg_) {
    prob_board_.Update(*last_move_, GetSimpleKriegspielUmpireMessage(*last_my_msg_));
  }
  if (last_opp_msg_) {
    prob_board_.Update(GetSimpleKriegspielUmpireMessage(*last_opp_msg_));
  }

  auto is_update_end = std::chrono::high_resolution_clock::now();
  std::cout << "Information set update took: "
            << std::chrono::duration_cast<std::chrono::milliseconds>(
                is_update_end - step_start).count() << std::endl;

  // find best move
  chess::Move best_move = RunSearch(step_end);

  //  update observations for updating infoset next turn
  auto kriegspielState = down_cast<const kriegspiel::KriegspielState &>(state);
  kriegspiel::KriegspielUmpireMessage msg = kriegspiel::GetUmpireMessage(kriegspielState.Board(), best_move);
  last_move_ = best_move;
  last_my_msg_ = msg;
  last_opp_msg_ = {};

  return chess::MoveToAction(best_move);
}

void DarkboardBot::InformAction(const State &state,
                           Player player_id,
                           Action action) {
  auto kriegspiel_state = down_cast<const kriegspiel::KriegspielState &>(state);
  chess::Move move = chess::ActionToMove(action, kriegspiel_state.Board());
  kriegspiel::KriegspielUmpireMessage msg = kriegspiel::GetUmpireMessage(kriegspiel_state.Board(), move);
  last_opp_msg_ = msg;
}

void DarkboardBot::SetCurrentPlayer(const State &state) {
  // assert valid current player for the bot
  if (!player_) player_ = state.CurrentPlayer();
  else if (player_ != state.CurrentPlayer())
    SpielFatalError("Cannot change the player during the game");
}

chess::Move DarkboardBot::RunSearch(std::chrono::time_point<std::chrono::system_clock> step_end) {
  Reset();

  my_turn_node_pool_.push_back({});

  int a = 0;
  while(std::chrono::system_clock::now() < step_end) {
    RunSimulation();
    a++;
  }
  std::cout << "Number of simulations: " << a << std::endl;

  return GetBestMove();

}

chess::Move DarkboardBot::GetBestMove() const {
  ActionsAndProbs policy;

  const MyTurnNode& root_node = my_turn_node_pool_[root_node_idx];

  switch (final_policy_type_) {

    case ISMCTSFinalPolicyType::kMaxVisitCount: {
      chess::Move best_move;
      int max_visits = -std::numeric_limits<int>::infinity();
      for (const auto& [move, msg_node_idx] : root_node.child_nodes) {
        const MyMessageNode& msg_node = my_msg_turn_node_pool_[msg_node_idx];
        if (msg_node.visits > max_visits) {
          max_visits = msg_node.visits;
          best_move = move;
        }
      }
      SPIEL_CHECK_GT(max_visits, 0);
      return best_move;
    }

    case ISMCTSFinalPolicyType::kMaxValue: {
      chess::Move best_move;
      double max_value = -std::numeric_limits<double>::infinity();
      for (const auto& [move, msg_node_idx] : root_node.child_nodes) {
        const MyMessageNode& msg_node = my_msg_turn_node_pool_[msg_node_idx];
        double value = msg_node.return_sum / msg_node.visits;

        if (value > max_value) {
          max_value = value;
          best_move = move;
        }
      }
      SPIEL_CHECK_GT(max_value, -std::numeric_limits<double>::infinity());
      return best_move;
    }
  }
}

chess::Move DarkboardBot::SelectMoveUCB(MyTurnNode* node, const std::vector<chess::Move>& legal_moves) {
  SPIEL_CHECK_GE(legal_moves.size(), 1);
  std::vector<chess::Move> candidates;
  double max_value = -std::numeric_limits<double>::infinity();

  for (const chess::Move& move : legal_moves) {
    // Every child should have at least one visit because the child is only
    // created when the action took it in a simulation, which then increases
    // its visit count immediately.
    size_t child_node_idx = node->child_nodes[move];
    MyMessageNode* child_node = &my_msg_turn_node_pool_[child_node_idx];
    if (child_node->visits == 0) {
      continue;
    }

    double uct_val = child_node->return_sum / child_node->visits + // backlog operator - average, could be max or something
                     uct_c_ * std::sqrt(std::log(child_node->available_count) / // tady to muze byt bud kolikrat byl na vyber nebo kolikrat to bylo v rodicovi asi
                                        child_node->visits);

    if (uct_val > max_value + kTieTolerance) {
      candidates.clear();
      candidates.push_back(move);
      max_value = uct_val;
    } else if (uct_val > max_value - kTieTolerance &&
               uct_val < max_value + kTieTolerance) {
      candidates.push_back(move);
      max_value = uct_val;
    }
  }

  SPIEL_CHECK_GE(candidates.size(), 1);

  if (candidates.size() == 1) {
    return candidates[0];
  } else {
    return candidates[absl::Uniform(rng_, 0u, candidates.size())];
  }
}

void DarkboardBot::RunSimulation() {
  SPIEL_CHECK_EQ(prob_board_.board_.ToPlay(), color_);
  size_t node_idx = root_node_idx;
  MyTurnNode* node = &my_turn_node_pool_[root_node_idx];
  std::optional<OppTurnNode*> opp_turn_node = std::nullopt;
  ProbabilityBoard working_prob_board = prob_board_;

  SimpleKriegspielUmpireMessage last_msg;
  chess::Square last_square;

  while (true) {

    // find unexplored nodes
    std::vector<chess::Move> pseudo_legal_moves = working_prob_board.LegalMoves();
    std::vector<chess::Move> legal_moves;
    for (const chess::Move& move : pseudo_legal_moves) {
      for (Action action : legal_actions) {
        if (chess::MoveToAction(move, 8) == action) {
          legal_moves.push_back(move);
        }
      }
    }


    std::vector<chess::Move> unexplored_moves;

    for (const chess::Move& move : legal_moves) {
      if (!node->child_nodes.count(move)) {
        unexplored_moves.push_back(move);

        MyMessageNode my_msg_node = {};
        my_msg_node.parent = node_idx;
        my_msg_node.available_count++;
        my_msg_turn_node_pool_.push_back(my_msg_node);

        node->child_nodes[move] = my_msg_turn_node_pool_.size() - 1;
      }
      else {
        size_t my_msg_node_idx = node->child_nodes[move];
        MyMessageNode* my_msg_node = &my_msg_turn_node_pool_[my_msg_node_idx];
        if (my_msg_node->visits == 0) {
          unexplored_moves.push_back(move);
        }
        my_msg_node->available_count++;
      }
    }

    // if there are unexplored nodes, evaluate and backpropagate
    if (!unexplored_moves.empty()) {
      size_t my_move_idx = absl::Uniform(rng_, 0u, unexplored_moves.size());
      chess::Move my_move = unexplored_moves[my_move_idx];
      size_t my_msg_node_idx = node->child_nodes[my_move];
      MyMessageNode* my_msg_node = &my_msg_turn_node_pool_[my_msg_node_idx];

      std::vector<std::pair<double, SimpleKriegspielUmpireMessage>> my_msgs
          = working_prob_board.MyProbabilisticMessage(my_move);
      SimpleKriegspielUmpireMessage my_msg = GetKriegspielMsg(rng_, my_msgs);
      working_prob_board.Update(my_move, my_msg);
      last_msg = my_msg;
      last_square = my_move.to;

      OppTurnNode new_opp_node = {};
      new_opp_node.parent = my_msg_node_idx;
      opp_turn_node_pool_.push_back(new_opp_node);
      size_t new_opp_node_idx = opp_turn_node_pool_.size() - 1;
      my_msg_node->child_nodes[my_msg] = new_opp_node_idx;
      opp_turn_node = &opp_turn_node_pool_[new_opp_node_idx];

      break;
    }

    // select the move and descend with the move
    chess::Move my_move = SelectMoveUCB(node, legal_moves);
    size_t my_msg_node_idx = node->child_nodes[my_move];
    MyMessageNode* my_msg_node = &my_msg_turn_node_pool_[my_msg_node_idx];

    // generate probabilistic message and descend with the message
    std::vector<std::pair<double, SimpleKriegspielUmpireMessage>> my_msgs
      = working_prob_board.MyProbabilisticMessage(my_move);
    SimpleKriegspielUmpireMessage my_msg = GetKriegspielMsg(rng_, my_msgs);

    // update the probabilities with the move and message
    working_prob_board.Update(my_move, my_msg);

    if (!my_msg_node->child_nodes.count(my_msg)) {
      OppTurnNode new_opp_node = {};
      new_opp_node.parent = my_msg_node_idx;
      opp_turn_node_pool_.push_back(new_opp_node);
      size_t new_opp_node_idx = opp_turn_node_pool_.size() - 1;
      my_msg_node->child_nodes[my_msg] = new_opp_node_idx;
    }
    size_t opp_node_idx = my_msg_node->child_nodes[my_msg];
    opp_turn_node = &opp_turn_node_pool_[opp_node_idx];

    // generate probabilistic opponent message, descend and update the probabilities

    std::vector<std::pair<double, SimpleKriegspielUmpireMessage>> opp_msgs
        = working_prob_board.OppProbabilisticMessage(my_msg, my_msg.check ? my_move.to : chess::kInvalidSquare);
    SimpleKriegspielUmpireMessage opp_msg = GetKriegspielMsg(rng_, opp_msgs);
    working_prob_board.Update(opp_msg);

    if (!(*opp_turn_node)->child_nodes.count(opp_msg)) {
      MyTurnNode new_my_node = {};
      new_my_node.parent = opp_node_idx;

      my_turn_node_pool_.push_back(new_my_node);
      size_t my_new_node_idx = my_turn_node_pool_.size() - 1;
      MyTurnNode* inserted_my_new_node = &my_turn_node_pool_[my_new_node_idx];

      std::vector<chess::Move> legal_moves = working_prob_board.LegalMoves();

      for (const chess::Move& move : legal_moves) {
        MyMessageNode new_msg_node = {};
        new_msg_node.parent = my_new_node_idx;

        my_msg_turn_node_pool_.push_back(new_msg_node);
        size_t my_new_msg_node_idx = my_msg_turn_node_pool_.size() - 1;

        inserted_my_new_node->child_nodes[move] = my_new_msg_node_idx;
      }

      break;
    }

    node_idx = (*opp_turn_node)->child_nodes[opp_msg];
    node = &my_turn_node_pool_[node_idx];

  }

  double value = evaluator_(working_prob_board, last_msg, last_square);

  SPIEL_CHECK_TRUE(opp_turn_node);

  MyMessageNode* my_msg_node;
  while ((*opp_turn_node)->parent) {

    my_msg_node = &my_msg_turn_node_pool_[*(*opp_turn_node)->parent];
    my_msg_node->visits++;
    my_msg_node->return_sum += value;
    MyTurnNode* my_turn_node = &my_turn_node_pool_[*my_msg_node->parent];
    if (!my_turn_node->parent) {
      break;
    }
    opp_turn_node = &opp_turn_node_pool_[*my_turn_node->parent];

  }

}

SimpleKriegspielUmpireMessage GetSimpleKriegspielUmpireMessage(const kriegspiel::KriegspielUmpireMessage &orig_msg) {
  return {orig_msg.illegal,
      orig_msg.capture_type,
      orig_msg.square,
      orig_msg.check_types
      != std::pair{kriegspiel::KriegspielCheckType::kNoCheck,
                   kriegspiel::KriegspielCheckType::kNoCheck}};
}

}  // namespace algorithms
}  // namespace open_spiel
