// 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/abseil-cpp/absl/random/distributions.h"
#include "absl/random/random.h"

#include "open_spiel/algorithms/kriegspiel/infoset.h"
#include "open_spiel/algorithms/kriegspiel/common.h"

namespace open_spiel {
namespace algorithms {


AOSPInfoset::AOSPInfoset(int seed,
                         int pool_max_size)
    : rng_(seed),
      pool_max_size_(pool_max_size),
      pool0_({chess::MakeDefaultBoard()}, pool_max_size_),
      pool1_(pool_max_size_),
      pool_(0),
      v_pool_({chess::MakeDefaultBoard()}) {
}

chess::ChessBoard AOSPInfoset::SampleState() {
  // todo implement answer No.5 https://stackoverflow.com/questions/12761315/random-element-from-unordered-set-in-o1/31522686
  uint r = absl::Uniform(rng_, 0u, v_pool_.size());
  return v_pool_[r];
}

void AOSPInfoset::Reset() {
  pool0_.clear();
  pool0_.insert(chess::MakeDefaultBoard());
  pool1_.clear();
  pool_ = 0;
  v_pool_.clear();
  v_pool_.push_back(chess::MakeDefaultBoard());

  my_last_move_ = std::nullopt;
  after_my_turn_observation_ = std::nullopt;
  after_opp_turn_observations_.clear();
}

void AOSPInfoset::AddState(chess::ChessBoard board) {

  std::unordered_set<chess::ChessBoard, ChessBoardHash> *pool =
      pool_ == 0 ?
          pool = &pool0_ :
          pool = &pool1_;
  if (!pool->count(board)) {
    v_pool_.push_back(board);
    pool->insert(board);
  }
}

void AOSPInfoset::Inform(const chess::Move& move, const kriegspiel::KriegspielUmpireMessage& message) {
  if (my_last_move_ || after_my_turn_observation_) {
    SpielFatalError("Not updated before informing about another my move");
  }
  my_last_move_ = move;
  after_my_turn_observation_ = message;
}

void AOSPInfoset::Inform(const kriegspiel::KriegspielUmpireMessage& message, Player player) {
  if (!message.illegal) {
    after_opp_turn_observations_.emplace_back(std::make_pair(player, message));
  }
}

void AOSPInfoset::Update() {
  if (my_last_move_ && after_my_turn_observation_) {
    Update(*my_last_move_, *after_my_turn_observation_);
  }

  SPIEL_CHECK_LE(after_opp_turn_observations_.size(), 1);
  for (auto &[player, message] : after_opp_turn_observations_) {
    Update(message, player);
  }

  my_last_move_ = std::nullopt;
  after_my_turn_observation_ = std::nullopt;
  after_opp_turn_observations_.clear();
}

void AOSPInfoset::Update(const chess::Move& move,
                         const kriegspiel::KriegspielUmpireMessage& message) {
  std::unordered_set<chess::ChessBoard, ChessBoardHash> *orig_pool;
  std::unordered_set<chess::ChessBoard, ChessBoardHash> *new_pool;
  if (pool_ == 0) {
    orig_pool = &pool0_;
    new_pool = &pool1_;
  } else {
    orig_pool = &pool1_;
    new_pool = &pool0_;
  }

  new_pool->clear();
  for (chess::ChessBoard board : *orig_pool) {

    if (board.IsMoveLegal(move) == message.illegal) {
      continue;
    }
    if (message.illegal) {
      new_pool->insert(board);
      continue;
    }

    chess::PieceType capture_type = board.at(move.to).type;
    bool capture_type_checks_out = true;
    switch (capture_type) {
      case chess::PieceType::kEmpty:
        if (message.capture_type != kriegspiel::KriegspielCaptureType::kNoCapture) {
          capture_type_checks_out = false;
        }
        break;
      case chess::PieceType::kPawn:
        if (message.capture_type != kriegspiel::KriegspielCaptureType::kPawn) {
          capture_type_checks_out = false;
        }
        break;
      default:
        if (message.capture_type != kriegspiel::KriegspielCaptureType::kPiece) {
          capture_type_checks_out = false;
        }
    }
    if (!capture_type_checks_out) {
      continue;
    }

    chess::ChessBoard board_copy = board;
    board_copy.ApplyMove(move);

    if (IsTerminal(board_copy)) {
      continue;
    }

    if (message.check_types != kriegspiel::GetCheckType(board_copy)) {
      continue;
    }

    int pawnTries = 0;
    board_copy.GenerateLegalPawnCaptures([&pawnTries](const chess::Move &move) {
      pawnTries++;
      return true;
    }, board_copy.ToPlay());
    if (message.pawn_tries != pawnTries) {
      continue;
    }
    if (message.to_move != board_copy.ToPlay()) {
      continue;
    }
    new_pool->insert(board_copy);
  }
  v_pool_.clear();
  for (chess::ChessBoard board : *new_pool) v_pool_.push_back(board);
  pool_ = 1 - pool_;
}

void AOSPInfoset::Update(const kriegspiel::KriegspielUmpireMessage& message,
                         Player player) {
  std::unordered_set<chess::ChessBoard, ChessBoardHash> *orig_pool;
  std::unordered_set<chess::ChessBoard, ChessBoardHash> *new_pool;
  if (pool_ == 0) {
    orig_pool = &pool0_;
    new_pool = &pool1_;
  } else {
    orig_pool = &pool1_;
    new_pool = &pool0_;
  }

  new_pool->clear();
  for (const chess::ChessBoard &board : *orig_pool) {

    board.GenerateLegalMoves([&board, &message, &new_pool](const chess::Move &move) {

      // todo optimize generate pseudolegal, check legal and use the board_copy

      SPIEL_CHECK_FALSE(message.illegal);

      chess::PieceType capture_type = board.at(move.to).type;
      switch (capture_type) {
        case chess::PieceType::kEmpty:
          if (message.capture_type != kriegspiel::KriegspielCaptureType::kNoCapture) {
            return true;
          }
          break;
        case chess::PieceType::kPawn:
          if (message.capture_type != kriegspiel::KriegspielCaptureType::kPawn || move.to != message.square) {
            return true;
          }
          break;
        default:
          if (message.capture_type != kriegspiel::KriegspielCaptureType::kPiece || move.to != message.square) {
            return true;
          }
      }

      chess::ChessBoard board_copy = board;
      board_copy.ApplyMove(move);

      if (IsTerminal(board_copy)) {
        return true;
      }

      if (message.check_types != kriegspiel::GetCheckType(board_copy)) {
        return true;
      }

      int pawnTries = 0;
      board_copy.GenerateLegalPawnCaptures([&pawnTries](const chess::Move &move) {
        pawnTries++;
        return true;
      }, board_copy.ToPlay());
      if (message.pawn_tries != pawnTries) {
        return true;
      }
      if (message.to_move != board_copy.ToPlay()) {
        return true;
      }
      new_pool->insert(board_copy);

      return true;
    });
  }

  // trimm the pool
  if (new_pool->size() > pool_max_size_) {

    orig_pool->clear();
    size_t pool_size = new_pool->size();
    for (const chess::ChessBoard &board : *new_pool) {
      int r = absl::Uniform(rng_, 0u, pool_size);
      if (r < pool_max_size_) {
        orig_pool->insert(board);
      }
    }
    v_pool_.clear();
    for (const chess::ChessBoard &board : *orig_pool) v_pool_.push_back(board);
  } else {
    v_pool_.clear();
    for (const chess::ChessBoard &board : *new_pool) v_pool_.push_back(board);
    pool_ = 1 - pool_;
  }
}


LOSInfoset::LOSInfoset(int seed, chess::Color color)
    : rng_(seed),
      color_(color) {
  stats_ = ReadStatistics();

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

struct PieceProb {
  chess::PieceType piece_type;
  chess::Square sq;
  double prob;
};

PieceProb RandomPieceType(const std::vector<PieceProb>& probs) {

  if (probs.empty()) {
    SpielFatalError("there are no options for a piece there");
  }

  double sum = 0.0;

  for (const PieceProb& prob : probs) {
    sum += prob.prob;
  }


  double r = ((double) rand()/RAND_MAX) * sum;

  double sum2 = 0.0;
  for (const PieceProb& prob : probs) {
    sum2 += prob.prob;
    if (sum2 >= r) {
      return prob;
    }
  }
  return probs[probs.size() - 1];
}

chess::ChessBoard LOSInfoset::SampleState() {

  std::string fen = board_.ToFEN();

  int pieces_to_place = 7 - captured_pieces_;
  int pawns_to_place = 8 - captured_pawns_;
  int rooks_to_place = 2;
  int knights_to_place = 2;
  int bishops_to_place = 2;
  int queens_to_place = 1;
  int kings_to_place = 1;

  chess::Color opp_color = chess::OppColor(color_);
  Player opp = chess::ColorToPlayer(opp_color);

  std::array<std::array<double, 8>, 8> king_dist_ {};
  std::array<std::array<double, 8>, 8> pawn_dist_ {};
  std::array<std::array<double, 8>, 8> rook_dist_ {};
  std::array<std::array<double, 8>, 8> knight_dist_ {};
  std::array<std::array<double, 8>, 8> bishop_dist_ {};
  std::array<std::array<double, 8>, 8> queen_dist_ {};

  chess::ChessBoard result_board = board_;

  for (int8_t x = 0; x < chess::kMaxBoardSize; ++x) {
    for (int8_t y = 0; y < chess::kMaxBoardSize; ++y) {
      chess::Square sq = {x, y};
      if (result_board.at(sq).type == chess::PieceType::kEmpty) {
        king_dist_[x][y] = ((double) stats_.king_counts[opp][ply_][x][y]) / stats_.counts[opp][ply_];
        pawn_dist_[x][y] = ((double) stats_.pawn_counts[opp][ply_][x][y]) / stats_.counts[opp][ply_];
        rook_dist_[x][y] = ((double) stats_.rook_counts[opp][ply_][x][y]) / stats_.counts[opp][ply_];
        knight_dist_[x][y] = ((double) stats_.knight_counts[opp][ply_][x][y]) / stats_.counts[opp][ply_];
        bishop_dist_[x][y] = ((double) stats_.bishop_counts[opp][ply_][x][y]) / stats_.counts[opp][ply_];
        queen_dist_[x][y] = ((double) stats_.queen_counts[opp][ply_][x][y]) / stats_.counts[opp][ply_];
      }
    }
  }

  // fill with the king
  if (kings_to_place > 0) {
    std::vector<PieceProb> probs;
    for (int8_t x = 0; x < chess::kMaxBoardSize; ++x) {
      for (int8_t y = 0; y < chess::kMaxBoardSize; ++y) {
        chess::Square sq = {x, y};
        if (result_board.UnderAttack(sq, opp_color)) {
          continue;
        }
        if (king_dist_[x][y] > 0.0) {
          probs.push_back({chess::PieceType::kKing, sq, king_dist_[x][y]});
        }
      }
    }
    if (!probs.empty()) {

      PieceProb result_piece = RandomPieceType(probs);
      chess::Square sq = result_piece.sq;
      result_board.set_square(sq, {opp_color, result_piece.piece_type});
      king_dist_[sq.x][sq.y] = 0.0;
      pawn_dist_[sq.x][sq.y] = 0.0;
      rook_dist_[sq.x][sq.y] = 0.0;
      knight_dist_[sq.x][sq.y] = 0.0;
      bishop_dist_[sq.x][sq.y] = 0.0;
      queen_dist_[sq.x][sq.y] = 0.0;

      kings_to_place--;
    }

  }

  if (my_last_move_ && after_my_turn_observation_) {
    SPIEL_CHECK_TRUE(after_my_turn_observation_->illegal);

    if (my_last_move_->piece.type == chess::PieceType::kPawn) {

      if (my_last_move_->to.x != my_last_move_->from.x) {
        int8_t x = my_last_move_->to.x;
        int8_t y = my_last_move_->to.y;
        king_dist_[x][y] = 0.0;
        pawn_dist_[x][y] = 0.0;
        rook_dist_[x][y] = 0.0;
        knight_dist_[x][y] = 0.0;
        bishop_dist_[x][y] = 0.0;
        queen_dist_[x][y] = 0.0;
      }
      else {
        chess::Square sq = my_last_move_->to;
        if (sq == chess::kInvalidSquare) {
          int a = 5;
        }
        std::vector<PieceProb> probs;
        if (pawns_to_place > 0) {
          probs.push_back({chess::PieceType::kPawn, sq, pawn_dist_[sq.x][sq.y]});
        }
        if (pieces_to_place > 0) {
          if (rooks_to_place > 0) {
            probs.push_back({chess::PieceType::kRook, sq, rook_dist_[sq.x][sq.y]});
          }
          if (knights_to_place > 0) {
            probs.push_back({chess::PieceType::kKnight, sq, knight_dist_[sq.x][sq.y]});
          }
          if (bishops_to_place > 0) {
            probs.push_back({chess::PieceType::kBishop, sq, bishop_dist_[sq.x][sq.y]});
          }
          if (queens_to_place > 0) {
            probs.push_back({chess::PieceType::kQueen, sq, queen_dist_[sq.x][sq.y]});
          }
        }
        if (kings_to_place > 0 && !result_board.UnderAttack(sq, opp_color)) {
          probs.push_back({chess::PieceType::kKing, sq, king_dist_[sq.x][sq.y]});
        }

        PieceProb result_piece = RandomPieceType(probs);

        switch (result_piece.piece_type) {
          case chess::PieceType::kPawn:
            pawns_to_place--;
            break;
          case chess::PieceType::kKing:
            kings_to_place--;
            break;
          case chess::PieceType::kQueen:
            queens_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kRook:
            rooks_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kBishop:
            bishops_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kKnight:
            knights_to_place--;
            pieces_to_place--;
            break;
        }
        king_dist_[sq.x][sq.y] = 0.0;
        pawn_dist_[sq.x][sq.y] = 0.0;
        rook_dist_[sq.x][sq.y] = 0.0;
        knight_dist_[sq.x][sq.y] = 0.0;
        bishop_dist_[sq.x][sq.y] = 0.0;
        queen_dist_[sq.x][sq.y] = 0.0;

        result_board.set_square(sq, {opp_color, result_piece.piece_type});
      }

    }
    else if (my_last_move_->piece.type == chess::PieceType::kKing) {
      std::array<chess::Offset, 8> knight_offsets = {chess::Offset{1, 2},
                                                     chess::Offset{1, -2},
                                                     chess::Offset{-1, 2},
                                                     chess::Offset{-1, -2},
                                                     chess::Offset{2, 1},
                                                     chess::Offset{2, -1},
                                                     chess::Offset{-2, 1},
                                                     chess::Offset{-2, -1}};

      std::array<chess::Offset, 8> diagonal_offsets = {chess::Offset{1, 1},
                                                       chess::Offset{-1, -1},
                                                       chess::Offset{1, -1},
                                                       chess::Offset{-1, 1}};

      std::array<chess::Offset, 8> nondiagonal_offsets = {chess::Offset{1, 0},
                                                          chess::Offset{-1, 0},
                                                          chess::Offset{0, 1},
                                                          chess::Offset{0, -1}};


      std::vector<PieceProb> probs;
      if (knights_to_place > 0) {
        for (const chess::Offset &offset : knight_offsets) {
          chess::Square sq = my_last_move_->to + offset;
          if (result_board.InBoardArea(sq) && result_board.at(sq) == chess::kEmptyPiece) {

            probs.push_back({chess::PieceType::kKnight, sq,
                             knight_dist_[sq.x][sq.y]});
          }
        }
      }

      // todo add pawns
      for (const chess::Offset &offset : diagonal_offsets) {
        chess::Square sq = my_last_move_->to + offset;
        while (result_board.InBoardArea(sq) && result_board.at(sq) == chess::kEmptyPiece) {

          if (bishops_to_place > 0) {
            probs.push_back({chess::PieceType::kBishop, sq,
                             bishop_dist_[sq.x][sq.y]});
          }
          if (queens_to_place > 0) {
            probs.push_back({chess::PieceType::kQueen, sq,
                             rook_dist_[sq.x][sq.y]});
          }
          if (kings_to_place > 0 && !result_board.UnderAttack(sq, opp_color)) {
            probs.push_back({chess::PieceType::kKing, sq,
                             king_dist_[sq.x][sq.y]});
          }

          sq += offset;
        }
      }

      for (const chess::Offset &offset : nondiagonal_offsets) {
        chess::Square sq = my_last_move_->to + offset;
        while (result_board.InBoardArea(sq) && result_board.at(sq) == chess::kEmptyPiece) {

          if (rooks_to_place > 0) {
            probs.push_back({chess::PieceType::kBishop, sq,
                             bishop_dist_[sq.x][sq.y]});
          }
          if (queens_to_place > 0) {
            probs.push_back({chess::PieceType::kQueen, sq,
                             rook_dist_[sq.x][sq.y]});
          }
          if (kings_to_place > 0 && !result_board.UnderAttack(sq, opp_color)) {
            probs.push_back({chess::PieceType::kKing, sq,
                             king_dist_[sq.x][sq.y]});
          }

          sq += offset;
        }
      }


      PieceProb result_piece = RandomPieceType(probs);
      chess::Square sq = result_piece.sq;

      switch (result_piece.piece_type) {
        case chess::PieceType::kPawn:
          pawns_to_place--;
          break;
        case chess::PieceType::kKing:
          kings_to_place--;
          break;
        case chess::PieceType::kQueen:
          queens_to_place--;
          pieces_to_place--;
          break;
        case chess::PieceType::kRook:
          rooks_to_place--;
          pieces_to_place--;
          break;
        case chess::PieceType::kBishop:
          bishops_to_place--;
          pieces_to_place--;
          break;
        case chess::PieceType::kKnight:
          knights_to_place--;
          pieces_to_place--;
          break;
      }
      king_dist_[sq.x][sq.y] = 0.0;
      pawn_dist_[sq.x][sq.y] = 0.0;
      rook_dist_[sq.x][sq.y] = 0.0;
      knight_dist_[sq.x][sq.y] = 0.0;
      bishop_dist_[sq.x][sq.y] = 0.0;
      queen_dist_[sq.x][sq.y] = 0.0;

      result_board.set_square(sq, {opp_color, result_piece.piece_type});

    }
    else {
      // todo it also could be reveal check and not trying to jump over someone
      if (my_last_move_->piece.type != chess::PieceType::kKnight) {
        int offset_x = 0;
        int offset_y = 0;

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

        int diff_y = my_last_move_->to.y - my_last_move_->from.y;
        if (diff_y > 0)
          offset_y = 1;
        else if (diff_y < 0)
          offset_y = -1;
        chess::Offset offset_step = {int8_t(offset_x), int8_t(offset_y)};

        std::vector<PieceProb> probs;
        for (chess::Square dest = my_last_move_->from + offset_step; dest != my_last_move_->to;
             dest += offset_step) {

          if (pawns_to_place > 0) {
            probs.push_back({chess::PieceType::kPawn, dest, pawn_dist_[dest.x][dest.y]});
          }
          if (pieces_to_place > 0) {
            if (rooks_to_place > 0) {
              probs.push_back({chess::PieceType::kRook, dest, rook_dist_[dest.x][dest.y]});
            }
            if (knights_to_place > 0) {
              probs.push_back({chess::PieceType::kKnight, dest, knight_dist_[dest.x][dest.y]});
            }
            if (bishops_to_place > 0) {
              probs.push_back({chess::PieceType::kBishop, dest, bishop_dist_[dest.x][dest.y]});
            }
            if (queens_to_place > 0) {
              probs.push_back({chess::PieceType::kQueen, dest, queen_dist_[dest.x][dest.y]});
            }
          }

        }

        PieceProb result_piece = RandomPieceType(probs);
        chess::Square sq = result_piece.sq;

        switch (result_piece.piece_type) {
          case chess::PieceType::kPawn:
            pawns_to_place--;
            break;
          case chess::PieceType::kKing:
            kings_to_place--;
            break;
          case chess::PieceType::kQueen:
            queens_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kRook:
            rooks_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kBishop:
            bishops_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kKnight:
            knights_to_place--;
            pieces_to_place--;
            break;
        }
        king_dist_[sq.x][sq.y] = 0.0;
        pawn_dist_[sq.x][sq.y] = 0.0;
        rook_dist_[sq.x][sq.y] = 0.0;
        knight_dist_[sq.x][sq.y] = 0.0;
        bishop_dist_[sq.x][sq.y] = 0.0;
        queen_dist_[sq.x][sq.y] = 0.0;

        result_board.set_square(sq, {opp_color, result_piece.piece_type});

      }
    }

  }

  if (after_opp_turn_observation_) {
    if (after_opp_turn_observation_->capture_type != kriegspiel::KriegspielCaptureType::kNoCapture) {
      chess::Square sq = after_opp_turn_observation_->square;
      if (sq == chess::kInvalidSquare) {
        int a = 5;
      }
      std::vector<PieceProb> probs;
      if (pawns_to_place > 0) {
        probs.push_back({chess::PieceType::kPawn, sq, pawn_dist_[sq.x][sq.y]});
      }
      if (pieces_to_place > 0) {
        if (rooks_to_place > 0) {
          probs.push_back({chess::PieceType::kRook, sq, rook_dist_[sq.x][sq.y]});
        }
        if (knights_to_place > 0) {
          probs.push_back({chess::PieceType::kKnight, sq, knight_dist_[sq.x][sq.y]});
        }
        if (bishops_to_place > 0) {
          probs.push_back({chess::PieceType::kBishop, sq, bishop_dist_[sq.x][sq.y]});
        }
        if (queens_to_place > 0) {
          probs.push_back({chess::PieceType::kQueen, sq, queen_dist_[sq.x][sq.y]});
        }
      }
      if (kings_to_place > 0 && !result_board.UnderAttack(sq, opp_color)) {
        probs.push_back({chess::PieceType::kKing, sq, king_dist_[sq.x][sq.y]});
      }

      if (!probs.empty()) {
        PieceProb result_piece = RandomPieceType(probs);

        switch (result_piece.piece_type) {
          case chess::PieceType::kPawn:
            pawns_to_place--;
            break;
          case chess::PieceType::kKing:
            kings_to_place--;
            break;
          case chess::PieceType::kQueen:
            queens_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kRook:
            rooks_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kBishop:
            bishops_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kKnight:
            knights_to_place--;
            pieces_to_place--;
            break;
        }
        king_dist_[sq.x][sq.y] = 0.0;
        pawn_dist_[sq.x][sq.y] = 0.0;
        rook_dist_[sq.x][sq.y] = 0.0;
        knight_dist_[sq.x][sq.y] = 0.0;
        bishop_dist_[sq.x][sq.y] = 0.0;
        queen_dist_[sq.x][sq.y] = 0.0;


        result_board.set_square(sq, {opp_color, result_piece.piece_type});

      }
    }

    if (after_opp_turn_observation_->check_types.first != kriegspiel::KriegspielCheckType::kNoCheck) {
      chess::Square king_sq = result_board.find({color_, chess::PieceType::kKing});
      kriegspiel::KriegspielCheckType check_type = after_opp_turn_observation_->check_types.first;

      std::vector<PieceProb> probs;

      if (check_type == kriegspiel::KriegspielCheckType::kKnight) {

        std::array<chess::Offset, 8> offsets = {chess::Offset{1, 2},
                                                chess::Offset{1, -2},
                                                chess::Offset{-1, 2},
                                                chess::Offset{-1, -2},
                                                chess::Offset{2, 1},
                                                chess::Offset{2, -1},
                                                chess::Offset{-2, 1},
                                                chess::Offset{-2, -1}};

        for (const chess::Offset &offset : offsets) {
          chess::Square sq = king_sq + offset;
          if (result_board.InBoardArea(sq)
              && result_board.at(sq) == chess::kEmptyPiece) {

            probs.push_back({chess::PieceType::kKnight, sq,
                             rook_dist_[sq.x][sq.y]});

          }
        }


      }
      else if (check_type == kriegspiel::KriegspielCheckType::kRank) {
        std::array<chess::Offset, 2> offsets = {chess::Offset{1, 0},
                                                chess::Offset{-1, 0}};

        for (const chess::Offset &offset : offsets) {
          chess::Square sq = king_sq + offset;
          while (result_board.InBoardArea(sq)
                 && result_board.at(sq) == chess::kEmptyPiece) {

            if (rooks_to_place > 0) {
              probs.push_back({chess::PieceType::kRook, sq,
                               rook_dist_[sq.x][sq.y]});
            }
            if (queens_to_place > 0) {
              probs.push_back({chess::PieceType::kQueen, sq,
                               rook_dist_[sq.x][sq.y]});
            }

            sq += offset;
          }
        }
      }
      else if (check_type == kriegspiel::KriegspielCheckType::kFile) {
        std::array<chess::Offset, 2> offsets = {chess::Offset{0, 1},
                                                chess::Offset{0, -1}};

        for (const chess::Offset &offset : offsets) {
          chess::Square sq = king_sq + offset;
          while (result_board.InBoardArea(sq)
                 && result_board.at(sq) == chess::kEmptyPiece) {

            if (rooks_to_place > 0) {
              probs.push_back({chess::PieceType::kRook, sq,
                               rook_dist_[sq.x][sq.y]});
            }
            if (queens_to_place > 0) {
              probs.push_back({chess::PieceType::kQueen, sq,
                               rook_dist_[sq.x][sq.y]});
            }

            sq += offset;
          }
        }
      }
      else {
        bool right_left_is_long = chess::IsLongDiagonal(king_sq,
                                                        king_sq + chess::Offset{1, 1},
                                                        8);
        std::array<chess::Offset, 2> offsets {};
        if (right_left_is_long) {
          offsets = {chess::Offset{1, 1}, chess::Offset{-1, -1}};
        }
        else {
          offsets = {chess::Offset{1, -1}, chess::Offset{-1, 1}};
        }

        for (const chess::Offset &offset : offsets) {
          chess::Square sq = king_sq + offset;
          while (result_board.InBoardArea(sq)
                 && result_board.at(sq) == chess::kEmptyPiece) {

            if (bishops_to_place > 0) {
              probs.push_back({chess::PieceType::kBishop, sq,
                               bishop_dist_[sq.x][sq.y]});
            }
            if (queens_to_place > 0) {
              probs.push_back({chess::PieceType::kQueen, sq,
                               rook_dist_[sq.x][sq.y]});
            }

            sq += offset;
          }
        }

      }

      if (!probs.empty()) {

        PieceProb result_piece = RandomPieceType(probs);
        chess::Square sq = result_piece.sq;

        switch (result_piece.piece_type) {
          case chess::PieceType::kPawn:
            pawns_to_place--;
            break;
          case chess::PieceType::kKing:
            kings_to_place--;
            break;
          case chess::PieceType::kQueen:
            queens_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kRook:
            rooks_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kBishop:
            bishops_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kKnight:
            knights_to_place--;
            pieces_to_place--;
            break;
        }
        king_dist_[sq.x][sq.y] = 0.0;
        pawn_dist_[sq.x][sq.y] = 0.0;
        rook_dist_[sq.x][sq.y] = 0.0;
        knight_dist_[sq.x][sq.y] = 0.0;
        bishop_dist_[sq.x][sq.y] = 0.0;
        queen_dist_[sq.x][sq.y] = 0.0;


        result_board.set_square(sq, {opp_color, result_piece.piece_type});
      }

    }

    if (after_opp_turn_observation_->check_types.second != kriegspiel::KriegspielCheckType::kNoCheck) {

      chess::Square king_sq = result_board.find({color_, chess::PieceType::kKing});
      kriegspiel::KriegspielCheckType check_type = after_opp_turn_observation_->check_types.second;

      std::vector<PieceProb> probs;

      if (check_type == kriegspiel::KriegspielCheckType::kKnight) {

        std::array<chess::Offset, 8> offsets = {chess::Offset{1, 2},
                                                chess::Offset{1, -2},
                                                chess::Offset{-1, 2},
                                                chess::Offset{-1, -2},
                                                chess::Offset{2, 1},
                                                chess::Offset{2, -1},
                                                chess::Offset{-2, 1},
                                                chess::Offset{-2, -1}};

        for (const chess::Offset &offset : offsets) {
          chess::Square sq = king_sq + offset;
          if (result_board.InBoardArea(sq)
              && result_board.at(sq) == chess::kEmptyPiece) {

            probs.push_back({chess::PieceType::kKnight, sq,
                             rook_dist_[sq.x][sq.y]});

          }
        }

      }
      else if (check_type == kriegspiel::KriegspielCheckType::kRank) {
        std::array<chess::Offset, 2> offsets = {chess::Offset{1, 0},
                                                chess::Offset{-1, 0}};

        for (const chess::Offset &offset : offsets) {
          chess::Square sq = king_sq + offset;
          while (result_board.InBoardArea(sq)
                 && result_board.at(sq) == chess::kEmptyPiece) {

            probs.push_back({chess::PieceType::kRook, sq,
                             rook_dist_[sq.x][sq.y]});
            probs.push_back({chess::PieceType::kQueen, sq,
                             rook_dist_[sq.x][sq.y]});

            sq += offset;
          }
        }
      }
      else if (check_type == kriegspiel::KriegspielCheckType::kFile) {
        std::array<chess::Offset, 2> offsets = {chess::Offset{0, 1},
                                                chess::Offset{0, -1}};

        for (const chess::Offset &offset : offsets) {
          chess::Square sq = king_sq + offset;
          while (result_board.InBoardArea(sq)
                 && result_board.at(sq) == chess::kEmptyPiece) {

            probs.push_back({chess::PieceType::kRook, sq,
                             rook_dist_[sq.x][sq.y]});
            probs.push_back({chess::PieceType::kQueen, sq,
                             rook_dist_[sq.x][sq.y]});

            sq += offset;
          }
        }
      }
      else {
        bool right_left_is_long = chess::IsLongDiagonal(king_sq,
                                                        king_sq + chess::Offset{1, 1},
                                                        8);
        std::array<chess::Offset, 2> offsets {};
        if (right_left_is_long) {
          offsets = {chess::Offset{1, 1}, chess::Offset{-1, -1}};
        }
        else {
          offsets = {chess::Offset{1, -1}, chess::Offset{-1, 1}};
        }

        for (const chess::Offset &offset : offsets) {
          chess::Square sq = king_sq + offset;
          while (result_board.InBoardArea(sq)
                 && result_board.at(sq) == chess::kEmptyPiece) {

            probs.push_back({chess::PieceType::kBishop, sq,
                             bishop_dist_[sq.x][sq.y]});
            probs.push_back({chess::PieceType::kQueen, sq,
                             rook_dist_[sq.x][sq.y]});

            sq += offset;
          }
        }

      }

      if (!probs.empty()) {
        PieceProb result_piece = RandomPieceType(probs);
        chess::Square sq = result_piece.sq;

        switch (result_piece.piece_type) {
          case chess::PieceType::kPawn:
            pawns_to_place--;
            break;
          case chess::PieceType::kKing:
            kings_to_place--;
            break;
          case chess::PieceType::kQueen:
            queens_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kRook:
            rooks_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kBishop:
            bishops_to_place--;
            pieces_to_place--;
            break;
          case chess::PieceType::kKnight:
            knights_to_place--;
            pieces_to_place--;
            break;
        }
        king_dist_[sq.x][sq.y] = 0.0;
        pawn_dist_[sq.x][sq.y] = 0.0;
        rook_dist_[sq.x][sq.y] = 0.0;
        knight_dist_[sq.x][sq.y] = 0.0;
        bishop_dist_[sq.x][sq.y] = 0.0;
        queen_dist_[sq.x][sq.y] = 0.0;

        result_board.set_square(sq, {opp_color, result_piece.piece_type});
      }

    }


    std::vector<chess::Move> legal_pawn_captures;
    result_board.GeneratePseudoLegalPawnCaptures([&legal_pawn_captures] (const chess::Move& move) {
      legal_pawn_captures.push_back(move);
      return true;
    }, color_, true);

    if (after_opp_turn_observation_->pawn_tries == 0) {
      for (const chess::Move& move : legal_pawn_captures) {
        king_dist_[move.to.x][move.to.y] = 0.0;
        pawn_dist_[move.to.x][move.to.y] = 0.0;
        rook_dist_[move.to.x][move.to.y] = 0.0;
        knight_dist_[move.to.x][move.to.y] = 0.0;
        bishop_dist_[move.to.x][move.to.y] = 0.0;
        queen_dist_[move.to.x][move.to.y] = 0.0;
      }
    }
    else {

      int pawn_tries = 0;
      for (const chess::Move& move : legal_pawn_captures) {
        if (result_board.at(move.to) != chess::kEmptyPiece) {
          pawn_tries++;
        }
      }

      for (int i = 0; i < after_opp_turn_observation_->pawn_tries - pawn_tries; ++i) {
        std::vector<PieceProb> probs;
        for (const chess::Move& move : legal_pawn_captures) {
          if (result_board.at(move.to) != chess::kEmptyPiece) {
            continue;
          }
          chess::Square sq = move.to;

          if (pawns_to_place > 0) {
            probs.push_back({chess::PieceType::kPawn, sq, pawn_dist_[sq.x][sq.y]});
          }
          if (pieces_to_place > 0) {
            if (rooks_to_place > 0) {
              probs.push_back({chess::PieceType::kRook, sq, rook_dist_[sq.x][sq.y]});
            }
            if (knights_to_place > 0) {
              probs.push_back({chess::PieceType::kKnight, sq, knight_dist_[sq.x][sq.y]});
            }
            if (bishops_to_place > 0) {
              probs.push_back({chess::PieceType::kBishop, sq, bishop_dist_[sq.x][sq.y]});
            }
            if (queens_to_place > 0) {
              probs.push_back({chess::PieceType::kQueen, sq, queen_dist_[sq.x][sq.y]});
            }
          }

        }

        if (!probs.empty()) {
          PieceProb result_piece = RandomPieceType(probs);
          chess::Square sq = result_piece.sq;

          switch (result_piece.piece_type) {
            case chess::PieceType::kPawn:
              pawns_to_place--;
              break;
            case chess::PieceType::kKing:
              kings_to_place--;
              break;
            case chess::PieceType::kQueen:
              queens_to_place--;
              pieces_to_place--;
              break;
            case chess::PieceType::kRook:
              rooks_to_place--;
              pieces_to_place--;
              break;
            case chess::PieceType::kBishop:
              bishops_to_place--;
              pieces_to_place--;
              break;
            case chess::PieceType::kKnight:
              knights_to_place--;
              pieces_to_place--;
              break;
          }
          king_dist_[sq.x][sq.y] = 0.0;
          pawn_dist_[sq.x][sq.y] = 0.0;
          rook_dist_[sq.x][sq.y] = 0.0;
          knight_dist_[sq.x][sq.y] = 0.0;
          bishop_dist_[sq.x][sq.y] = 0.0;
          queen_dist_[sq.x][sq.y] = 0.0;

          result_board.set_square(sq, {opp_color, result_piece.piece_type});
        }
      }
    }

  }

  // fill with pawns
  for (int i = 0; i < pawns_to_place; ++i) {
    std::vector<PieceProb> probs;
    for (int8_t x = 0; x < chess::kMaxBoardSize; ++x) {
      for (int8_t y = 0; y < chess::kMaxBoardSize; ++y) {
        chess::Square sq = {x, y};
        if (pawn_dist_[x][y] > 0.0) {
          probs.push_back({chess::PieceType::kPawn, sq, pawn_dist_[x][y]});
        }
      }
    }

    if (probs.empty()) {
      continue;
    }
    PieceProb result_piece = RandomPieceType(probs);
    chess::Square sq = result_piece.sq;
    result_board.set_square(sq, {opp_color, result_piece.piece_type});
    king_dist_[sq.x][sq.y] = 0.0;
    pawn_dist_[sq.x][sq.y] = 0.0;
    rook_dist_[sq.x][sq.y] = 0.0;
    knight_dist_[sq.x][sq.y] = 0.0;
    bishop_dist_[sq.x][sq.y] = 0.0;
    queen_dist_[sq.x][sq.y] = 0.0;


  }

  // fill with pieces
  for (int i = 0; i < pieces_to_place; ++i) {
    std::vector<PieceProb> probs;
    for (int8_t x = 0; x < chess::kMaxBoardSize; ++x) {
      for (int8_t y = 0; y < chess::kMaxBoardSize; ++y) {
        chess::Square sq = {x, y};
        if (pawns_to_place > 0) {
          probs.push_back({chess::PieceType::kPawn, sq, pawn_dist_[sq.x][sq.y]});
        }
        if (pieces_to_place > 0) {
          if (rooks_to_place > 0) {
            probs.push_back({chess::PieceType::kRook, sq, rook_dist_[sq.x][sq.y]});
          }
          if (knights_to_place > 0) {
            probs.push_back({chess::PieceType::kKnight, sq, knight_dist_[sq.x][sq.y]});
          }
          if (bishops_to_place > 0) {
            probs.push_back({chess::PieceType::kBishop, sq, bishop_dist_[sq.x][sq.y]});
          }
          if (queens_to_place > 0) {
            probs.push_back({chess::PieceType::kQueen, sq, queen_dist_[sq.x][sq.y]});
          }
        }

      }
    }

    if (probs.empty()) {
      continue;
    }
    PieceProb result_piece = RandomPieceType(probs);

    switch (result_piece.piece_type) {
      case chess::PieceType::kPawn:
        pawns_to_place--;
        break;
      case chess::PieceType::kKing:
        kings_to_place--;
        break;
      case chess::PieceType::kQueen:
        queens_to_place--;
        pieces_to_place--;
        break;
      case chess::PieceType::kRook:
        rooks_to_place--;
        pieces_to_place--;
        break;
      case chess::PieceType::kBishop:
        bishops_to_place--;
        pieces_to_place--;
        break;
      case chess::PieceType::kKnight:
        knights_to_place--;
        pieces_to_place--;
        break;
    }
    chess::Square sq = result_piece.sq;
    king_dist_[sq.x][sq.y] = 0.0;
    pawn_dist_[sq.x][sq.y] = 0.0;
    rook_dist_[sq.x][sq.y] = 0.0;
    knight_dist_[sq.x][sq.y] = 0.0;
    bishop_dist_[sq.x][sq.y] = 0.0;
    queen_dist_[sq.x][sq.y] = 0.0;


    result_board.set_square(sq, {opp_color, result_piece.piece_type});
  }

  // fill with the king
  if (kings_to_place > 0) {
    std::vector<PieceProb> probs;
    for (int8_t x = 0; x < chess::kMaxBoardSize; ++x) {
      for (int8_t y = 0; y < chess::kMaxBoardSize; ++y) {
        chess::Square sq = {x, y};
        if (result_board.UnderAttack(sq, opp_color)) {
          continue;
        }
        if (king_dist_[x][y] > 0.0) {
          probs.push_back({chess::PieceType::kKing, sq, king_dist_[x][y]});
        }
      }
    }
    if (probs.empty()) {
      return result_board;
    }
    PieceProb result_piece = RandomPieceType(probs);
    chess::Square sq = result_piece.sq;
    result_board.set_square(sq, {opp_color, result_piece.piece_type});
    king_dist_[sq.x][sq.y] = 0.0;
    pawn_dist_[sq.x][sq.y] = 0.0;
    rook_dist_[sq.x][sq.y] = 0.0;
    knight_dist_[sq.x][sq.y] = 0.0;
    bishop_dist_[sq.x][sq.y] = 0.0;
    queen_dist_[sq.x][sq.y] = 0.0;

  }

  return result_board;
}

void LOSInfoset::Reset() {

  my_last_move_ = std::nullopt;
  after_my_turn_observation_ = std::nullopt;
  after_opp_turn_observation_ = std::nullopt;

  std::string FEN = kPartialFENs[chess::ColorToPlayer(color_)];
  std::optional<chess::ChessBoard> board = chess::ChessBoard::BoardFromFEN(FEN);
  SPIEL_CHECK_TRUE(board);
  board_ = *board;
  captured_pieces_ = 0;
  captured_pawns_ = 0;
  ply_ = 0;
}

void LOSInfoset::Inform(const chess::Move& move, const kriegspiel::KriegspielUmpireMessage& message) {

  my_last_move_ = move;
  after_my_turn_observation_ = message;
//  std::cout << "before" << std::endl;
//  std::cout << board_.DebugString() << std::endl;
  if (!message.illegal) {
    ply_++;
    board_.ApplyMove(move);
    after_opp_turn_observation_ = std::nullopt;
    if (message.capture_type == kriegspiel::KriegspielCaptureType::kPawn) {
      captured_pawns_++;
    }
    else if (message.capture_type == kriegspiel::KriegspielCaptureType::kPiece) {
      captured_pieces_++;
    }
    SPIEL_CHECK_TRUE(board_.at(move.to).type != chess::PieceType::kEmpty);
    SPIEL_CHECK_TRUE(board_.at(move.from).type == chess::PieceType::kEmpty);
  }
//  std::cout << "after" << std::endl;
//  std::cout << board_.DebugString() << std::endl;
}

void LOSInfoset::Inform(const kriegspiel::KriegspielUmpireMessage& message, Player player) {
  if (!message.illegal) {

    after_opp_turn_observation_ = message;
    my_last_move_ = std::nullopt;
    after_my_turn_observation_ = std::nullopt;
    ply_++;
    board_.SetToPlay(color_);
    if (message.capture_type != kriegspiel::KriegspielCaptureType::kNoCapture) {
//      if (message.capture_type == kriegspiel::KriegspielCaptureType::kPawn) {
//        SPIEL_CHECK_TRUE(board_.at(message.square).type == chess::PieceType::kPawn);
//      }
//      else {
//        SPIEL_CHECK_TRUE(board_.at(message.square).type != chess::PieceType::kEmpty && board_.at(message.square).type != chess::PieceType::kPawn && board_.at(message.square).type != chess::PieceType::kKing );
//      }
      board_.set_square(message.square, chess::kEmptyPiece);
    }
  }
}

HSInfoset::HSInfoset(int seed, chess::Color color, int max_pool_size) :
    aosp_infoset_(seed, max_pool_size),
    los_infoset_(seed, color),
    max_pool_size(max_pool_size),
    color_(color) { }

chess::ChessBoard HSInfoset::SampleState() {
  return aosp_infoset_.SampleState();
}

void HSInfoset::Inform(const chess::Move& move, const kriegspiel::KriegspielUmpireMessage& message) {
  aosp_infoset_.Inform(move, message);
  los_infoset_.Inform(move, message);
}

void HSInfoset::Inform(const kriegspiel::KriegspielUmpireMessage& message, Player player) {
  aosp_infoset_.Inform(message, player);
  los_infoset_.Inform(message, player);
}

void HSInfoset::Update() {
  aosp_infoset_.Update();
  if (!los_infoset_.after_my_turn_observation_ || !los_infoset_.after_my_turn_observation_->illegal) {

    if (aosp_infoset_.Size() < max_pool_size / 2
    && already_was_higher_
    && los_infoset_.captured_pawns_ + los_infoset_.captured_pieces_ < 15) {
      for (int i = 0; i < max_pool_size - aosp_infoset_.Size(); ++i) {
        chess::ChessBoard sample = los_infoset_.SampleState();

        if (sample.board_size_ != 8) {
          int a = 5;
        }
        if (sample.find({chess::OppColor(color_), chess::PieceType::kKing}) != chess::kInvalidSquare && !IsTerminal(sample)) {
          aosp_infoset_.AddState(sample);
        }
      }
    }
  }
  if (aosp_infoset_.Size() > max_pool_size / 10 * 9) {
    already_was_higher_ = true;
  }
}

void HSInfoset::Reset() {
  aosp_infoset_.Reset();
  los_infoset_.Reset();
  already_was_higher_ = false;
}


}  // namespace algorithms
}  // namespace open_spiel
