text: ''' MCTSClass.py: Class for a basic Monte Carlo Tree Search Planner. '''
  
  # Python imports.
  import math as m
  import random
  from collections import defaultdict
  
  # Other imports.
  from simple_rl.planning.PlannerClass import Planner
  
  class MCTS(Planner):
  
      def __init__(self, mdp, name="mcts", explore_param=m.sqrt(2), rollout_depth=20, num_rollouts_per_step=10):
          Planner.__init__(self, mdp, name=name)
  
          self.rollout_depth = rollout_depth
          self.num_rollouts_per_step = num_rollouts_per_step
          self.value_total = defaultdict(lambda : defaultdict(float))
          self.explore_param = explore_param
          self.visitation_counts = defaultdict(lambda : defaultdict(lambda : 0))
  
      def plan(self, cur_state, horizon=20):
          '''
          Args:
              cur_state (State)
              horizon (int)
  
          Returns:
              (list): List of actions
          '''
          action_seq = []
          state_seq = [cur_state]
          steps = 0
          while not cur_state.is_terminal() and steps < horizon:
              action = self._next_action(cur_state)
              # Do the rollouts...
              cur_state = self.transition_func(cur_state, action)
              action_seq.append(action)
              state_seq.append(cur_state)
              steps += 1
  
          self.has_planned = True
  
          return action_seq, state_seq
  
      def policy(self, state):
          '''
          Args:
              state (State)
  
          Returns:
              (str)
          '''
          if not self.has_planned:
              self.plan(state)
  
          return self._next_action(state)
  
      def _next_action(self, state):
          '''
          Args;
              state (State)
  
          Returns:
              (str)
  
          Summary:
              Performs a single step of the MCTS rollout.
          '''
          best_action = self.actions[0]
          best_score = 0
          total_visits = [self.visitation_counts[state][a] for a in self.actions]
  
          print(total_visits)
  
          if 0 in total_visits:
              # Insufficient stats, return random.
                  # Should choose randomly AMONG UNSAMPLED.
              unsampled_actions = [self.actions[i] for i, x in enumerate(total_visits) if x == 0]
              next_action = random.choice(unsampled_actions)
              self.visitation_counts[state][next_action] += 1
              return next_action
  
          total = sum(total_visits)
  
          # Else choose according to the UCT explore method.
          for cur_action in self.actions:
              s_a_value_tot = self.value_total[state][cur_action]
              s_a_visit = self.visitation_counts[state][cur_action]
              score = s_a_value_tot / s_a_visit + self.explore_param * m.sqrt(m.log(total) / s_a_visit)
  
              if score > best_score:
                  best_action = cur_action
                  best_score = score
  
          return best_action
  
      def _rollout(self, cur_state, action):
          '''
          Args:
              cur_state (State)
              action (str)
  
          Returns:
              (float): Discounted reward from the rollout.
          '''
          trajectory = []
          total_discounted_reward = []
          for i in range(self.rollout_depth):
  
              # Simulate next state.
              next_action = self._next_action(cur_state)
              cur_state = self.transition_func(cur_state, next_action)
              next_reward = self.reward_func(cur_state, next_action)
  
              # Track rewards and states.
              total_discounted_reward.append(self.gamma**i * next_reward)
              trajectory.append((cur_state, next_action))
  
              if cur_state.is_terminal():
                  # Break terminal.
                  break
  
          # Update all visited nodes.
          for i, experience in enumerate(trajectory):
              s, a = experience
              self.visitation_counts[s][a] += 1
              self.value_total[s][a] += sum(total_discounted_reward[i:])
  
          return total_discounted_reward
  
  
  