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

namespace open_spiel {
namespace algorithms {

MOISMCTSBot::MOISMCTSBot(int seed,
                         int move_time,
                         StatisticalSamplingInfoset *information_set,
                         std::function<double(chess::ChessBoard, chess::Color)> evaluator,
                         double uct_c,
                         ISMCTSFinalPolicyType final_policy_type) : MyISMCTSBot(
    seed,
    move_time,
    information_set,
    std::move(evaluator),
    uct_c,
    final_policy_type) { }

void MOISMCTSBot::RunSimulation(chess::ChessBoard board) {

  size_t my_tree_my_node_idx = root_node_idx;
  size_t opp_tree_opp_node_idx = opp_root_node_idx;
  size_t opp_tree_my_node_idx = std::numeric_limits<size_t>::max();
  double value;
  while (true) {
    MyTurnNode* my_tree_my_node = &my_turn_node_pool_[my_tree_my_node_idx];
    OppTurnNode* opp_tree_opp_node = &opp_turn_node_pool_[opp_tree_opp_node_idx];

    // if it's terminal, evaluate it, update opp tree and backpropagate
    if (IsTerminal(board)) {
      value = GetReturn(board, chess::PlayerToColor(*player_));
      break;
    }

    // check if we should expand in my tree
    std::vector<chess::Move> my_legal_moves = board.GenerateLegalMoveVector();
    std::vector<chess::Move> my_unexpanded_moves;
    for (const chess::Move& my_move : my_legal_moves) {
      if (!my_tree_my_node->child_nodes.count(my_move)) {
        my_unexpanded_moves.push_back(my_move);

        MyMessageNode new_my_tree_msg_node = {};
        new_my_tree_msg_node.parent = my_tree_my_node_idx;
        new_my_tree_msg_node.available_count++;
        my_msg_node_pool_.push_back(new_my_tree_msg_node);

        my_tree_my_node->child_nodes[my_move] = my_msg_node_pool_.size() - 1;
      }
      else {
        size_t my_tree_msg_node_idx = my_tree_my_node->child_nodes[my_move];
        MyMessageNode* my_tree_msg_node = &my_msg_node_pool_[my_tree_msg_node_idx];
        if (my_tree_msg_node->visits == 0) {
          my_unexpanded_moves.push_back(my_move);
        }
        my_tree_msg_node->available_count++;
      }
    }

    // if we should expand in my tree, do it and then evaluate and backpropagate
    if (!my_unexpanded_moves.empty()) {
      // choose move to expand and evaluate the result board
      chess::Move my_move = my_unexpanded_moves[absl::Uniform(rng_, 0u, my_unexpanded_moves.size())];
      kriegspiel::KriegspielUmpireMessage my_msg = kriegspiel::GetUmpireMessage(board, my_move);
      board.ApplyMove(my_move);
      if (IsTerminal(board)) {
        value = GetReturn(board, chess::PlayerToColor(*player_));
      }
      else {
        value = evaluator_(board, chess::PlayerToColor(*player_));
      }

      // create new node in my tree and update value
      size_t new_my_tree_msg_node_idx = my_tree_my_node->child_nodes[my_move];
      MyMessageNode* new_my_tree_msg_node = &my_msg_node_pool_[new_my_tree_msg_node_idx];
      new_my_tree_msg_node->visits++;
      new_my_tree_msg_node->return_sum += value;

      // check the opp tree and add a node if necessary
      if (!opp_tree_opp_node->child_nodes.count(my_msg)) {
        MyTurnNode new_opp_tree_my_node = {};
        new_opp_tree_my_node.parent = opp_tree_opp_node_idx;
        my_turn_node_pool_.push_back(new_opp_tree_my_node);
        opp_tree_opp_node->child_nodes[my_msg] = my_turn_node_pool_.size() - 1;
      }
      opp_tree_my_node_idx = opp_tree_opp_node->child_nodes[my_msg];

      break;
    }

    // we select in my tree with UCB
    chess::Move my_selected_move = SelectMoveUCB(my_tree_my_node, my_legal_moves);
    kriegspiel::KriegspielUmpireMessage my_msg = kriegspiel::GetUmpireMessage(board, my_selected_move);
    board.ApplyMove(my_selected_move);
    size_t my_tree_msg_node_idx = my_tree_my_node->child_nodes[my_selected_move];
    MyMessageNode* my_tree_msg_node = &my_msg_node_pool_[my_tree_msg_node_idx];

    if (!my_tree_msg_node->child_nodes.count(my_msg)) {
      OppTurnNode new_my_tree_opp_node = {};
      new_my_tree_opp_node.parent = my_tree_msg_node_idx;

      opp_turn_node_pool_.push_back(new_my_tree_opp_node);
      my_tree_msg_node->child_nodes[my_msg] = opp_turn_node_pool_.size() - 1;

    }

    size_t my_tree_opp_node_idx = my_tree_msg_node->child_nodes[my_msg];
    OppTurnNode* my_tree_opp_node = &opp_turn_node_pool_[my_tree_opp_node_idx];

    // check the opp tree and add a node if necessary
    if (!opp_tree_opp_node->child_nodes.count(my_msg)) {
      MyTurnNode new_opp_tree_my_node = {};
      new_opp_tree_my_node.parent = opp_tree_opp_node_idx;
      my_turn_node_pool_.push_back(new_opp_tree_my_node);
      opp_tree_opp_node->child_nodes[my_msg] = my_turn_node_pool_.size() - 1;
    }
    // descend the opp tree
    opp_tree_my_node_idx = opp_tree_opp_node->child_nodes[my_msg];
    MyTurnNode* opp_tree_my_node = &my_turn_node_pool_[opp_tree_my_node_idx];

    // if terminal, just evaluate, update my tree node and backpropagate
    if (IsTerminal(board)) {
      value = GetReturn(board, chess::PlayerToColor(*player_));
      my_tree_msg_node->visits++;
      my_tree_msg_node->return_sum += value;
      break;
    }

    // check if we should expand in opp tree
    std::vector<chess::Move> opp_legal_moves = board.GenerateLegalMoveVector();
    std::vector<chess::Move> opp_unexpanded_moves;
    for (const chess::Move& opp_move : opp_legal_moves) {
      if (!opp_tree_my_node->child_nodes.count(opp_move)) {
        opp_unexpanded_moves.push_back(opp_move);

        MyMessageNode new_opp_tree_msg_node = {};
        new_opp_tree_msg_node.parent = opp_tree_my_node_idx;
        new_opp_tree_msg_node.available_count++;
        my_msg_node_pool_.push_back(new_opp_tree_msg_node);

        opp_tree_my_node->child_nodes[opp_move] = my_msg_node_pool_.size() - 1;
      }
      else {
        size_t new_opp_tree_msg_node_idx = opp_tree_my_node->child_nodes[opp_move];
        MyMessageNode* new_opp_tree_msg_node = &my_msg_node_pool_[new_opp_tree_msg_node_idx];
        if (new_opp_tree_msg_node->visits == 0) {
          opp_unexpanded_moves.push_back(opp_move);
        }
        new_opp_tree_msg_node->available_count++;
      }
    }

    // if we should expand in opp tree, do it and then evaluate and backpropagate
    if (!opp_unexpanded_moves.empty()) {
      // choose move to expand and evaluate the result board
      chess::Move opp_move = opp_unexpanded_moves[absl::Uniform(rng_, 0u, opp_unexpanded_moves.size())];
      kriegspiel::KriegspielUmpireMessage opp_msg = kriegspiel::GetUmpireMessage(board, opp_move);
      board.ApplyMove(opp_move);
      if (IsTerminal(board)) {
        value = GetReturn(board, chess::PlayerToColor(*player_));
      }
      else {
        value = evaluator_(board, chess::PlayerToColor(*player_));
      }

      // create new node in opp tree and update value
      size_t new_opp_tree_msg_node_idx = opp_tree_my_node->child_nodes[opp_move];
      MyMessageNode* new_opp_tree_msg_node = &my_msg_node_pool_[new_opp_tree_msg_node_idx];
      new_opp_tree_msg_node->visits++;
      new_opp_tree_msg_node->return_sum += 1.0 - value;

      // check my tree and add a node if necessary
      if (!my_tree_opp_node->child_nodes.count(opp_msg)) {
        MyTurnNode new_my_tree_my_node = {};
        new_my_tree_my_node.parent = my_tree_opp_node_idx;
        my_turn_node_pool_.push_back(new_my_tree_my_node);
        my_tree_opp_node->child_nodes[opp_msg] = my_turn_node_pool_.size() - 1;
      }
      my_tree_my_node_idx = my_tree_opp_node->child_nodes[opp_msg];

      break;
    }

    // we select opp move with UCB
    chess::Move opp_selected_move = SelectMoveUCB(opp_tree_my_node, opp_legal_moves);
    kriegspiel::KriegspielUmpireMessage opp_msg = kriegspiel::GetUmpireMessage(board, opp_selected_move);
    board.ApplyMove(opp_selected_move);
    size_t opp_tree_msg_node_idx = opp_tree_my_node->child_nodes[opp_selected_move];
    MyMessageNode* opp_tree_msg_node = &my_msg_node_pool_[opp_tree_msg_node_idx];
    opp_tree_opp_node_idx = opp_tree_msg_node->child_nodes[opp_msg];


    if (!my_tree_opp_node->child_nodes.count(opp_msg)) {
      MyTurnNode new_my_tree_my_node = {};
      new_my_tree_my_node.parent = my_tree_opp_node_idx;
      my_turn_node_pool_.push_back(new_my_tree_my_node);
      my_tree_opp_node->child_nodes[opp_msg] = my_turn_node_pool_.size() - 1;
    }
    //descend my tree
    my_tree_my_node_idx = my_tree_opp_node->child_nodes[opp_msg];

    // continue the next iteration
  }

  // backpropagate in my tree
  MyTurnNode* my_turn_node = &my_turn_node_pool_[my_tree_my_node_idx];
  while (my_turn_node->parent) {
    size_t opp_turn_node_idx = *my_turn_node->parent;
    OppTurnNode* opp_turn_node = &opp_turn_node_pool_[opp_turn_node_idx];
    size_t msg_node_idx = *opp_turn_node->parent;
    MyMessageNode* msg_node = &my_msg_node_pool_[msg_node_idx];
    msg_node->return_sum += value;
    msg_node->visits++;

    // should never be empty
    my_tree_my_node_idx = *msg_node->parent;
    my_turn_node = &my_turn_node_pool_[my_tree_my_node_idx];
  }
  SPIEL_CHECK_EQ(my_tree_my_node_idx, 0);


  // backpropagate in opp tree
  if (opp_tree_my_node_idx != std::numeric_limits<size_t>::max()) {
    MyTurnNode* opp_tree_my_node = &my_turn_node_pool_[opp_tree_my_node_idx];
    while (opp_tree_my_node->parent) {
      size_t opp_turn_node_idx = *opp_tree_my_node->parent;
      OppTurnNode* opp_turn_node = &opp_turn_node_pool_[opp_turn_node_idx];
      if (!opp_turn_node->parent) {
        break;
      }
      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 += 1.0 - value;
      my_msg_node->visits++;

      if (my_msg_node->parent) {
        opp_tree_my_node_idx = *my_msg_node->parent;
        opp_tree_my_node = &my_turn_node_pool_[opp_tree_my_node_idx];
      }
      else {
        break;
      }
    }
    SPIEL_CHECK_NE(opp_tree_my_node_idx, 0);
  }

}

}  // namespace algorithms
}  // namespace open_spiel
