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

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

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

#include "open_spiel/spiel.h"
#include "open_spiel/spiel_utils.h"
#include "open_spiel/games/chess.h"
#include "open_spiel/games/kriegspiel.h"
#include "open_spiel/abseil-cpp/absl/random/distributions.h"

namespace open_spiel {
namespace algorithms {

PIMCBot::PIMCBot(int move_time, int min_state_eval_time,
                 StatisticalSamplingInfoset *information_set,
                 std::function<std::unordered_map<chess::Move, double, ChessMoveHash>(const chess::ChessBoard &, int)> action_evaluator,
                 ActionStatsWeights action_stats_weights)
    : move_time_(move_time),
      min_state_eval_time_(min_state_eval_time),
      information_set_(information_set),
      move_evaluator_(std::move(action_evaluator)),
      action_stats_weights_(action_stats_weights) {}

void PIMCBot::Restart() {
  information_set_->Reset();
}

Action PIMCBot::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_);

  SetCurrentPlayer(state);
  information_set_->Update();

  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;
//  std::cout << "Information set size: " << information_set_->Size() << std::endl;

  if (information_set_->Empty()) {
    std::mt19937 rng(time(nullptr));
    return RandomStep(state, rng, *player);
  }

  // find best move
  std::unordered_map<chess::Move, MoveStats, ChessMoveHash> move_stats_map = ComputeMoveStatsMap(step_end);
  chess::Move best_move = GetBestMove(move_stats_map);

  //  update observations for updating infoset next turn
  auto kriegspielState = down_cast<const kriegspiel::KriegspielState &>(state);
  kriegspiel::KriegspielUmpireMessage msg = kriegspiel::GetUmpireMessage(kriegspielState.Board(), best_move);
  information_set_->Inform(best_move, msg);

  return chess::MoveToAction(best_move);
}

void PIMCBot::ForceAction(const State &state,
                          Action action) {
  SetCurrentPlayer(state);
  information_set_->Update();

  //  update observations for updating infoset next turn
  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);
  information_set_->Inform(move, msg);
}

void PIMCBot::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);
  information_set_->Inform(msg, player_id);
}

void PIMCBot::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");
}

std::unordered_map<chess::Move, MoveStats, ChessMoveHash> PIMCBot::ComputeMoveStatsMap(std::chrono::time_point<std::chrono::system_clock> step_end) {
  std::unordered_map<chess::Move, MoveStats, ChessMoveHash> move_stats_map;

  uint move_time = std::chrono::duration_cast<std::chrono::milliseconds>(step_end - std::chrono::system_clock::now()).count();

  int infoset_size = (int) information_set_->Size();
//  int infoset_size = 1;
  uint board_eval_time = move_time / infoset_size;
  int i = 0;

  if (board_eval_time > min_state_eval_time_ && step_end > std::chrono::system_clock::now()) {
    for (const chess::ChessBoard &board : information_set_->Vector()) {
//    chess::ChessBoard board = information_set_->SampleState();
      i++;
      std::unordered_map<chess::Move, double, ChessMoveHash> curr_move_stats = move_evaluator_(board, board_eval_time);
      chess::Move best_move;
      double best_value = -100000000;
      for (const auto &[move, value] : curr_move_stats) {
        if (value > best_value) {
          best_move = move;
          best_value = value;
        }
        if (move_stats_map.count(move)) {
          MoveStats orig_action_stats = move_stats_map[move];
          orig_action_stats.count++;
          orig_action_stats.sum += value;
          if (value > orig_action_stats.max) orig_action_stats.max = value;
          if (value < orig_action_stats.min) orig_action_stats.min = value;
          move_stats_map[move] = orig_action_stats;
        } else {
          move_stats_map[move] = {value, value, value, 1, 0};
        }
      }
      MoveStats orig = move_stats_map[best_move];
      orig.wins++;
      move_stats_map[best_move] = orig;
    }
  }
  else {
    do {
      i++;
      chess::ChessBoard board = information_set_->SampleState();
      std::unordered_map<chess::Move, double, ChessMoveHash> curr_move_stats = move_evaluator_(board, min_state_eval_time_);
      chess::Move best_move;
      double best_value = -100000000;
      for (const auto &[move, value] : curr_move_stats) {
        if (value > best_value) {
          best_move = move;
          best_value = value;
        }
        if (move_stats_map.count(move)) {
          MoveStats orig_action_stats = move_stats_map[move];
          orig_action_stats.count++;
          orig_action_stats.sum += value;
          if (value > orig_action_stats.max) orig_action_stats.max = value;
          if (value < orig_action_stats.min) orig_action_stats.min = value;
          move_stats_map[move] = orig_action_stats;
        } else {
          move_stats_map[move] = {value, value, value, 1, 0};
        }
      }
      MoveStats orig = move_stats_map[best_move];
      orig.wins++;
      move_stats_map[best_move] = orig;
    } while(std::chrono::system_clock::now() < step_end);
  }
//  std::cout << i << std::endl;
  return move_stats_map;
}

chess::Move PIMCBot::GetBestMove(const std::unordered_map<chess::Move,
                                           MoveStats, ChessMoveHash> &action_stats_map) {

  chess::Move best_move;
  double best_value = -10000000000000; //todo set minimum

  for (const auto &[move, stats] : action_stats_map) {
    if (stats.count == 0) {
      continue;
    }
    double curr_value = stats.min * action_stats_weights_.min +
                        stats.max * action_stats_weights_.max +
                        stats.wins * action_stats_weights_.wins +
                        (stats.sum / stats.count)
                        * action_stats_weights_.average;
    if (curr_value > best_value) {
      best_value = curr_value;
      best_move = move;
    }
  }
  return best_move;
}

}  // namespace algorithms
}  // namespace open_spiel
