Tree.py

"""
Defines the Tree of the MCTS: main object for RetroPath 3
"""

# General utilities
import os
import sys
import time
import logging
import argparse
import pickle
import string
import numpy as np
import csv

import random

# RP3 specific packages
from compound import (
    Compound,
    unpickle,
    CompoundDefinitionException,
    ChemConversionError,
)
from chemical_compounds_state import ChemicalCompoundState
from MCTS_node import MCTS_node
from representation import Test_representation, Test_to_file
from UCT_policies import (
    Biochemical_UCT_1,
    Nature_UCT,
    Classical_UCT_RAVE,
    Classical_UCT_with_bias,
    Classical_UCT,
    Biochemical_UCT_1_with_RAVE,
    Biochemical_UCT_with_progressive_bias,
    Chemical_UCT_1,
    Biological_UCT_1,
    Biochemical_UCT_with_toxicity,
)
from Rollout_policies import (
    Rollout_policy_first,
    Rollout_policy_random_uniform_on_biochemical_multiplication_score,
    Rollout_policy_random_uniform,
)
from rewarding import Basic_Rollout_Reward, RolloutRewards
from rule_sets_examples import applicable_rules_10_dict
from rule_sets_similarity import (
    get_rules_and_score,
    full_rules_forward_H,
    full_rules_retro_H,
    full_rules_forward_no_H,
    full_rules_retro_no_H,
)
from pathway import Pathway
from pathway_scoring import (
    RandomPathwayScorer,
    constant_pathway_scoring,
    null_pathway_scoring,
    biological_pathway_scoring,
    biochemical_pathway_scoring,
)
from tree_viewer import Tree_viewer
from organisms import (

    # With Hs
    bsubtilis_H,
    core_ecoli_H,
    detectable_cmpds_H,
    ecoli_chassis_H,
    iJO1366_chassis_H,
    Test_organism_H,

    # Without Hs
    core_ecoli_noH,
    bsubtilis_noH,
    detectable_cmpds_noH,
    ecoli_chassis_noH,
    iJO1366_chassis_noH,
    Test_organism_noH,

    # Import function
    import_organism_from_csv,
)

# General Configuration
from config import *
if use_toxicity:
    from compound_scoring import toxicity_scorer

sys.setrecursionlimit(50000)


class RunModeError(Exception):
    """Class for Tree Mode exception."""

    def __init__(self, retrosynthesis, biosensor):
        self.message = "Choose between retrosynthesis ({}) and biosensor ({})".format(retrosynthesis, biosensor)
        self.reason = "RunModeError"


class IncorrectTreeLoading(Exception):
    """Class for Conflicts between trees when loading a tree for search extension."""

    def __init__(self, message):
        self.message = message
        self.reason = "IncorrectTreeLoading"


class myTimeoutException(Exception):
    """Class for Tree Timeout exception."""

    def __init__(self, iteration=-1):
        self.iteration = iteration
        self.time = time.time()
        self.message = "Timeout at time {}, at iteration {}".format(self.time, self.iteration)


class FirstResultException(Exception):
    """Class for Tree First Result Stop exception."""

    def __init__(self, file_to_save, name, time_spent, iteration=-1):
        self.message = "First result after {}s for compound {}".format(time_spent, name)
        self.time = time_spent
        self.iteration = iteration
        with open(file_to_save, "w") as writer:
            writer.write("{};{}".format(name, round(time_spent, 2)))


class CompoundInSink(Exception):
    """Class for raising exception if compound already in sink."""

    def __init__(self, folder_to_save, name):
        self.message = "Compound {} already in organism".format(name)
        self.reason = "CompoundInSink"
        with open("{}/in_sink".format(folder_to_save), "w") as results_file:
            pass


class InvalidSink(Exception):
    """Class for raising exception if sink is invalid."""

    def __init__(self):
        self.reason = "InvalidSink"
        self.message = "Invalid (empty) sink"


