text: from typing import Callable, Any, Optional
  from ai_controller.controller import AIController
  from ai_controller.mcts.jit import mcts
  from environment.game_state import GameState
  from .node import MCTSChildNode, MCTSRootNode, MCTSNode
  from typing import List, Tuple
  import shutil
  import os
  from pathos.multiprocessing import ProcessPool
  import pathlib
  
  
  class MCTSController(AIController):
    """
    A controller that implements MCTS with pluggable tree, action-selection & default policy
    """
    def __init__(self,
      tree_policy : Callable[[MCTSNode, bool], MCTSChildNode],
      default_policy : Callable[[MCTSNode], Tuple[List[float], float, List[float]]],
      action_selection_policy : Callable[[List[MCTSRootNode]], Any],
      num_iterations : int = 100,
      output_tree_diagram : bool = False,
      simulation_result_discount_strategy : Optional[str] = None,
      num_output_diagrams : int = 10,
      parallel : bool = False,
      num_process : int = 8,
      ply_discount_factor : float = 0.95,
      likelihood_discount_factor : float = 0.8,
      with_normalisation : bool = False
    ):
      self._tree_policy = tree_policy
      self._default_policy = default_policy
      self._action_selection_policy = action_selection_policy
  
      self._num_iterations = num_iterations
      self._output_tree_diagram = output_tree_diagram
      self._simulation_result_discount_strategy = simulation_result_discount_strategy
      self._num_output_diagrams = num_output_diagrams
      self._parallel = parallel
      self._num_process = num_process
  
      self._ply_discount_factor : float = ply_discount_factor
      self._likelihood_discount_factor = likelihood_discount_factor
      self._with_normalisation = with_normalisation
  
      if output_tree_diagram is True:
        self._clean_output_directory()
  
  
    def _clean_output_directory(self) -> None:
      if os.path.exists("output/mcts"):
        shutil.rmtree("output/mcts")
      pathlib.Path("output/mcts").mkdir(parents=True, exist_ok=True)
  
  
    def pick_move_with_likelihood(self, state : GameState) -> Tuple[List[float], float, List[float]]:
      if self._parallel:
        iteration_per_process : int = (self._num_iterations // self._num_process) + 1
        pool = ProcessPool()
        args : List[Tuple] = list(zip(*[[state, iteration_per_process, self._num_output_diagrams, self._output_tree_diagram, self._default_policy, self._tree_policy, self._simulation_result_discount_strategy, self._with_normalisation]] * self._num_process))
        roots = pool.map(mcts, *args)
  
      else:
        root = mcts(state, self._num_iterations, self._num_output_diagrams, self._output_tree_diagram, self._default_policy, self._tree_policy, self._simulation_result_discount_strategy, self._ply_discount_factor, self._likelihood_discount_factor, self._with_normalisation)
        roots = [root]
  
      move = self._action_selection_policy(roots)
      return move
  