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

MyISMCTSBot::MyISMCTSBot(int seed, int move_time,
                         StatisticalSamplingInfoset *information_set,
                         std::function<double(chess::ChessBoard, chess::Color)> evaluator,
                         double uct_c,
                         ISMCTSFinalPolicyType final_policy_type)
    : rng_(seed),
      move_time_(move_time),
      information_set_(information_set),
      evaluator_(evaluator),
      uct_c_(uct_c),
      final_policy_type_(final_policy_type) {
  my_turn_node_pool_.reserve(10000);
  opp_turn_node_pool_.reserve(10000);

}

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

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

Action MyISMCTSBot::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();

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

  // 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);
  information_set_->Inform(best_move, msg);

  return chess::MoveToAction(best_move);
}

void MyISMCTSBot::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 MyISMCTSBot::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 MyISMCTSBot::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 MyISMCTSBot::RunSearch(std::chrono::time_point<std::chrono::system_clock> step_end) {
  Reset();

  my_turn_node_pool_.push_back({});
  opp_turn_node_pool_.push_back({});
  my_msg_node_pool_.push_back({});

  int a = 0;
  do {
    chess::ChessBoard board = information_set_->SampleState();
    RunSimulation(board);
    a++;
  } while(std::chrono::system_clock::now() < step_end);

  return GetBestMove();

}

chess::Move MyISMCTSBot::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 = 0;
      for (const auto& [move, my_msg_node_idx] : root_node.child_nodes) {
        const MyMessageNode& my_msg_node = my_msg_node_pool_[my_msg_node_idx];
        if (my_msg_node.visits > max_visits) {
          max_visits = my_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, my_msg_node_idx] : root_node.child_nodes) {
        const MyMessageNode& my_msg_node = my_msg_node_pool_[my_msg_node_idx];
        double value = my_msg_node.return_sum / my_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 MyISMCTSBot::SelectMoveUCB(MyTurnNode* node, const std::vector<chess::Move>& legal_moves) {
  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_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 MyISMCTSBot::RunSimulation(chess::ChessBoard board) {

  size_t node_idx = root_node_idx;
  double value;
  while (true) {
    MyTurnNode* node = &my_turn_node_pool_[node_idx];

    if (IsTerminal(board)) {
      value = GetReturn(board, chess::PlayerToColor(*player_));
      break;
    }

    // check if we should expand
    std::vector<chess::Move> legal_moves = board.GenerateLegalMoveVector();
    std::vector<chess::Move> unexpanded_moves;
    for (const chess::Move& move : legal_moves) {
      if (!node->child_nodes.count(move)) {
        unexpanded_moves.push_back(move);

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

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

    // if we should expand
    if (!unexpanded_moves.empty()) {
      // choose move to expand
      chess::Move move_to_expand = unexpanded_moves[absl::Uniform(rng_, 0u, unexpanded_moves.size())];
      kriegspiel::KriegspielUmpireMessage msg_to_expand = kriegspiel::GetUmpireMessage(board, move_to_expand);
      board.ApplyMove(move_to_expand);

      if (IsTerminal(board)) {
        value = GetReturn(board, chess::PlayerToColor(*player_));
      }
      else {
        value = evaluator_(board, chess::PlayerToColor(*player_));
      }

      size_t my_msg_node_idx = node->child_nodes[move_to_expand];
      MyMessageNode* my_msg_node = &my_msg_node_pool_[my_msg_node_idx];
      my_msg_node->visits++;
      my_msg_node->return_sum += value;

      break;
    }

    // we select with UCB
    chess::Move selected_move = SelectMoveUCB(node, legal_moves);
    kriegspiel::KriegspielUmpireMessage selected_message = kriegspiel::GetUmpireMessage(board, selected_move);
    board.ApplyMove(selected_move);
    size_t my_msg_node_idx = node->child_nodes[selected_move];
    MyMessageNode* my_msg_node = &my_msg_node_pool_[my_msg_node_idx];

    if (IsTerminal(board)) {
      value = GetReturn(board, chess::PlayerToColor(*player_));
      my_msg_node->visits++;
      my_msg_node->return_sum += value;
      break;
    }

    if (!my_msg_node->child_nodes.count(selected_message)) {
      OppTurnNode new_opp_node = {};
      new_opp_node.parent = my_msg_node_idx;

      opp_turn_node_pool_.push_back(new_opp_node);
      my_msg_node->child_nodes[selected_message] = opp_turn_node_pool_.size() - 1;

    }

    size_t opp_turn_node_idx = my_msg_node->child_nodes[selected_message];
    OppTurnNode* opp_turn_node = &opp_turn_node_pool_[opp_turn_node_idx];


    // we select random opponent move
    std::vector<chess::Move> opp_legal_moves = board.GenerateLegalMoveVector();
    chess::Move opp_move;
    if (selected_message.capture_type == kriegspiel::KriegspielCaptureType::kNoCapture) {
      opp_move = opp_legal_moves[absl::Uniform(rng_, 0u, opp_legal_moves.size())];
    }
    else {

      std::vector<chess::Move> opp_recaptures;
      for (const chess::Move& move : opp_legal_moves) {
        if (move.to == selected_message.square) {
          opp_recaptures.push_back(move);
        }
      }
      if (opp_recaptures.empty()) {
        opp_move = opp_legal_moves[absl::Uniform(rng_, 0u, opp_legal_moves.size())];
      }
      else {
        opp_move = opp_recaptures[absl::Uniform(rng_, 0u, opp_recaptures.size())];
      }

    }

    kriegspiel::KriegspielUmpireMessage msg = kriegspiel::GetUmpireMessage(board, opp_move);

    if (opp_turn_node->child_nodes.count(msg)) {
      node_idx = opp_turn_node->child_nodes[msg];
      board.ApplyMove(opp_move);
    }
    else {
      value = evaluator_(board, chess::PlayerToColor(*player_));
      MyTurnNode my_turn_node = {};
      my_turn_node.parent = opp_turn_node_idx;

      my_turn_node_pool_.push_back(my_turn_node);
      opp_turn_node->child_nodes[msg] = my_turn_node_pool_.size() - 1;
      node_idx = my_turn_node_pool_.size() - 1;
      break;
    }
  }

  //backpropagate
  OppTurnNode* opp_turn_node;
  MyTurnNode* my_turn_node = &my_turn_node_pool_[node_idx];
  while (my_turn_node->parent) {
    size_t opp_node_idx = *my_turn_node->parent;
    opp_turn_node = &opp_turn_node_pool_[opp_node_idx];


    size_t my_msg_node_idx = *opp_turn_node->parent;
    MyMessageNode* my_msg_node = &my_msg_node_pool_[my_msg_node_idx];
    my_msg_node->return_sum += value;
    my_msg_node->visits++;


    size_t my_turn_node_idx = *my_msg_node->parent;
    my_turn_node = &my_turn_node_pool_[my_turn_node_idx];
  }

}


}  // namespace algorithms
}  // namespace open_spiel
