import time
from dsl.base_dsl import Node
from evaluation import Evaluation
from game import Game
from os.path import join
from players.rule_of_28_player import Rule_of_28_Player
from search.dsl_bus import VarList, VarScalar, VarScalarFromArray, NumericConstant, StringConstant
from typing import List, Tuple, Optional, Dict, Set, Union, Any, Generator


class ProgramList:
    """Represent all programs that are generated by depth following Bottom-Up Search steps."""

    def __init__(self) -> None:
        """
        - plist store info about programs expanded (number of programs expanded, generated by which
        DSL rule and at which depth).
        - number_programs accounts for the amount of programs added to self.plist.
        """
        self.plist: Dict[int, Dict[str, List[Node]]] = {}
        self.number_programs: int = 0

    def insert(self, program: Node) -> None:
        """Add a program to the dictionary based on its rule and size."""

        if program.get_size() not in self.plist:
            self.plist[program.get_size()] = {}

        if program.class_name() not in self.plist[program.get_size()]:
            self.plist[program.get_size()][program.class_name()] = []

        self.plist[program.get_size()][program.class_name()].append(program)
        self.number_programs += 1

    def get_programs(self, size: int) -> Dict[str, List[Node]]:
        """Get programs of size "size"."""

        if size in self.plist:
            return self.plist[size]
        return {}