class Tree:
    """
    Defines a search on the Tree according to
    - defined policy
    - organism
    Useful methods include:
    - running the search (key method of this object)
    - extracting the best pathway
    - extracting all pathways
    - visualising the tree
    - extending the tree with more tolerant settings or simply more iterations.
    - tree summary
    Numerous utilities has been developped for exporting results.
    """

    logger = logging.getLogger(__name__)

    def __init__(
        self,
        root_state,  # starting chemical compound state (contains the target Compound only)
        itermax=2,  # Iteration budget
        time_budget=20,  # Time budget
        expansion_width=10,  # Number of allwoed children
        available_rules=None,  # Rule dataset that can be used
        rewarding=Basic_Rollout_Reward,  # How to reward the chemical state after rollout
        max_depth=5,  # Maximum number of reactions to be used
        max_rollout=3,
        number_rollout=1,  # if parallel, number of different rollouts performed. NOT USED
        UCT_policy="Classical_UCT",
        UCT_parameters={"UCTK": 2, "bias_k": 0},
        Rollout_policy="Rollout_policy_first",
        Rollout_parameters={"rollout_timeout": 0.2},
        print_tree_information=True,
        parallel=False,
        main_layer_tree=False,
        main_layer_chassis=True,
        organism=ecoli_chassis_H,
        chemical_scorer="RandomChemicalScorer",
        biological_scorer="RandomBiologicalScorer",
        folder_to_save="tests/generated_jsons",
        virtual_visits=0,
        progressive_bias_strategy=None,
        progressive_widening=False,
        heavy_saving=False,
        minimal_visit_counts=1,
        use_RAVE=False,
        pathway_scoring="null_pathway_scoring",
    ):
        # Chemical equality between nodes in the stree: only main layer or full inchi
        self.main_layer_tree = main_layer_tree
        self.main_layer_chassis = main_layer_chassis
        self.organism = organism
        self.organism.set_main_layer(main_layer_chassis)
        self.root_state = root_state
        self.root_state.set_organism(self.organism)
        self.root_state.set_main_layer(self.main_layer_tree)  # Chassis only used in the chassis
        # Defaut rules for pytest - small subset
        if available_rules is None:
            self.available_rules = applicable_rules_10_dict
        else:
            self.available_rules = available_rules
        # Define rules and scores settings for the root state.
        self.root_state.set_available_rules(self.available_rules)
        self.root_state.set_chemical_scorer(chemical_scorer)
        self.root_state.set_biological_scorer(biological_scorer)

        self.scorer_info = (
            "Chemical scoring is {}\n. Biological scoring is {}\n".format(
                chemical_scorer, biological_scorer
            )
        )
        # Define initial node of the Tree with all previous information
        self.root_node = MCTS_node(
            root_state,
            rewarding=rewarding,
            expansion_width=expansion_width,
            maximum_depth=max_depth,
            main_layer_tree=self.main_layer_tree,
            main_layer_chassis=self.main_layer_chassis,
            chemical_scorer=chemical_scorer,
            biological_scorer=biological_scorer,
            virtual_visits=virtual_visits,
            progressive_bias_strategy=progressive_bias_strategy,
            minimal_visit_counts=minimal_visit_counts,
            use_RAVE=use_RAVE,
        )
        # Defining attributes from init
        self.expansion_width = expansion_width
        self.progressive_widening = progressive_widening
        self.itermax = itermax
        self.time_budget = time_budget
        self.rewarding = rewarding
        self.UCT_policy = UCT_policy
        self.UCT_parameters = UCT_parameters
        self.max_depth = max_depth
        self.max_rollout = max_rollout
        self.Rollout_policy = Rollout_policy
        self.Rollout_parameters = Rollout_parameters
        self.pathway_scoring = pathway_scoring
        # Saving settings
        self.folder_to_save = folder_to_save
        self.folder_to_save_pickles = "{}/{}".format(folder_to_save, "pickles")
        self.heavy_saving = heavy_saving

        self.print_tree_information = True

        self.parallel = parallel
        if self.parallel:
            # Should not be used at the moment.
            # Used to conduct multiple rollouts in parallel.
            # Impossible with current RDKit lack of timeout
            self.worker = worker_Rollout_with_array
            self.number_rollout = number_rollout
        else:
            self.number_rollout = 1  # When not doing in parallel
        self.recap = None  # Used for storing results and information that will be outputed in simplified foramt.

    def set_folder_to_save(self, folder_to_save):
        self.folder_to_save = folder_to_save
        self.folder_to_save_pickles = "{}/{}".format(folder_to_save, "pickles")

    def set_heavy_saving(self, heavy_saving):
        self.heavy_saving = heavy_saving

    def set_rules(self, rules):
        self.available_rules = rules
        # Change also for all nodes.
        nodes_to_treat = [self.root_node]
        while nodes_to_treat != []:
            node = nodes_to_treat[0]
            del nodes_to_treat[0]
            node.state.set_available_rules(rules)
            for child in node.children:
                nodes_to_treat.append(child)

    def __len__(self):
        length = 0
        nodes_to_treat = [self.root_node]
        while nodes_to_treat != []:
            node = nodes_to_treat[0]
            del nodes_to_treat[0]
            length = length + 1
            for child in node.children:
                nodes_to_treat.append(child)
        return length

    def set_itermax(self, itermax):
        self.itermax = itermax

    def run_search(self, stop_at_first_result=False):
        """
        Main method for using RetroPath 3.
        Runs a tree search for itermax iterations:
        - starts from root
        - while node have children, select the best
        - expand the leaf if possible. Select the best child of the expanded node
        - rollout from this node
        - update all results
        - save pathways as they are found
        """
        UCT_policy = self._get_UCT_policy()
        pathway_scoring = self._get_pathway_scoring()
        if self.print_tree_information:
            self._print_tree_information_beginning()
        # Signal below is dropped in favour of checking at the beginning of an iteration if there is time left.
        # signal.signal(signal.SIGALRM, timeout_tree_search)
        # signal.alarm(self.time_budget)
        selection_time = 0
        expansion_time = 0
        rollout_time = 0
        update_time = 0
        raise_stop = False
        start_tree_search_time = time.time()
        try:
            for i in range(
                self.root_node.visits - self.root_node.virtual_visits,
                self.itermax + self.root_node.visits - self.root_node.virtual_visits,
            ):
                # exploration and EXPANSION
                if time.time() - start_tree_search_time > self.time_budget:
                    raise myTimeoutException(iteration=i)
                self.logger.debug("Currently performing iteration {}".format(i))
                if i % int(round(self.itermax / 10)) == 0:
                    self.logger.info("Performing iteration {} out of {}".format(i, self.itermax))
                    if self.heavy_saving:
                        # Saving Tree for further analysis during the search.
                        self.find_full_tree(iteration=i)
                        self.save(
                            file_name="{}_iteration_{}".format(
                                self.root_state.compound_list[0], i
                            ),
                            folder_address=self.folder_to_save_pickles,
                        )
                        self.jsonify_full_tree(
                            file_name="{}_iteration_{}".format(
                                self.root_state.compound_list[0], i
                            )
                        )
                # A new iteration starts from the root node.
                node = self.root_node
                start_selection_time = time.time()
                while node.children != [] and not node.terminal:
                    if node.flag_for_extension:
                        # When extending a previous tree.
                        node.expand_later_iteration(
                            extension_length=node.flag_extension_length,
                            maximum_depth=node.flag_maximum_depth,
                            chemical_scoring_configuration=node.flag_chemical_scoring_configuration,
                        )

                    if not self.progressive_widening and node.moves == []:
                        # Go down the tree while all moves from a node has been expanded.
                        node = node.UCTSelectChild(UCT_policy=UCT_policy)
                        state = node.state
                    elif self.progressive_widening:
                        if node.moves == []:  # No moves are available for PW.
                            node = node.UCTSelectChild(UCT_policy=UCT_policy)
                            state = node.state
                        else:  # there are still available moves, apply PW condition
                            if node.visits <= len(node.children) ^ 2:
                                node = node.UCTSelectChild(UCT_policy=UCT_policy)
                                state = node.state
                            else:
                                break
                    else:
                        break
                selection_time = selection_time + time.time() - start_selection_time
                start_expansion_time = time.time()
                if node.moves != [] and not node.terminal:  # This node can be expanded.
                    move = node.moves[0]
                    looping_move = False
                    for product in move.product_list:
                        if product.in_list(node.history.compound_list, main_layer=node.main_layer_chassis):
                            looping_move = True
                            self.logger.info("Looping move in node {}-{} for move {}".format(node.state, node.level, move))
                            node.moves.remove(move)
                            break
                    if not looping_move:
                        node.AddChild(move, node.state)
                        node = node.UCTSelectChild(UCT_policy=UCT_policy)
                        state = node.state

                expansion_time = expansion_time + time.time() - start_expansion_time
                # Perform rollout on the selected node
                if not node.terminal:
                    start_rollout_time = time.time()
                    allowed_rollout = min(self.max_rollout, self.max_depth - node.level)
                    if self.parallel:
                        # Not in use at the moment.
                        # Made for performing multiple rollouts on the same node.
                        # Was temporarily suspended due to the necessity of encapsulating Rdkit preocesses.
                        shared_array_for_rollout = mp.Array('d', self.number_rollout)
                        kwargs_local = {"state": node.state,
                                        "rewarding": self.rewarding,
                                        "rollout_number": allowed_rollout,
                                        "RolloutPolicy_name": self.Rollout_policy,
                                        "main_layer": self.main_layer_chassis}
                        _pool = mp.Pool(processes=4, initializer=worker_init_array,
                                        initargs=(shared_array_for_rollout, kwargs_local))
                        try:
                            _pool.map(worker_Rollout_with_array, range(self.number_rollout))
                        except mp.TimeoutError:
                            self.logger.warning("Time out in parllel workers for node \n{}".format(node))
                            kill(_pool)
                        result, solved = np.mean(shared_array_for_rollout[:]), self.rewarding.full_state_reward in shared_array_for_rollout[:]
                    else:
                        # Perform rollout without parallelisation.
                        RolloutPolicy = self._get_Rollout_policy()
                        state = node.rollout(rollout_number=allowed_rollout, RolloutPolicy=RolloutPolicy)

                        if retrosynthesis:
                            result = state.GetResults_from_InChI_Keys(
                                rewarding=self.rewarding,
                                main_layer=self.main_layer_chassis,
                            )
                            solved = (
                                state.GetResults_from_InChI_Keys(
                                    rewarding=self.rewarding,
                                    main_layer=self.main_layer_chassis,
                                )
                                == self.rewarding.full_state_reward
                            )
                        elif biosensor:
                            result = state.GetResultsForBiosensors(
                                rewarding=self.rewarding,
                                main_layer=self.main_layer_chassis,
                            )
                            solved = (
                                state.GetResultsForBiosensors(
                                    rewarding=self.rewarding,
                                    main_layer=self.main_layer_chassis,
                                )
                                == self.rewarding.full_state_reward
                            )
                        else:
                            raise NotImplementedError
                    self.logger.info("Rollout results are {}, {}".format(result, solved))
                    rollout_time = rollout_time + time.time() - start_rollout_time
                else:
                    # Node is already terminal.
                    start_rollout_time = time.time()
                    state = node.state
                    if retrosynthesis:
                        result = state.GetResults_from_InChI_Keys(
                            rewarding=self.rewarding, main_layer=self.main_layer_chassis
                        )
                        solved = (
                            state.GetResults_from_InChI_Keys(
                                rewarding=self.rewarding,
                                main_layer=self.main_layer_chassis,
                            )
                            == self.rewarding.full_state_reward
                        )
                    elif biosensor:
                        result = state.GetResultsForBiosensors(
                            rewarding=self.rewarding, main_layer=self.main_layer_chassis
                        )
                        solved = (
                            state.GetResultsForBiosensors(
                                rewarding=self.rewarding,
                                main_layer=self.main_layer_chassis,
                            )
                            == self.rewarding.full_state_reward
                        )
                    else:
                        raise NotImplementedError
                    if solved and not node.has_a_solved_child:
                        # Extract pathway if encountered from the first time and fully solved.
                        self.logger.info("Pathway is found at iteration {}".format(i))
                        found_pathway = self.extract_pathway_from_bottom(node, iteration=i)
                        if len(found_pathway) == 1:
                            name = str(self.root_state.compound_list[0]) + "iteration_{}".format(i)
                            if self.heavy_saving:
                                found_pathway[0]["pathway"].save(name, folder_address=self.folder_to_save_pickles)
                            result = pathway_scoring.calculate(found_pathway[0]["pathway"])
                        else:
                            best_path = None
                            for path in found_pathway:
                                name = str(
                                    self.root_state.compound_list[0]
                                ) + "iteration_{}_{}".format(i, path["number"])
                                if self.heavy_saving:
                                    path["pathway"].save(name, folder_address=self.folder_to_save_pickles)
                                result = pathway_scoring.calculate(path["pathway"])
                                if best_path is None:
                                    best_path = path
                                    best_result = result
                                else:
                                    if best_result < result:
                                        best_path = path
                                        best_result = result
                        if stop_at_first_result:
                            raise_stop = True

                # Backpropagate:
                start_update_time = time.time()
                while node is not None:  # backpropagate from the expanded node and work back to the root node
                    node.update(result,
                                solved=solved,
                                visit_number=self.number_rollout,
                                iteration=i)
                    node = node.parent
                update_time = update_time + time.time() - start_update_time
                if raise_stop:
                    time_spent = time.time() - start_tree_search_time
                    raise FirstResultException(
                        file_to_save="{}/{}_execution_time.txt".format(self.folder_to_save, self.root_state.state_name),
                        name=self.root_state.state_name,
                        time_spent=time_spent,
                        iteration=i)
        except myTimeoutException as e:
            self.logger.warning(
                "Timeout happened on Tree search after {} iterations for a budget of {}s".format(i, self.time_budget))

            self.recap = {"TIME_EXECUTION": round(e.time - start_tree_search_time, 2),
                          "STOP_REASON": "timeout",
                          "NUMBER_ITERATION": e.iteration}
        except FirstResultException as e:
            self.logger.warning(e.message)
            self.recap = {"TIME_EXECUTION": round(e.time, 2),
                          "STOP_REASON": "first_result",
                          "NUMBER_ITERATION": e.iteration}
        if self.print_tree_information:
            self._print_tree_information_end(selection_time, expansion_time, rollout_time, update_time)
        logging.info("At the end of this tree search, the cache contains {}".format(len(home_made_cache.keys())))
        if self.recap is None:
            self.recap = {"TIME_EXECUTION": round(time.time() - start_tree_search_time, 2),
                          "STOP_REASON": "iteration",
                          "NUMBER_ITERATION": i}

        self.logger.warning("RECAP TIME_EXECUTION={}".format(self.recap["TIME_EXECUTION"]))
        self.logger.warning("RECAP STOP_REASON={}".format(self.recap["STOP_REASON"]))
        self.logger.warning("RECAP NUMBER_ITERATION={}".format(self.recap["NUMBER_ITERATION"]))

        if self.heavy_saving:
            self.find_full_tree(iteration=i)
            self.save(file_name="{}_iteration_{}".format(self.root_state.compound_list[0], i),
                      folder_address=self.folder_to_save_pickles)
            self.jsonify_full_tree(file_name="{}_iteration_{}".format(self.root_state.compound_list[0], i))

    def _get_UCT_policy(self):
        if self.UCT_policy == "Classical_UCT":
            policy = Classical_UCT(self.UCT_parameters)
        elif self.UCT_policy == "Classical_UCT_with_bias":
            policy = Classical_UCT_with_bias(self.UCT_parameters)
        elif self.UCT_policy == "Classical_UCT_RAVE":
            policy = Classical_UCT_RAVE(parameters=self.UCT_parameters)
        elif self.UCT_policy == "Nature_UCT":
            policy = Nature_UCT(self.UCT_parameters)
        elif self.UCT_policy == "Biochemical_UCT_1":
            policy = Biochemical_UCT_1(self.UCT_parameters)
        elif self.UCT_policy == "Biochemical_UCT_with_progressive_bias":
            policy = Biochemical_UCT_with_progressive_bias(self.UCT_parameters)
        elif self.UCT_policy == "Biochemical_UCT_1_with_RAVE":
            policy = Biochemical_UCT_1_with_RAVE(self.UCT_parameters)
        elif self.UCT_policy == "Chemical_UCT_1":
            policy = Chemical_UCT_1(self.UCT_parameters)
        elif self.UCT_policy == "Biological_UCT_1":
            policy = Biological_UCT_1(self.UCT_parameters)
        elif self.UCT_policy == "Biochemical_UCT_with_toxicity":
            policy = Biochemical_UCT_with_toxicity(self.UCT_parameters)
        else:
            raise NotImplementedError(self.UCT_policy)
        return policy

    def _get_Rollout_policy(self):
        if self.Rollout_policy == "Rollout_policy_first":
            policy = Rollout_policy_first()
        elif self.Rollout_policy == "Rollout_policy_chemical_best":
            policy = Rollout_policy_chemical_best()
        elif self.Rollout_policy == "Rollout_policy_biological_best":
            policy = Rollout_policy_biological_best()
        elif self.Rollout_policy == "Rollout_policy_biochemical_addition_best":
            policy = Rollout_policy_biochemical_addition_best()
        elif self.Rollout_policy == "Rollout_policy_biochemical_multiplication_best":
            policy = Rollout_policy_biochemical_multiplication_best()
        elif self.Rollout_policy == "Rollout_policy_random_uniform":
            policy = Rollout_policy_random_uniform()
        elif self.Rollout_policy == "Rollout_policy_random_uniform_on_chem_score":
            policy = Rollout_policy_random_uniform_on_chem_score()
        elif self.Rollout_policy == "Rollout_policy_random_uniform_on_bio_score":
            policy = Rollout_policy_random_uniform_on_bio_score()
        elif self.Rollout_policy == "Rollout_policy_random_uniform_on_biochemical_addition_score":
            policy = Rollout_policy_random_uniform_on_biochemical_addition_score()
        elif self.Rollout_policy == "Rollout_policy_random_uniform_on_biochemical_multiplication_score":
            policy = Rollout_policy_random_uniform_on_biochemical_multiplication_score()
        else:
            raise NotImplementedError(self.RolloutPolicy)
        return policy

    def _get_pathway_scoring(self):
        self.pathway_scoring_logs = False
        if self.pathway_scoring == "RandomPathwayScorer":
            pathway_scoring = RandomPathwayScorer
        elif self.pathway_scoring == "constant_pathway_scoring":
            pathway_scoring = constant_pathway_scoring
        elif self.pathway_scoring == "biological_pathway_scoring":
            pathway_scoring = biological_pathway_scoring
        elif self.pathway_scoring == "null_pathway_scoring":
            pathway_scoring = null_pathway_scoring
        else:
            raise NotImplementedError(self.pathway_scoring)
        return pathway_scoring

    def _print_tree_information_beginning(self):
        test_block = "---------------Beginning of search informations------------------ \n"
        text_block = test_block + "UCT policy is : {} \n".format(str(self.UCT_policy))
        text_block = text_block + "Rollout policy is : {} \n".format(str(self.Rollout_policy))
        text_block = text_block + "Rewarding policy is : {} \n".format(str(self.rewarding))
        text_block = text_block + 'Maximum number of iterations {} \n'.format(self.itermax)
        start_time = time.localtime()
        self.starttime = time.time()
        text_block = text_block + 'Tree search started running at {}h {}min {}s \n'.format(start_time.tm_hour,
                                                                                           start_time.tm_min,
                                                                                           start_time.tm_sec)
        text_block = text_block + 'Your tree search is currently running ... \n \n'
        self.logger.info(text_block)

    def _print_tree_information_end(self, selection_time=0, expansion_time=0, rollout_time=0, update_time=0):
        end_time = time.localtime()
        test_block = "---------------End of search informations------------------ \n"
        text_block = test_block + 'Tree search finished running at {}h {}min {}s \n'.format(end_time.tm_hour,
                                                                                            end_time.tm_min,
                                                                                            end_time.tm_sec)
        text_block = text_block + 'Your tree search took {:.1e}s \n \n'.format(time.time() - round(self.starttime))
        text_block = text_block + 'Time was spent in selection: {:.1e}s, expansion: {:.1e}s, rollout: {:.1e}s, update: {:.1e}s \n'.format(
            selection_time, expansion_time, rollout_time, update_time)
        if self.root_node.has_a_solved_child:
            text_block = text_block + "CONGRATULATIONS, THATS A WIN"
        else:
            text_block = text_block + "You just lost the game"
        self.logger.info(text_block)

    def __repr__(self):
        """
        Represents the tree horizontally for small cases
        """
        rep = ''
        node_list = [self.root_node]
        while node_list != []:
            current_node = node_list[-1]
            del node_list[-1]
            rep = rep + str(current_node.level) + str(current_node) + "\n"
            for child in current_node.children:
                node_list.append(child)
        return rep

    def __eq__(self, other):
        """
        Verifying whether 2 Trees are identical. Will be useful for comparison of UCT methods
        - Trees are identical if all their nodes are recursively identical
        """
        self_pile = [self.root_node]
        other_pile = [other.root_node]
        equal = (self.root_node == other.root_node)
        if not equal:
            return False
        while self_pile != [] and other_pile != [] and equal:
            current_node_self = self_pile[-1]
            del self_pile[-1]
            self_pile = self_pile + current_node_self.children
            found = False
            for node in other_pile:
                if node == current_node_self:
                    current_node_other = node
                    found = True
            if not found:
                # Quits the loop with the return statement.
                return False
            other_pile.remove(current_node_other)
            equal = (current_node_other == current_node_self)  # it's been selected for this
            other_pile = other_pile + current_node_other.children
        return equal

    def jsonify_full_tree(self, file_name=None, folder_to_save=None):
        """
        - prints the whole tree. Useful to see where the search got lost.
        - uses the json format for node and moves
        - create a Tree viewer object
        """

        if file_name is None:
            file_name = self.root_state.state_name
        if folder_to_save is None:
            folder_to_save = self.folder_to_save
        file_to_save = folder_to_save + "/{}_full_tree_for_MCTS.json".format(file_name)

        #  First, assign IDs to all moves and compounds
        id_move = 0
        id_node = 0
        nodes_to_treat = [self.root_node]
        while nodes_to_treat != []:
            node = nodes_to_treat[0]
            del nodes_to_treat[0]
            if node.move is not None:
                node.move.set_id(id_move)
                id_move = id_move + 1
            node.set_id(id_node)
            id_node = id_node + 1
            for child in node.children:
                nodes_to_treat.append(child)

        tree_viewer = Tree_viewer(file_to_save=file_to_save)
        nodes_to_treat = [self.root_node]
        while nodes_to_treat != []:
            node = nodes_to_treat[0]
            del nodes_to_treat[0]
            tree_viewer.add_node(node)
            for child in node.children:
                nodes_to_treat.append(child)
        tree_viewer.jsonify_tree_viewer()

    def save(self, file_name=None, folder_address="pickled_data"):
        if self.__class__.__module__ == "__main__":
            self.logger.warning("Tree should not be in main for pickling")
        if file_name is None:
            base_name = self.root_state.state_name
            i = 1
            self.logger.info("Tree is saved at {}/tree_{}_{}.pkl".format(folder_address, base_name, i))
            while os.path.exists("{}/tree_{}_{}.pkl".format(folder_address, base_name, i)):
                i += 1
            file_name = base_name + '_' + str(i)
        file_saving = open('{}/tree_{}.pkl'.format(folder_address, file_name), 'wb')
        pickle.dump(self, file_saving)

    def equality_visited_states(self, other):
        """
        Verifying whether 2 Trees explored the same states. Will be useful for comparison of UCT methods
        - Trees are identical if all their nodes are recursively identical
        - contrary to identify function, checks states but not visit counts.
        """
        self_pile = [self.root_node]
        other_pile = [other.root_node]
        equal = (self.root_node.state == other.root_node.state)
        if not equal:
            return False
        while self_pile != [] and other_pile != [] and equal:
            current_node_self = self_pile[-1]
            del self_pile[-1]
            self_pile = self_pile + current_node_self.children
            found = False
            for node in other_pile:
                if node.state == current_node_self.state:
                    current_node_other = node
                    found = True
            if not found:
                return False
            other_pile.remove(current_node_other)
            equal = (current_node_other.state == current_node_self.state)
            other_pile = other_pile + current_node_other.children
        return equal

    def extract_pathway_from_bottom(self, node, name=None, iteration=-1):
        """
        Extract pathways from found state (and not from the top down)
        """
        bottom_up_pathway = Pathway(first_iteration=iteration, target=None,
                                    compounds=[], moves=[],
                                    main_layer=self.main_layer_chassis,
                                    organism=self.organism,
                                    edges=[],
                                    nodes_compounds=[],
                                    nodes_transformations=[])
        target = self.root_state.compound_list[0]
        bottom_up_pathway.add_compound(target, is_source=1)
        if name is None:
            name = str(target)
        file_to_save = self.folder_to_save + "/{}_iteration_{}.json".format(name, iteration)
        bottom_up_pathway.set_file_to_save(file_to_save)
        if use_transpositions:
            complete_pathways = []
            pathways_to_treat = [{"pathway": bottom_up_pathway, "number": "a", "node_list": [node]}]
            while pathways_to_treat != []:
                current_pathway_node = pathways_to_treat[0]
                del pathways_to_treat[0]
                pathway = current_pathway_node["pathway"].clone()
                number = current_pathway_node["number"]
                node_list = current_pathway_node["node_list"]
                node = node_list[-1]
                if not node.parent is None:
                    node_parents = transposition_table[node.parent.hash]
                    if len(node_parents) == 1:
                        node_list.append(node.parent)
                        new_pathway = {"pathway": pathway,
                                       "number": number,
                                       "node_list": node_list}
                        pathways_to_treat.append(new_pathway)
                    else:
                        for i in range(len(node_parents)):
                            parent = node_parents[i]
                            new_node_list = copy.deepcopy(node_list)
                            new_pathway = pathway.clone()
                            new_node_list.append(parent)
                            new_pathway = {"pathway": new_pathway,
                                           "number": number + '_' + list(string.ascii_lowercase)[i],
                                           "node_list": new_node_list}
                            pathways_to_treat.append(new_pathway)
                else:
                    complete_pathways.append(current_pathway_node)
            for pathway_as_dict in complete_pathways:
                pathway = pathway_as_dict["pathway"]
                number = pathway_as_dict["number"]
                node_list = pathway_as_dict["node_list"]
                node_list.reverse()
                for top_down_node in node_list[1:]:
                    pathway.add_reaction(top_down_node.move, depth=top_down_node.level)
                file_to_save = self.folder_to_save + "/{}_iteration_{}_{}.json".format(name, iteration, number)
                pathway.set_file_to_save(file_to_save)
                pathway.jsonify_scope_viewer()
            return complete_pathways
        else:
            nodelist = [node]
            while node.parent is not None:
                # Going up the tree to get the parents
                nodelist.append(node.parent)
                node = node.parent
            nodelist.reverse()
            for top_down_node in nodelist[1:]:
                bottom_up_pathway.add_reaction(top_down_node.move, depth=top_down_node.level)
            bottom_up_pathway.set_file_to_save(file_to_save)
            bottom_up_pathway.jsonify_scope_viewer()
            return [{"pathway": bottom_up_pathway}]

    def find_single_best_pathway(self, policy="visits", folder_to_save="pathways", name=None):
        """
        Currently, best according to number of visits.
        """

        best_pathway = Pathway(
            first_iteration=-1,
            target=None,
            compounds=[],
            moves=[],
            main_layer=self.main_layer_chassis,
            organism=self.organism,
            edges=[],
            nodes_compounds=[],
            nodes_transformations=[],
        )
        target = self.root_state.compound_list[0]
        best_pathway.add_compound(target, is_source=1)
        if policy == "visits":
            self.logger.info(
                "According to Wikipedia, choosing the pathway with the most visists \n This could be replaced by the best score")
        if name is None:
            name = str(target)
        file_to_save = folder_to_save + "/{}_best.json".format(name)
        best_pathway.set_file_to_save(file_to_save)
        node = self.root_node
        while not node.terminal:
            # Could be modified with children cashing
            node = node.SelectBestChild(policy=policy)
            best_pathway.add_reaction(node.move, depth=node.level)
        best_pathway.jsonify_scope_viewer()

    def find_multiple_best_pathways(self, folder_to_save="pathways",
                                    name=None, return_result=False,
                                    return_pathways = False, biochem_sorted = False):
        """
        Returns all solved pathways during the search.
        They are sorted according to biochemical score.
        """
        # Have a pile of patwhays to test and add. Only add patwhays that are fully solved. Export them all
        pathways_to_print = []
        pathway_iteration = 1
        if self.root_node.has_a_solved_child:
            initial_pathway = Pathway(
                first_iteration=-1,
                target=None,
                compounds=[],
                moves=[],
                main_layer=self.main_layer_chassis,
                organism=self.organism,
                edges=[],
                nodes_compounds=[],
                nodes_transformations=[],
            )

            target = self.root_state.compound_list[0]
            initial_pathway.add_compound(target, is_source=1)
            if name is None:
                name = str(target)
            # Start by treating the initla pathway
            max_number = 0
            pathways_to_treat = [{"pathway": initial_pathway, "node": self.root_node, "number": 1}]
            while pathways_to_treat != []:
                # Obtain and delete patwhay
                current_pathway_node = pathways_to_treat[0]
                del pathways_to_treat[0]
                node = current_pathway_node["node"]
                pathway = current_pathway_node["pathway"]
                number = current_pathway_node["number"]
                if node.move is not None:
                    pathway.add_reaction(node.move, depth=node.level)
                if node.terminal:
                    # If the node is terminal, add the creating move in the pwathway and stop there
                    if pathway in pathways_to_print:
                        self.logger.debug("Pathway is already present: \n {}".format(pathway))
                    else:
                        pathways_to_print.append(pathway)
                else:
                    # Extracting the nodes that have solved children
                    nodes_to_treat = []
                    for child in node.children:
                        if child.has_a_solved_child:
                            nodes_to_treat.append(child)
                    children_with_solved_num = len(nodes_to_treat)
                    for i in range(len(nodes_to_treat)):
                        # for i in range(len(nodes_to_treat)):
                        new_pathway = pathway.clone()
                        new_number = max_number + i + number
                        pathways_to_treat.append(
                            {"pathway": new_pathway, "node": nodes_to_treat[i], "number": new_number})
                    max_number = new_number + 1
                current_pile_length = len(pathways_to_treat)
            self.logger.info("Found {} pathways for {}".format(len(pathways_to_print), str(target)))
        else:
            pathways_to_print = []
            self.logger.info("No pathway was found for this compound")
        if biochem_sorted:
            sorted_pathways_couples = [(pathway, round(biochemical_pathway_scoring.calculate(pathway), 2)) for pathway in pathways_to_print]
            sorted_pathways = sorted(sorted_pathways_couples, key=lambda item: item[1], reverse=True)
            for j in range(len(sorted_pathways)):
                pathway = sorted_pathways[j][0]
                file_to_save = folder_to_save + "{}_{}.json".format(name, j + 1)
                pathway.set_file_to_save(file_to_save)
                pathway.jsonify_scope_viewer()
            sorted_pathways = [patwhay for patwhay, score in sorted_pathways]
        else:
            sorted_pathways = pathways_to_print
            for j in range(len(sorted_pathways)):
                pathway = sorted_pathways[j]
                file_to_save = folder_to_save + "/{}_{}.json".format(name, j + 1)
                pathway.set_file_to_save(file_to_save)
                pathway.jsonify_scope_viewer()
        if return_result:
            return len(sorted_pathways)
        if return_pathways:
            return sorted_pathways

    def find_full_scope(self, folder_to_save="pathways", name=None):
        """
        Returns the scope of the compound: all pathways leading to the chassis.
        """
        pile_to_treat = []
        pathways_to_print = []
        pathway_iteration = 1

        full_scope = Pathway(
            first_iteration=-1,
            target=None,
            compounds=[],
            moves=[],
            main_layer=self.main_layer_chassis,
            organism=self.organism,
            edges=[],
            nodes_compounds=[],
            nodes_transformations=[],
        )

        target = self.root_state.compound_list[0]
        full_scope.add_compound(target, is_source=1)
        if name is None:
            name = str(target)
        max_number = 0
        nodes_to_treat = [self.root_node]
        while nodes_to_treat != []:
            node = nodes_to_treat[0]
            del nodes_to_treat[0]
            if node.move is not None:
                full_scope.add_reaction(node.move, depth=node.level)
            if node.terminal:
                pass
            else:
                for child in node.children:
                    if child.has_a_solved_child:
                        nodes_to_treat.append(child)
        file_to_save = folder_to_save + "/{}_full_scope.json".format(name)
        full_scope.set_file_to_save(file_to_save)
        full_scope.jsonify_scope_viewer()

    def find_full_tree(self, folder_to_save=None, name=None, iteration=0):
        """
        Returns all compounds and moves visited during the search. Can be a big file.
        """
        pile_to_treat = []
        pathways_to_print = []
        if folder_to_save is None:
            folder_to_save = self.folder_to_save

        full_tree = Pathway(
            first_iteration=iteration,
            target=None,
            compounds=[],
            moves=[],
            main_layer=self.main_layer_chassis,
            organism=self.organism,
            edges=[],
            nodes_compounds=[],
            nodes_transformations=[],
        )

        target = self.root_state.compound_list[0]
        full_tree.add_compound(target, is_source=1)
        if name is None:
            name = str(target)
        max_number = 0
        nodes_to_treat = [self.root_node]
        while nodes_to_treat != []:
            node = nodes_to_treat[0]
            del nodes_to_treat[0]
            if node.move is not None:
                full_tree.add_reaction(node.move, depth=node.level)
            if node.terminal:
                pass
            else:
                for child in node.children:
                    nodes_to_treat.append(child)
        file_to_save = folder_to_save + "/{}_full_tree_{}.json".format(name, iteration)
        full_tree.set_file_to_save(file_to_save)
        full_tree.jsonify_scope_viewer()

    def find_all_scopes(self, folder_to_save="pathways", name=None):
        for depth in range(1, self.max_depth + 1):
            self.find_specific_depth_scope(folder_to_save=folder_to_save, name=name, depth=depth)

    def find_specific_depth_scope(self, folder_to_save="pathways", name=None, depth=3):
        """
        Finds the scope up to a specific depth (less than n chemical reactions)
        """
        pile_to_treat = []
        pathways_to_print = []
        pathway_iteration = 1

        full_scope = Pathway(
            first_iteration=-1,
            target=None,
            compounds=[],
            moves=[],
            main_layer=self.main_layer_chassis,
            organism=self.organism,
            edges=[],
            nodes_compounds=[],
            nodes_transformations=[],
        )

        target = self.root_state.compound_list[0]
        full_scope.add_compound(target, is_source=1)
        if name is None:
            name = str(target)
        max_number = 0
        nodes_to_treat = [self.root_node]
        while nodes_to_treat != []:
            node = nodes_to_treat[0]
            del nodes_to_treat[0]
            if node.move is not None:
                full_scope.add_reaction(node.move, depth=node.level)
            if node.terminal or node.level >= depth:
                pass
            else:
                for child in node.children:
                    if child.has_a_solved_child:
                        nodes_to_treat.append(child)
        file_to_save = folder_to_save + "/{}_depth_{}_scope.json".format(name, depth)
        full_scope.set_file_to_save(file_to_save)
        full_scope.jsonify_scope_viewer()

    def flag_nodes_for_extension(self, extension_length=10, maximum_depth=None, chemical_scoring_configuration = None):
        """
        Flag the nodes that will have to be extended from the Loaded Tree.
        Function used only where starting search from an already explored tree
        """
        total_len = len(self)
        logging.info("Length of tree is {}".format(total_len))
        flagged_nodes = 0
        if maximum_depth is None:
            maximum_depth = self.max_depth
        else:
            self.logger.info("Changing maximum depth from {} to {}".format(self.max_depth, maximum_depth))

        nodes_to_treat = [self.root_node]
        while nodes_to_treat != []:
            node = nodes_to_treat[0]
            del nodes_to_treat[0]
            flagged_nodes = flagged_nodes + 1
            for child in node.children:
                nodes_to_treat.append(child)
            node.flag_node_for_extension(
                extension_length=extension_length,
                maximum_depth=maximum_depth,
                chemical_scoring_configuration=chemical_scoring_configuration,
            )
        assert flagged_nodes == total_len
        logging.info("Flagged {} nodes".format(flagged_nodes))


