// 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.

#ifndef OPEN_SPIEL_GAMES_KRIEGSPIEL_H_
#define OPEN_SPIEL_GAMES_KRIEGSPIEL_H_

#include <array>
#include <cstring>
#include <map>
#include <memory>
#include <string>
#include <vector>
#include <open_spiel/fog/observation_history.h>

#include "open_spiel/abseil-cpp/absl/container/flat_hash_map.h"
#include "open_spiel/abseil-cpp/absl/algorithm/container.h"
#include "open_spiel/games/chess/chess_board.h"
#include "open_spiel/games/chess.h"
#include "open_spiel/spiel.h"
#include "open_spiel/spiel_utils.h"

// Kriegspiel - imperfect information version of chess:
// https://en.wikipedia.org/wiki/Kriegspiel
// The implementation follows ICC rules (with a few exceptions):
// https://www.chessclub.com/help/Kriegspiel
// The exceptions are 50-move rule and threefold repetition, which under ICC
// rules are not automatically enforced, but can be claimed by the player. This
// implementation does not support claiming or offering draws so these rules'
// automatic enforcement can be turned on and off
//
// Parameters:
//   "board_size"           int     Number of squares in each row and column
//                                  (default: 8)
//   "fen"                  string  String describing the chess board position
//                                  in Forsyth-Edwards Notation. The FEN has to
//                                  match the board size. Default values are
//                                  available for board sizes 4 and 8.
//   "threefold_repetition" bool    Whether threefold repetition rule should be
//                                  automatically enforced (default: false)
//   "50_move_rule"         bool    Whether 50 move rule should be automatically
//                                  enforced (default: false)

namespace open_spiel {
namespace kriegspiel {

// Constants.
inline constexpr int NumPlayers() { return 2; }
inline constexpr double LossUtility() { return -1; }
inline constexpr double DrawUtility() { return 0; }
inline constexpr double WinUtility() { return 1; }

// See action encoding below.
inline constexpr int NumDistinctActions() { return 4672; }

// https://math.stackexchange.com/questions/194008/how-many-turns-can-a-chess-game-take-at-maximum
inline constexpr int MaxGameLength() { return 17695; }

enum KriegspielCaptureType : int8_t {
  kNoCapture = 0,
  kPawn = 1,
  kPiece = 2
};

std::string CaptureTypeToString(KriegspielCaptureType capture_type);

enum KriegspielCheckType : int8_t {
  kNoCheck = 0,
  kFile = 1,
  kRank = 2,
  kLongDiagonal = 3,
  kShortDiagonal = 4,
  kKnight = 5
};

std::string CheckTypeToString(KriegspielCheckType check_type);

std::pair<KriegspielCheckType, KriegspielCheckType> GetCheckType(const chess::ChessBoard &board);

struct KriegspielUmpireMessage {

  bool illegal = false;
  KriegspielCaptureType capture_type = KriegspielCaptureType::kNoCapture;
  chess::Square square = chess::kInvalidSquare;
  // there can be max two checks at a time so a pair is enough
  std::pair<KriegspielCheckType, KriegspielCheckType> check_types =
      {KriegspielCheckType::kNoCheck, KriegspielCheckType::kNoCheck};
  chess::Color to_move = chess::Color::kEmpty;
  int pawn_tries = 0;

  std::string ToString() const;

  bool operator==(const KriegspielUmpireMessage &other) const {
    return illegal == other.illegal &&
           capture_type == other.capture_type &&
           square == other.square &&
           check_types == other.check_types &&
           to_move == other.to_move &&
           pawn_tries == other.pawn_tries;
  }
};

KriegspielUmpireMessage GetUmpireMessage(const chess::ChessBoard &chess_board,
                                         const chess::Move &move);

bool GeneratesUmpireMessage(const chess::ChessBoard &chess_board,
                            const chess::Move &move,
                            const KriegspielUmpireMessage &orig_msg);

class KriegspielGame;
class KriegspielObserver;

// State of an in-play game.
class KriegspielState : public State {
 public:

  // Constructs a chess state at the given position in Forsyth-Edwards Notation.
  // https://en.wikipedia.org/wiki/Forsyth%E2%80%93Edwards_Notation
  KriegspielState(std::shared_ptr<const Game> game,
                 int board_size, const std::string& fen, bool threefold_repetition,
                 bool rule_50_move, bool till_end);
  KriegspielState(const KriegspielState&) = default;

  KriegspielState& operator=(const KriegspielState&) = default;

  Player CurrentPlayer() const override {
    return IsTerminal() ? kTerminalPlayerId : ColorToPlayer(Board().ToPlay());
  }
  std::vector<Action> LegalActions() const override;
  std::string ActionToString(Player player, Action action) const override;
  std::string ToString() const override;

  bool IsTerminal() const override {
    return static_cast<bool>(MaybeFinalReturns());
  }

  std::vector<double> Returns() const override;
  std::string InformationStateString(Player player) const override;
  std::string ObservationString(Player player) const override;
  void ObservationTensor(Player player,
                         absl::Span<float> values) const override;
  std::unique_ptr<State> Clone() const override;
  void UndoAction(Player player, Action action) override;