class BottomUpSearch:
    """Implementation of BUS applied to program synthesis represented as an AST."""

    def __init__(self, log_file: str, program_file: str,
                 log_results: Optional[bool] = True) -> None:
        """
        - log_results represents if logging will happen or not.
        - log_folder is the folder where self.log_file will be written.
        - program_folder is the folder where self.program_file will be written.
        - log_file is the name of the log file.
        - program_file is the name of file where the program will be written as text.
        - state_action_pairs stores the states and the actions taken in a match.
        - detect_equivalence checks or not if it should ignore new programs that produce the same
          output of previously seen programs.
        - closed_list is a list of already expanded/seen programs during search.
        - programs_outputs stores the output of programs on the data used for checking equivalence.
        - plist is a ProgramList used to store info about programs expanded (number of programs
          expanded, generated by which DSL rule and at which depth).
        """
        self.log_results = log_results

        if self.log_results:
            self.log_folder = 'logs/'
            self.program_folder = 'programs/'

            self.log_file = 'bus-' + log_file
            self.program_file = 'bus-' + program_file

        self.state_action_pairs: List[Tuple[Game, Tuple[Union[str, int]]]] = []
        self.detect_equivalence: bool = True
        self.closed_list: Set[str] = set()
        self.programs_outputs: Set[str] = set()
        self.plist: ProgramList = ProgramList()

    def generate_observational_data(self) -> None:
        """Generate data used to check if a program is equivalent (i.e.: produce the same output
        given the same input) as other programs already expanded during search.
        """
        player = Rule_of_28_Player()
        evaluation = Evaluation()

        self.state_action_pairs = []

        for _ in range(1):
            evaluation.play_match(player, player, save_state_action_pairs=True)
            self.state_action_pairs += evaluation.state_action_pairs

    def generate_initial_set_of_programs(self, numeric_constant_values: List[int],
                                         string_constant_values: List[str],
                                         variables_scalar: List[str], variables_list: List[str],
                                         variables_scalar_from_array: List[str],
                                         functions_scalars: List[Node]) -> List[Node]:
        """Create the initial set of programs (this is DSL specific) for BUS to work on."""

        set_of_initial_programs = []

        for i in variables_scalar:
            p = VarScalar.new(i)
            if self.detect_equivalence and self.has_equivalent(p):
                continue
            set_of_initial_programs.append(p)

        for i in variables_scalar_from_array:
            p = VarScalarFromArray.new(i)
            if self.detect_equivalence and self.has_equivalent(p):
                continue
            set_of_initial_programs.append(p)

        for i in variables_list:
            p = VarList.new(i)
            if self.detect_equivalence and self.has_equivalent(p):
                continue
            set_of_initial_programs.append(p)

        for i in numeric_constant_values:
            constant = NumericConstant.new(i)
            if self.detect_equivalence and self.has_equivalent(constant):
                continue
            set_of_initial_programs.append(constant)

        for i in string_constant_values:
            constant = StringConstant.new(i)
            if self.detect_equivalence and self.has_equivalent(constant):
                continue
            set_of_initial_programs.append(constant)

        for i in functions_scalars:
            p = i.new()
            if self.detect_equivalence and self.has_equivalent(p):
                continue
            set_of_initial_programs.append(p)

        return set_of_initial_programs

    @staticmethod
    def init_env(state: Game) -> Dict[str, Any]:
        """Store relevant information about the current state."""

        actions = state.available_moves()
        progress_value = [0, 0, 7, 7, 3, 2, 2, 1, 2, 2, 3, 7, 7]
        move_value = [0, 0, 7, 0, 2, 0, 4, 3, 4, 0, 2, 0, 7]
        marker = 6

        env = {'state': state, 'progress_value': progress_value, 'actions': actions,
               'marker': marker, 'move_value': move_value,
               'neutrals': [col[0] for col in state.neutral_positions]}

        return env

    def run_program_on_equivalence_data(self, p: Node) -> Tuple[Optional[Any], bool]:
        """Generate p's output based on the data generated in self.generate_observational_data."""

        outputs = []
        for pair in self.state_action_pairs:
            env = self.init_env(pair[0])
            try:
                out = p.interpret(env)

                if isinstance(out, list):
                    if isinstance(out[0], list):
                        out = tuple(tuple(li) for li in out)
                    else:
                        out = tuple(out)

                outputs.append(out)
            except Exception as e:
                template = "An exception of type {0} occurred. Arguments:\n{1!r}"
                template.format(type(e).__name__, e.args)
                return None, True
        return tuple(outputs), False

    def has_equivalent(self, p: Node) -> bool:
        """Check if program p is equivalent to previously expanded programs."""

        outputs_p, error = self.run_program_on_equivalence_data(p)

        if not error:
            if outputs_p in self.programs_outputs:
                return True
            else:
                self.programs_outputs.add(outputs_p)
                return False
        return False

    def grow(self, operations: List[Node], size: int) -> Generator:
        """Grow BUS tree respecting the DSL's expansion rules."""

        new_programs = []
        for op in operations:
            for p in op.grow(self.plist, size):
                if p.to_string() not in self.closed_list:
                    if self.detect_equivalence:
                        if self.has_equivalent(p):
                            self.closed_list.add(p.to_string())
                            continue
                    self.closed_list.add(p.to_string())
                    new_programs.append(p)
                    yield p

        for p in new_programs:
            self.plist.insert(p)

    def search(self, bound: int, operations: List[Node], numeric_constant_values: List[int],
               string_constant_values: List[str], variables_scalar: List[str],
               variables_list: List[str], variables_scalar_from_array: List[str],
               functions_scalars: List[Node], eval_function: Evaluation, use_triage: bool,
               time_limit: float, detect_equivalence: Optional[bool] = False,
               collect_library: Optional[bool] = False) \
            -> Union[Dict[int, Dict[str, List[Node]]], Tuple[float, Node, int, int]]:
        """Main routine of BUS.
        - bound is the max size of a program generated.
        - operations stores all operations allowed by the DSL.
        - numeric_constant_values stores all integers allowed by the DSL.
        - string_constant_values stores all strings allowed by the DSL.
        - variables_list stores two list names that give information about the state of a game.
          'neutrals' stores the positions of the neutral pins and 'actions' stores the available
          actions.
        - variables_scalar store a list name 'marker' that give information about the position of
          the neutral markers. (Defunct attribute, not used - pass an empty list to search).
        - variables_scalar_from_array stores two list names that give information about the weights
          of each column according to the type of the current action. Refer to the article
          "A Generalized Heuristic for Can't Stop".
        - functions_scalars stores methods that extract specific information about the current
          state of a game.
        - eval_function store an evaluation function used to calculate a score of a program.
        - use_triage checks if it will use triage for evaluation.
        - time_limit stores for how long BUS will run.
        - detect_equivalence checks if BUS should check for programs that produce the same output.
        - collect_library check if BUS should store programs while it's running.
        """

        time_start = time.time()

        self.detect_equivalence = detect_equivalence

        NumericConstant.accepted_types = [set(numeric_constant_values)]
        StringConstant.accepted_types = [set(string_constant_values)]
        VarList.accepted_types = [set(variables_list)]
        VarScalar.accepted_types = [set(variables_scalar)]
        VarScalarFromArray.accepted_types = [set(variables_scalar_from_array)]

        Node.filter_production_rules(operations, numeric_constant_values, string_constant_values,
                                     variables_scalar, variables_list, variables_scalar_from_array,
                                     functions_scalars)

        self.closed_list = set()
        self.programs_outputs = set()

        if self.detect_equivalence:
            self.generate_observational_data()

        init_set_of_programs = self.generate_initial_set_of_programs(numeric_constant_values,
                                                                     string_constant_values,
                                                                     variables_scalar,
                                                                     variables_list,
                                                                     variables_scalar_from_array,
                                                                     functions_scalars)
        self.plist = ProgramList()
        for p in init_set_of_programs:
            self.plist.insert(p)

        number_programs_evaluated = 0
        number_games_played = 0
        current_size = 0
        id_log = 1

        best_score = 0.0
        best_program = None

        while current_size <= bound:

            number_evaluations_bound = 0

            for p in self.grow(operations, current_size):
                time_end = time.time()
                if time_end - time_start > time_limit - 60:
                    if self.log_results:
                        with open(join(self.log_folder + self.log_file), 'a') as results_file:
                            str_format = "{:d}, {:f}, {:d}, {:f} \n"
                            results_file.write((str_format.format(id_log, best_score,
                                                                  number_games_played,
                                                                  time_end - time_start)))
                    return best_score, best_program, number_programs_evaluated, number_games_played

                number_programs_evaluated += 1
                number_evaluations_bound += 1

                if collect_library:
                    score = 0
                else:
                    if use_triage:
                        score, _, number_matches_played = eval_function.eval_triage(p, best_score)
                    else:
                        score, _, number_matches_played = eval_function.eval(p)
                    number_games_played += number_matches_played

                if best_program is None or score > best_score:
                    best_score = score
                    best_program = p

                    print('Best Program: ')
                    print(best_program.to_string())
                    print('Best Score: ', best_score)
                    print()

                    if self.log_results:
                        with open(join(self.log_folder + self.log_file), 'a') as results_file:
                            str_format = "{:d}, {:f}, {:d}, {:f} \n"
                            results_file.write((str_format.format(id_log, best_score,
                                                                  number_games_played,
                                                                  number_programs_evaluated,
                                                                  time.time() - time_start)))

                        with open(join(self.program_folder + self.program_file),
                                  'a') as results_file:
                            results_file.write(("{:d} \n".format(id_log)))
                            results_file.write(best_program.to_string())
                            results_file.write('\n')

                        id_log += 1
            current_size += 1

        if collect_library:
            return self.plist.plist

        return best_score, best_program, number_programs_evaluated, number_games_played