def __cli():
    """
    Command line interface.
    All options for running RetroPath 3
    """

    def get_progressive_bias(progressive_bias_strategy):
        try:
            progressive_bias_strategy = int(progressive_bias_strategy)
        except ValueError:
            progressive_bias_strategy = progressive_bias_strategy
            if progressive_bias_strategy == "None":
                progressive_bias_strategy = None
        return progressive_bias_strategy

    def define_folder_to_save(folder):
        if folder is None:
            folder_to_save = os.path.join('debugging_results', args.c_name)
        else:
            folder_to_save = folder
        if not os.path.exists(folder_to_save):
            os.makedirs(folder_to_save, exist_ok=True)
        if not os.path.exists(os.path.join(folder_to_save, 'pickles')):
            os.mkdir(os.path.join(folder_to_save, 'pickles'))
        return folder_to_save

    def get_representation(representation):
        if representation:
            representation = Test_representation
        else:
            representation = Test_to_file
        return representation

    def get_organism(biosensor, organism_name="none", complementary_sink=None, add_Hs=True):
        """
        Imports sinks.
        - detectable compounds for biosensors
        - all avaialble sinks.
        - complementary sink for media supplementation
        """
        if biosensor:
            logging.info("Using detectable compounds as sink as running in biosensor mode.")
            if add_Hs:
                organism = detectable_cmpds_H
            else:
                detectable_cmpds_noH
        else:
            if organism_name == "none":
                if complementary_sink is None:
                    logging.warning("Need to specify a sink")
                    raise InvalidSink
                else:
                    organism = import_organism_from_csv(complementary_sink, add_Hs=add_Hs)
                    logging.info("File {} is the sink".format(complementary_sink))
            elif organism_name == "test":
                if add_Hs:
                    organism = Test_organism_H
                else:
                    organism = Test_organism_noH
            elif organism_name == "ecoli":
                if add_Hs:
                    organism = ecoli_chassis_H
                else:
                    organism = ecoli_chassis_noH
            elif organism_name == "core_ecoli":
                if add_Hs:
                    organism = core_ecoli_H
                else:
                    organism = core_ecoli_noH
            elif organism_name == "ijo1366":
                if add_Hs:
                    organism = iJO1366_chassis_H
                else:
                    organism = iJO1366_chassis_noH
            elif organism_name == "bsubtilis":
                if add_Hs:
                    organism = bsubtilis_H
                else:
                    organism = bsubtilis_noH
            else:
                logging.warning("This organism is not implemented yet: {}".format(organism_name))
                raise NotImplementedError
        if complementary_sink is not None and organism_name != "none":
            cmpds_to_add = import_organism_from_csv(complementary_sink, add_Hs=add_Hs)
            organism.merge_states(cmpds_to_add)
            logging.info("Add compounds from {} to the sink".format(complementary_sink))
        return organism

    d = "All arguments to run a nice MCTS"
    parser = argparse.ArgumentParser(description=d)
    # Logs and saving information
    parser.add_argument(
        "--verbose",
        help="Default logger is INFO, switch to DEBUG is specified",
        dest="verbose",
        action="store_true",
        default=False,
    )
    parser.add_argument(
        "--log_file",
        help="Default logger is stderr, switch to log_file if specified",
        default=None,
    )
    parser.add_argument(
        "--folder_to_save",
        help="Folder to store results. Default: temp",
        default="temp",
    )
    parser.add_argument(
        "--heavy_saving",
        help="If True, will save the tree each max_iteration/10",
        default=False,
        type=lambda x: (str(x).lower() == "true"),
    )
    parser.add_argument(
        "--stop_at_first_result",
        help="If True, will stop the first time it encounters a fully solved pathway",
        default=False,
        type=lambda x: (str(x).lower() == "true"),
    )
    # Compound information
    parser.add_argument(
        "--c_name", help="Compound name. Defaults to None (InchiKey)", default=None
    )
    # One of the next 2 arguments has to be specified to give a target structure.
    parser.add_argument("--c_smiles", help="Compound smiles", default=None)
    parser.add_argument("--c_inchi", help="Compound inchi", default=None)
    # Timeouts on rdkit processes
    parser.add_argument(
        "--fire_timeout",
        help="Time allowed for firing one rule on one substrate. If -1, no timeout. Default: %(default)s",
        metavar="SECONDS",
        type=int,
        default=1,
    )
    parser.add_argument(
        "--standardisation_timeout",
        help="Time allowed for standardising results from one application. If -1, no timeout. Default: %(default)s",
        metavar="SECONDS",
        type=int,
        default=5,
    )
    # Complementary sink: if we need to supplement the media or use another sink than
    # the provided ones.
    parser.add_argument(
        "--organism_name",
        default="ecoli",
        choices=["none", "test", "ecoli", "core_ecoli", "bsubtilis", "ijo1366"],
    )
    parser.add_argument(
        "--complementary_sink",
        help="address of a csv file containing compounds to add",
        default=None,
    )
    # Visualisation
    parser.add_argument(
        "--representation",
        help="If activated, uses colors for representation. Otherwise, gedit compatible",
        action="store_true",
        default=False,
    )
    #  General MCTS parameters
    parser.add_argument(
        "--itermax", help="Maximum number of tree iterations", default=1000, type=int
    )
    parser.add_argument(
        "--parallel",
        help="Using rollout parallelisation. Default is False. Should not be used at the moment.",
        type=lambda x: (str(x).lower() == "true"),
        default=False,
    )
    parser.add_argument(
        "--expansion_width", help="Maximum number of children", default=5, type=int
    )
    parser.add_argument("--time_budget", help="Time budget", default=300, type=int)
    parser.add_argument(
        "--max_depth", help="Maximum depth of search", default=7, type=int
    )
    parser.add_argument(
        "--minimal_visit_counts",
        default=1,
        type=int,
        help="Minimal number of times a node has to be rolled out before his brothers can be expanded",
    )
    # UCT parameters
    parser.add_argument(
        "--UCT_policy",
        help="UCT policy for the tree search.",
        choices=[
            "Classical_UCT",
            "Biological_UCT_1",
            "Classical_UCT_with_bias",
            "Classical_UCT_RAVE",
            "Biochemical_UCT_with_progressive_bias",
            "Biochemical_UCT_1_with_RAVE",
            "Biochemical_UCT_1",
            "Nature_UCT",
            "Chemical_UCT_1",
            "Biological_UCT_1",
            "Biochemical_UCT_with_toxicity",
        ],
        default="Classical_UCT",
    )
    parser.add_argument(
        "--UCTK", help="UCTK for exploration/exploitation", type=float, default=2
    )
    parser.add_argument(
        "--bias_k", help="bias_k for exploration/exploitation", type=float, default=1
    )
    parser.add_argument(
        "--k_rave",
        help="k_rave for weighting of RAVE/MCTS. Number of visits before MCTS/RAVE = 50%%",
        type=float,
        default=50,
    )
    parser.add_argument(
        "--use_RAVE",
        help="Use rave or not",
        type=lambda x: (str(x).lower() == "true"),
        default=False,
    )
    # Rewarding
    parser.add_argument(
        "--penalty", help="penalty for fully unsolved state", type=int, default=-1
    )
    parser.add_argument(
        "--full_state_reward",
        help="full_state_reward for fully solved state",
        type=int,
        default=2,
    )
    parser.add_argument(
        "--pathway_scoring",
        help="pathway scoring function",
        choices=[
            "RandomPathwayScorer",
            "constant_pathway_scoring",
            "null_pathway_scoring",
            "biological_pathway_scoring",
        ],
        default="constant_pathway_scoring",
    )
    # Rollout parameters
    parser.add_argument(
        "--Rollout_policy",
        help="Rollout_policy for the tree search.",
        choices=[
            "Rollout_policy_chemical_best",
            "Rollout_policy_random_uniform_on_biochemical_multiplication_score",
            "Rollout_policy_biological_best",
            "Rollout_policy_biochemical_addition_best",
            "Rollout_policy_biochemical_multiplication_best",
            "Rollout_policy_random_uniform",
            "Rollout_policy_random_uniform_on_chem_score",
            "Rollout_policy_random_uniform_on_bio_score",
            "Rollout_policy_random_uniform_on_biochemical_addition_score",
            "Rollout_policy_first",
        ],
        default="Rollout_policy_first",
    )
    parser.add_argument("--max_rollout", help="Max rollout number", type=int, default=3)

    # Chemical and biological scoring
    parser.add_argument(
        "--chemical_scoring",
        help="Chemical scoring policy.",
        choices=[
            "RandomChemicalScorer",
            "SubstrateChemicalScorer",
            "SubandprodChemicalScorer",
            "ConstantChemicalScorer",
        ],
        default="SubandprodChemicalScorer",
    )
    parser.add_argument("--biological_score_cut_off", default=0.1, type=float)
    parser.add_argument("--substrate_only_score_cut_off", default=0.3, type=float)
    parser.add_argument("--chemical_score_cut_off", default=0.3, type=float)
    # Bias parameters
    parser.add_argument("--virtual_visits", help="Virtual visits", type=int, default=0)
    parser.add_argument(
        "--progressive_bias_strategy",
        help="Progressive bias strategy",
        default="max_reward",
    )
    parser.add_argument(
        "--progressive_widening",
        help="progressive_widening",
        type=lambda x: (str(x).lower() == "true"),
        default=False,
    )
    parser.add_argument(
        "--diameter", nargs="+", help="Diameters to consider", default=[16], type=int
    )
    parser.add_argument(
        "--EC_filter",
        nargs="+",
        help="EC numbers to consider for rules",
        default=None,
        type=str,
    )
    parser.add_argument(
        "--small",
        help="Use only a small subset",
        type=lambda x: (str(x).lower() == "true"),
        default=False,
    )
    parser.add_argument("--seed", help="Seed", type=int, default=None)
    # Load from a previously run tree
    parser.add_argument(
        "--tree_to_complete", help="Tree to restart the search from", default=None
    )
    parser.add_argument(
        "--folder_tree_to_complete",
        help="Tree to restart the search from",
        default=None,
    )

    args = parser.parse_args()
    # Config folder where to save data
    folder_to_save = define_folder_to_save(args.folder_to_save)

    if args.verbose:
        logging_level = logging.DEBUG
        log_format = "%(asctime)s - %(pathname)s (%(lineno)d) > %(funcName)s - %(levelname)s - %(message)s"
    else:
        logging_level = logging.INFO
        log_format = "%(asctime)s -- %(levelname)s -- %(message)s"
    if args.log_file is None:
        logging.basicConfig(
            stream=sys.stderr,
            level=logging_level,
            datefmt='%d/%m/%Y %H:%M:%S',
            format=log_format,
        )
    else:
        if "log" not in args.log_file:
            log_file = "log_" + args.log_file
        else:
            log_file = args.log_file
        log_writer = open("{}/{}".format(folder_to_save, log_file), "w")
        logging.basicConfig(
            stream=log_writer,
            level=logging_level,
            datefmt='%d/%m/%Y %H:%M:%S',
            format=log_format,
        )
    if args.seed is not None:
        random.seed(args.seed)
        logging.warning("Setting the seed at {}".format(args.seed))

    # Information from argparse
    representation = get_representation(args.representation)
    rules, biological_scoring = get_rules_and_score(
        full_rules_forward_H=full_rules_forward_H,
        full_rules_retro_H=full_rules_retro_H,
        full_rules_forward_no_H=full_rules_forward_no_H,
        full_rules_retro_no_H=full_rules_retro_no_H,
        add_Hs=add_Hs,
        retro=retrosynthesis,
        diameters=args.diameter,
        small=args.small,
        c_name=args.c_name,
        filtering_EC=args.EC_filter,
    )
    if args.EC_filter is not None:
        logging.info("Filtering rules based on EC number. Currently {} rules".format(len(rules.keys())))

    progressive_bias_strategy = get_progressive_bias(args.progressive_bias_strategy)
    # Rewarding:
    rollout_rewards = RolloutRewards(penalty=args.penalty, full_state_reward=args.full_state_reward)
    # Stop mode
    if args.stop_at_first_result:
        stop_mode_config = "Stopping at first result"
    else:
        stop_mode_config = "Full search - no stopping at first result"
    # Minimal visits
    minimal_visit_counts_config = f"Setting the minimal visit count for a node at {args.minimal_visit_counts}"
    # RAVE_config
    RAVE_config = f"Using RAVE: {args.use_RAVE}"
    # Scoring configuration
    chemical_scoring = args.chemical_scoring
    chemical_scoring_configuration = {
        "biological_score_cut_off": args.biological_score_cut_off,
        "substrate_only_score_cut_off": args.substrate_only_score_cut_off,
        "chemical_score_cut_off": args.chemical_score_cut_off}

    biological_score_config = f"Using biological cut off at {args.biological_score_cut_off}"
    substrate_only_score_config = f"Using substrate only cut off at {args.substrate_only_score_cut_off}"
    chemical_score_config = f"Using chemical score cut off at {args.chemical_score_cut_off}"

    # Setting chemistry info
    logging.info("Stating global parameters from configuration file")
    logging.info(tree_mode_config)
    logging.info(stop_mode_config)
    logging.info(DB_config)
    logging.info(cache_config)
    logging.info("-------------------Scoring configurations -------------------------")
    logging.info(biological_score_config)
    logging.info(substrate_only_score_config)
    logging.info(chemical_score_config)
    logging.info("-------------------MCTS configurations -------------------------")
    logging.info(minimal_visit_counts_config)
    logging.info(RAVE_config)
    logging.info("-------------------Chemical configurations -------------------------")
    logging.info(hydrogen_config)

    if use_toxicity:
        logging.info("-------------------Toxicity configurations -------------------------")
        logging.info(toxicity_scorer.log_loading)
        logging.info(toxicity_scorer.log_score)
    try:
        organism = get_organism(
            biosensor,
            organism_name=args.organism_name,
            complementary_sink=args.complementary_sink,
            add_Hs=add_Hs,
        )
        if retrosynthesis and biosensor:
            raise RunModeError(retrosynthesis, biosensor)
        try:
            if args.tree_to_complete is None:
                root_compound = Compound(
                    csmiles=args.c_smiles,
                    InChI=args.c_inchi,
                    name=args.c_name,
                    max_moves=args.expansion_width,
                    stereo=False,
                    heavy_standardisation=True,
                    fire_timeout=args.fire_timeout,
                    chemical_scoring_configuration=chemical_scoring_configuration,
                    standardisation_timeout=args.standardisation_timeout,
                )
                state = ChemicalCompoundState([root_compound], representation=representation)  # state is not sanitised
                if biosensor:
                    present_in_state_detectable = organism.compound_in_state(root_compound)
                    if present_in_state_detectable:
                        logging.warning("Removed compound from the detectable set to force enzymatic detection")
                        organism.remove_cmpd_from_state(root_compound)
                else:
                    present_in_state_sink = organism.compound_in_state(root_compound)
                    if present_in_state_sink:
                        raise CompoundInSink(folder_to_save, root_compound)

                search_tree = Tree(
                    root_state=state,
                    itermax=args.itermax,
                    parallel=args.parallel,
                    available_rules=rules,
                    rewarding=rollout_rewards,
                    expansion_width=args.expansion_width,
                    time_budget=args.time_budget,
                    max_depth=args.max_depth,
                    UCT_policy=args.UCT_policy,
                    UCT_parameters={
                        "UCTK": args.UCTK,
                        "bias_k": args.bias_k,
                        "k_rave": args.k_rave,
                    },
                    Rollout_policy=args.Rollout_policy,
                    max_rollout=args.max_rollout,
                    organism=organism,
                    main_layer_tree=True,
                    main_layer_chassis=True,
                    biological_scorer=biological_scoring,
                    chemical_scorer=chemical_scoring,
                    folder_to_save=folder_to_save,
                    virtual_visits=args.virtual_visits,
                    progressive_bias_strategy=progressive_bias_strategy,
                    progressive_widening=args.progressive_widening,
                    heavy_saving=args.heavy_saving,
                    minimal_visit_counts=args.minimal_visit_counts,
                    use_RAVE=args.use_RAVE,
                    pathway_scoring=args.pathway_scoring,
                )
            else:
                search_tree = unpickle(file_name=args.tree_to_complete,
                                       type='tree',
                                       folder_address=f"{args.folder_tree_to_complete}/pickles")
                # Check compound compatibility
                current_root_state = search_tree.root_state
                try:
                    root_compound = Compound(
                        csmiles=args.c_smiles,
                        InChI=args.c_inchi,
                        name=args.c_name,
                        max_moves=args.expansion_width,
                        stereo=False,
                        heavy_standardisation=True,
                    )
                    state = ChemicalCompoundState([root_compound], representation=representation)  # state is not sanitised
                    if state != current_root_state:
                        raise IncorrectTreeLoading(
                            "New root {} is different from old root {} when loading tree".format(
                                root_compound, current_root_state.compound_list[0]
                            )
                        )
                except CompoundDefinitionException:
                    logging.warning("Use compound information from previous Tree")
                    root_compound = current_root_state
                if biosensor:
                    present_in_state_detectable = organism.compound_in_state(root_compound)
                    if present_in_state_detectable:
                        logging.warning("Removed compound from the detectable set to force enzymatic detection")
                        organism.remove_cmpd_from_state(root_compound)

                search_tree.set_heavy_saving(args.heavy_saving)
                search_tree.set_folder_to_save(folder_to_save)
                search_tree.find_full_scope(folder_to_save=folder_to_save, name="after_unpickling")
                search_tree.jsonify_full_tree(file_name="after_unpickling")

                search_tree.set_rules(rules)  # Resetting the rules to new standards and not former tree
                start_time = time.time()
                logging.info("Starting flagging for extension at {}".format(start_time))
                search_tree.flag_nodes_for_extension(
                    extension_length=args.expansion_width,
                    maximum_depth=args.max_depth,
                    chemical_scoring_configuration=chemical_scoring_configuration,
                )
                logging.info("Finished flagging for extension at {}".format(time.time() - start_time))
                search_tree.find_full_scope(folder_to_save=folder_to_save, name="after_extending_nodes")
                search_tree.jsonify_full_tree(file_name="after_extending_nodes")
            # Running the search for both trees (new or loaded)
            search_tree.run_search(args.stop_at_first_result)
        except KeyboardInterrupt:
            logging.warning("Keyboard interruption")

        search_tree.jsonify_full_tree()
        search_tree.find_full_scope(folder_to_save=folder_to_save)
        if args.heavy_saving:
            search_tree.find_all_scopes(folder_to_save=folder_to_save)
        try:
            search_tree.find_single_best_pathway(folder_to_save=folder_to_save)
        except IndexError:
            logging.info("No best pathway was found")
        nbr = search_tree.find_multiple_best_pathways(folder_to_save=folder_to_save, return_result=True)
        logging.info("{} pathways were found".format(nbr))
        logging.info("At the end of this tree search, the cache contains {}".format(len(home_made_cache.keys())))
        search_tree.save(folder_address=folder_to_save + '/pickles', file_name="end_search")
    except (RunModeError, IncorrectTreeLoading, CompoundInSink, InvalidSink) as e:
        logging.error(e.message)
        nbr = 0
        loading_recap = {"TIME_EXECUTION": 0, "STOP_REASON": e.reason, "NUMBER_ITERATION": 0}
        logging.warning("RECAP TIME_EXECUTION={}".format(loading_recap["TIME_EXECUTION"]))
        logging.warning("RECAP STOP_REASON={}".format(loading_recap["STOP_REASON"]))
        logging.warning("RECAP NUMBER_ITERATION={}".format(loading_recap["NUMBER_ITERATION"]))
    except ChemConversionError as e:
        logging.error(e)
        nbr = 0
        loading_recap = {"TIME_EXECUTION": 0, "STOP_REASON": "incorrect_input", "NUMBER_ITERATION": 0}
        logging.error("Verify the input SMILES or InChI is valid")
        logging.warning("RECAP TIME_EXECUTION={}".format(loading_recap["TIME_EXECUTION"]))
        logging.warning("RECAP STOP_REASON={}".format(loading_recap["STOP_REASON"]))
        logging.warning("RECAP NUMBER_ITERATION={}".format(loading_recap["NUMBER_ITERATION"]))

    # Exporting information on obtained results
    with open(f"{folder_to_save}/results.csv", "w") as csv_file:
        csv_writer = csv.DictWriter(csv_file, fieldnames=["parameter", "value"])
        csv_writer.writeheader()
        csv_writer.writerow({"parameter": "stop_at_first_result", "value": args.stop_at_first_result})
        csv_writer.writerow({"parameter": "c_name", "value": args.c_name})
        csv_writer.writerow({"parameter": "c_smiles", "value": args.c_smiles})
        csv_writer.writerow({"parameter": "c_inchi", "value": args.c_inchi})
        csv_writer.writerow({"parameter": "fire_timeout", "value": args.fire_timeout})
        csv_writer.writerow({"parameter": "organism_name", "value": args.organism_name})
        csv_writer.writerow({"parameter": "complementary_sink", "value": args.complementary_sink})
        csv_writer.writerow({"parameter": "itermax", "value": args.itermax})
        csv_writer.writerow({"parameter": "expansion_width", "value": args.expansion_width})
        csv_writer.writerow({"parameter": "time_budget", "value": args.time_budget})
        csv_writer.writerow({"parameter": "max_depth", "value": args.max_depth})
        csv_writer.writerow({"parameter": "minimal_visit_counts", "value": args.minimal_visit_counts})
        csv_writer.writerow({"parameter": "UCT_policy", "value": args.UCT_policy})
        csv_writer.writerow({"parameter": "UCTK", "value": args.UCTK})
        csv_writer.writerow({"parameter": "bias_k", "value": args.bias_k})
        csv_writer.writerow({"parameter": "k_rave", "value": args.k_rave})
        csv_writer.writerow({"parameter": "use_RAVE", "value": args.use_RAVE})
        csv_writer.writerow({"parameter": "penalty", "value": args.penalty})
        csv_writer.writerow({"parameter": "full_state_reward", "value": args.full_state_reward})
        csv_writer.writerow({"parameter": "Rollout_policy", "value": args.Rollout_policy})
        csv_writer.writerow({"parameter": "max_rollout", "value": args.max_rollout})
        csv_writer.writerow({"parameter": "chemical_scoring", "value": args.chemical_scoring})
        csv_writer.writerow({"parameter": "biological_score_cut_off", "value": args.biological_score_cut_off})
        csv_writer.writerow({"parameter": "substrate_only_score_cut_off", "value": args.substrate_only_score_cut_off})
        csv_writer.writerow({"parameter": "chemical_score_cut_off", "value": args.chemical_score_cut_off})
        csv_writer.writerow({"parameter": "virtual_visits", "value": args.virtual_visits})
        csv_writer.writerow({"parameter": "progressive_bias_strategy", "value": args.progressive_bias_strategy})
        csv_writer.writerow({"parameter": "progressive_widening", "value": args.progressive_widening})
        csv_writer.writerow({"parameter": "diameter", "value": args.diameter})
        csv_writer.writerow({"parameter": "EC_filter", "value": args.EC_filter})
        csv_writer.writerow({"parameter": "tree_to_complete", "value": args.tree_to_complete})
        csv_writer.writerow({"parameter": "found_pathways", "value": nbr})
        try:
            recap = search_tree.recap
        except UnboundLocalError:
            recap = loading_recap
        csv_writer.writerow({"parameter": "TIME_EXECUTION", "value": recap["TIME_EXECUTION"]})
        csv_writer.writerow({"parameter": "STOP_REASON", "value": recap["STOP_REASON"]})
        csv_writer.writerow({"parameter": "NUMBER_ITERATION", "value": recap["NUMBER_ITERATION"]})

    with open(f"{folder_to_save}/{nbr}_results", "w") as results_file:
        # Only printing out results numbers.
        pass


if __name__ == "__main__":
    # from Tree import *
    __cli()