  // Current board.
  chess::ChessBoard& Board() { return current_board_; }
  const chess::ChessBoard& Board() const { return current_board_; }
  int BoardSize() const { return current_board_.BoardSize(); }

  // Starting board.
  chess::ChessBoard& StartBoard() { return start_board_; }
  const chess::ChessBoard& StartBoard() const { return start_board_; }

  std::vector<std::pair<chess::Move, KriegspielUmpireMessage>>& MoveMsgHistory() {
    return move_msg_history_;
  }
  const std::vector<std::pair<chess::Move, KriegspielUmpireMessage>>& MoveMsgHistory() const {
    return move_msg_history_;
  }

 protected:
  void DoApplyAction(Action action) override;

 private:

  friend class KriegspielObserver;

  // Draw can be claimed under the FIDE threefold repetition rule (the current
  // board position has already appeared twice in the history).
  bool IsThreefoldRepetitionDraw() const;

  // Calculates legal actions and caches them. This is separate from
  // LegalActions() as there are a number of other methods that need the value
  // of LegalActions. This is a separate method as it's called from
  // IsTerminal(), which is also called by LegalActions().
  void MaybeGenerateLegalActions() const;

  absl::optional<std::vector<double>> MaybeFinalReturns() const;

  // We have to store every move made to check for repetitions and to implement
  // undo. We store the current board position as an optimization.
  std::vector<std::pair<chess::Move, KriegspielUmpireMessage>> move_msg_history_;
  // We store this info as an optimisation so that we don't have to compute it
  // from move_msg_history for observations
  std::optional<KriegspielUmpireMessage> last_public_msg_{};
  std::optional<KriegspielUmpireMessage> before_last_public_msg_{};
  // Moves that the player tried and were illegal. We don't let player try them
  // again on the same board because they are clearly still illegal;
  std::vector<chess::Move> illegal_tried_moves_;
  // We store the start board for history to support games not starting
  // from the start position.
  chess::ChessBoard start_board_;
  // We store the current board position as an optimization.
  chess::ChessBoard current_board_;

  // cached ActionObservationHistory for each player
  std::vector<open_spiel::ActionObservationHistory> aohs_;

  bool threefold_repetition_;
  bool rule_50_move_;
  bool till_end_;

  // RepetitionTable records how many times the given hash exists in the history
  // stack (including the current board).
  // We are already indexing by board hash, so there is no need to hash that
  // hash again, so we use a custom passthrough hasher.
  class PassthroughHash {
   public:
    std::size_t operator()(uint64_t x) const {
      return static_cast<std::size_t>(x);
    }
  };
  using RepetitionTable = absl::flat_hash_map<uint64_t, int, PassthroughHash>;
  RepetitionTable repetitions_;
  mutable absl::optional<std::vector<Action>> cached_legal_actions_;
};

// Game object.
class KriegspielGame : public Game {
 public:
  explicit KriegspielGame(const GameParameters& params);
  int NumDistinctActions() const override {
    return kriegspiel::NumDistinctActions();
  }
  std::unique_ptr<State> NewInitialState() const override {
    return absl::make_unique<KriegspielState>(shared_from_this(),
                                             board_size_, fen_,
                                             threefold_repetition_, rule_50_move_, till_end_);
  }
  int NumPlayers() const override { return kriegspiel::NumPlayers(); }
  double MinUtility() const override { return LossUtility(); }
  double UtilitySum() const override { return DrawUtility(); }
  double MaxUtility() const override { return WinUtility(); }
  std::vector<int> ObservationTensorShape() const override {
    std::vector<int> shape{
        14 // private boards: piece types * colours + empty + unknown
         * board_size_ * board_size_ +
        3 + // public: repetitions count, one-hot encoding
        2 + // public: side to play
        1 + // public: irreversible move counter -- a fraction of $n over 100
        2 * ( // public: last two umpire messages
            3 + // capture type
            6 + // check type one
            6 + // check type two
            3 + // player to move
            16 + // pawn tries
            board_size_ * board_size_) + // capture square
        2*2 + // private: left/right castling rights, one-hot encoded.
        2 + // private: whether last move was illegal
        2 * (board_size_ * board_size_)// private: last move (from, two)
          + 6 // private: last move (promotion type)
    };
    return shape;
  }
  int MaxGameLength() const override { return kriegspiel::MaxGameLength(); }
  std::shared_ptr<Observer> MakeObserver(
      absl::optional<IIGObservationType> iig_obs_type,
      const GameParameters& params) const;


  std::shared_ptr<KriegspielObserver> default_observer_;
  std::shared_ptr<KriegspielObserver> info_state_observer_;

 private:
  const int board_size_;
  const std::string fen_;
  const bool threefold_repetition_;
  const bool rule_50_move_;
  const bool till_end_;
};

}  // namespace kriegspiel
}  // namespace open_spiel

#endif  // OPEN_SPIEL_GAMES_KRIEGSPIEL_H_
