// 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/evaluation.h"

#include "open_spiel/spiel_bots.h"
#include "open_spiel/algorithms/evaluate_bots.h"
//#include "open_spiel/algorithms/kriegspiel/darkboard.h"
#include "open_spiel/algorithms/kriegspiel/pimc.h"
#include "open_spiel/algorithms/kriegspiel/mo_ismcts.h"
#include "open_spiel/algorithms/kriegspiel/random_bot.h"
#include "open_spiel/abseil-cpp/absl/random/distributions.h"
#include "open_spiel/bots/stockfish/stockfish_bot.h"
#include "open_spiel/games/chess/pgn.h"

namespace open_spiel {
namespace algorithms {

void EvaluateBots(Bot* bot1, Bot* bot2, int seed, int count) {

  int wins1 = 0;
  int draws1 = 0;
  int losses1 = 0;

  auto game = open_spiel::LoadGame("kriegspiel");
  std::vector<Bot *> bot_ptrs;

  bot_ptrs.push_back(bot1);
  bot_ptrs.push_back(bot2);

  std::cout << "Playing " << count << " games" << std::endl;
  std::cout << "Black: " << "bot1" << std::endl;
  std::cout << "White: " << "bot2" << std::endl << std::endl;

  for (int i = 0; i < count; ++i) {

    std::unique_ptr<open_spiel::State> state = game->NewInitialState();
    auto returns = EvaluateBots(state.get(), bot_ptrs, seed);

    int ret = returns[0];
    if (ret == kriegspiel::WinUtility()) wins1++;
    if (ret == kriegspiel::LossUtility()) losses1++;
    if (ret == kriegspiel::DrawUtility()) draws1++;
  }

  std::cout << "Black wins: " << wins1 << std::endl;
  std::cout << "Draw: " << draws1 << std::endl;
  std::cout << "White wins: " << losses1 << std::endl;
  std::cout << std::endl;

  int wins2 = 0;
  int draws2 = 0;
  int losses2 = 0;

  bot_ptrs.clear();

  bot_ptrs.push_back(bot2);
  bot_ptrs.push_back(bot1);

  std::cout << "Playing " << count << " games" << std::endl;
  std::cout << "Black: " << "bot2" << std::endl;
  std::cout << "White: " << "bot1" << std::endl << std::endl;

  for (int i = 0; i < count; ++i) {

    std::unique_ptr<open_spiel::State> state = game->NewInitialState();
    auto returns = EvaluateBots(state.get(), bot_ptrs, seed);

    int ret = returns[0];
    if (ret == kriegspiel::WinUtility()) wins2++;
    if (ret == kriegspiel::LossUtility()) losses2++;
    if (ret == kriegspiel::DrawUtility()) draws2++;
  }


  std::cout << "Black wins: " << wins2 << std::endl;
  std::cout << "Draw: " << draws2 << std::endl;
  std::cout << "White wins: " << losses2 << std::endl;

  std::cout << std::endl;

  std::cout << "Overall:" << std::endl;
  std::cout << "bot1" << std::endl;
  std::cout << "\t\t" << wins1 + losses2 << std::endl;
  std::cout << "bot2" << std::endl;
  std::cout << "\t\t" << losses1 + wins2 << std::endl;
  std::cout << "Draws" << std::endl;
  std::cout << "\t\t" << draws1 + draws2 << std::endl;

}

void Evaluate(std::array<std::string, 2> players) {

  std::vector<Bot*> bots;
  int seed = time(NULL);
  std::mt19937 rng(seed);
  std::vector<HSInfoset> iss;
  std::vector<std::unique_ptr<stockfish::StockfishBot>> sfs;

  for (int i = 0; i < 2; ++i) {
    std::string player_str = players[i];

    if (player_str == "random") {
      bots.push_back(new RandomBot(absl::Uniform(rng, 0, std::numeric_limits<int>::max())));
    }
    else if (player_str == "random_h") {
      bots.push_back(new RandomBotHeuristic(absl::Uniform(rng, 0, std::numeric_limits<int>::max())));
    }
    else if (player_str == "pimc") {

      sfs.push_back(stockfish::MakeStockfishBot(10));
      stockfish::StockfishBot* sf_bot = sfs[sfs.size() - 1].get();
      bots.push_back(new PIMCBot(
          2000,
          20,
          new HSInfoset(absl::Uniform(rng, 0, std::numeric_limits<int>::max()), chess::PlayerToColor(i), 1000),
          [&sf_bot](const chess::ChessBoard &board, int eval_time) {

            std::unordered_map<chess::Move, double, algorithms::ChessMoveHash> action_value_map;
            uci::ResultMove result_move = sf_bot->Step(board, eval_time);
            absl::optional<chess::Move> move = board.ParseLANMove(result_move.move);
            if (!move || !board.IsMoveLegal(*move)) {
              uci::ResultMove r = sf_bot->ReadBestMove();
              result_move = sf_bot->Step(board, eval_time);
              move = board.ParseLANMove(result_move.move);
            }
            action_value_map[*move] = result_move.score;

            return action_value_map;
          },
          {0.0, 0.0, 0.0, 1.0}));
    }
    else if (player_str == "ismcts") {
      auto random_rollout_evaluator = [&rng] (const chess::ChessBoard& board, chess::Color color) {
        chess::ChessBoard working_board = board;

        chess::Square capture_sq = chess::kInvalidSquare;
        while (!algorithms::IsTerminal(working_board)) {
          std::vector<chess::Move> legal_moves = working_board.GenerateLegalMoveVector();
          chess::Move move;
          if (capture_sq == chess::kInvalidSquare) {
            move = legal_moves[absl::Uniform(rng, 0u, legal_moves.size())];
          }
          else {
            std::vector<chess::Move> recapture_moves;
            for (const chess::Move& loc_move : legal_moves) {
              if (loc_move.to == capture_sq) {
                recapture_moves.push_back(loc_move);
              }
            }
            if (recapture_moves.empty()) {
              move = legal_moves[absl::Uniform(rng, 0u, legal_moves.size())];
            }
            else {
              move = recapture_moves[absl::Uniform(rng, 0u, recapture_moves.size())];
            }

          }
          working_board.ApplyMove(move);
        }

        double reward = algorithms::GetReturn(working_board, color);
        return reward;
      };

      bots.push_back(new MyISMCTSBot(
          absl::Uniform(rng, 0, std::numeric_limits<int>::max()),
          2000,
          new HSInfoset(absl::Uniform(rng, 0, std::numeric_limits<int>::max()), chess::PlayerToColor(i), 1000),
          random_rollout_evaluator,
          0.5,
          algorithms::ISMCTSFinalPolicyType::kMaxValue));
    }
    else if (player_str == "ismcts_sf") {

      int board_eval_time = 20;
      sfs.push_back(stockfish::MakeStockfishBot(10));
      stockfish::StockfishBot* sf_bot = sfs[sfs.size() - 1].get();
      auto sf_evaluator = [&sf_bot, &board_eval_time] (const chess::ChessBoard& board, chess::Color color) {
        uci::ResultMove result_move = sf_bot->Step(board, board_eval_time);
        absl::optional<chess::Move> move = board.ParseLANMove(result_move.move);
        if (!move || !board.IsMoveLegal(*move)) {
          uci::ResultMove r = sf_bot->ReadBestMove();
          result_move = sf_bot->Step(board, board_eval_time);
          move = board.ParseLANMove(result_move.move);
        }

        int multiplier = board.ToPlay() == color ? -1 : 1;
        double win_prob = 1.0 / (1.0 + std::pow(10, ((double) multiplier* result_move.score) / 400));

        return win_prob;
      };

      bots.push_back(new MyISMCTSBot(
          absl::Uniform(rng, 0, std::numeric_limits<int>::max()),
          2000,
          new HSInfoset(absl::Uniform(rng, 0, std::numeric_limits<int>::max()), chess::PlayerToColor(i), 1000),
          sf_evaluator,
          0.25,
          algorithms::ISMCTSFinalPolicyType::kMaxValue));
    }
    else if (player_str == "mo_ismcts") {
      auto random_rollout_evaluator = [&rng] (const chess::ChessBoard& board, chess::Color color) {
        chess::ChessBoard working_board = board;

        chess::Square capture_sq = chess::kInvalidSquare;
        while (!algorithms::IsTerminal(working_board)) {
          std::vector<chess::Move> legal_moves = working_board.GenerateLegalMoveVector();
          chess::Move move;
          if (capture_sq == chess::kInvalidSquare) {
            move = legal_moves[absl::Uniform(rng, 0u, legal_moves.size())];
          }
          else {
            std::vector<chess::Move> recapture_moves;
            for (const chess::Move& loc_move : legal_moves) {
              if (loc_move.to == capture_sq) {
                recapture_moves.push_back(loc_move);
              }
            }
            if (recapture_moves.empty()) {
              move = legal_moves[absl::Uniform(rng, 0u, legal_moves.size())];
            }
            else {
              move = recapture_moves[absl::Uniform(rng, 0u, recapture_moves.size())];
            }

          }
          working_board.ApplyMove(move);
        }

        double reward = algorithms::GetReturn(working_board, color);
        return reward;
      };

      bots.push_back(new MOISMCTSBot(
          absl::Uniform(rng, 0, std::numeric_limits<int>::max()),
          2000,
          new HSInfoset(absl::Uniform(rng, 0, std::numeric_limits<int>::max()), chess::PlayerToColor(i), 1000),
          random_rollout_evaluator,
          0.25,
          algorithms::ISMCTSFinalPolicyType::kMaxValue));
    }
    else if (player_str == "mo_ismcts_sf") {

      int board_eval_time = 5;
      sfs.push_back(stockfish::MakeStockfishBot(10));
      stockfish::StockfishBot* sf_bot = sfs[sfs.size() - 1].get();
      auto sf_evaluator = [&sf_bot, &board_eval_time] (const chess::ChessBoard& board, chess::Color color) {
        uci::ResultMove result_move = sf_bot->Step(board, board_eval_time);
        absl::optional<chess::Move> move = board.ParseLANMove(result_move.move);
        if (!move || !board.IsMoveLegal(*move)) {
          uci::ResultMove r = sf_bot->ReadBestMove();
          result_move = sf_bot->Step(board, board_eval_time);
          move = board.ParseLANMove(result_move.move);
        }

        int multiplier = board.ToPlay() == color ? -1 : 1;
        double win_prob = 1.0 / (1.0 + std::pow(10, ((double) multiplier* result_move.score) / 400));

        return win_prob;
      };

      bots.push_back(new MOISMCTSBot(
          absl::Uniform(rng, 0, std::numeric_limits<int>::max()),
          2000,
          new HSInfoset(absl::Uniform(rng, 0, std::numeric_limits<int>::max()), chess::PlayerToColor(i), 1000),
          sf_evaluator,
          0.5,
          algorithms::ISMCTSFinalPolicyType::kMaxValue));

    }
    else if (player_str == "darkboard") {
      SpielFatalError("Darkboard not implemented yet");

    }


  }

  for (int i = 0; i < 1; ++i) {

    auto game = open_spiel::LoadGame("kriegspiel");
    std::unique_ptr<open_spiel::State> state = game->NewInitialState();
    auto returns = EvaluateBots(state.get(), bots, absl::Uniform<int>(rng, 0, std::numeric_limits<int>::max()));


    pgn::KriegspielPGN pgn = pgn::KriegspielPGN(
        down_cast<kriegspiel::KriegspielState>(*state),
        "Vojtech Foret Master Thesis",
        i,
        players[1],
        players[0],
        chess::Color::kEmpty);

    std::cout << pgn << std::endl;
  }


}


}  // namespace algorithms
}  // namespace open_spiel
