from typing import List, Set, Union, Tuple, Collection

import clingo
import networkx as nx

from clingo import Control, Symbol

from vizlo.types import ASTRuleSet, FlatASTProgram, ASTProgram
from vizlo.util import log

EMERGENCY_EXIT_COUNTER = 0


def get_all_trues_from_assumption(assumptions: Collection[Tuple[Symbol, bool]]) -> Set[Symbol]:
    return set(atom for atom, is_true in assumptions if is_true)


class SolverState:
    """
    Represents a single solver state that is created during execution.
    """

    def __init__(self, model: Set, is_still_active, step: int = -1, falses=None, adds=None):
        """
        Initializer
        :param model: a Set of clingo.Symbol that are know to be true in this step.
        :param is_still_active: whether this SolverState was generated by successful solving or a constraint
        :param step: the iterative solver step it was generated
        :param falses: a Set of clingo.Symbol that are know to be false in this step.
        :param adds: a Set of clingo.Symbol that were added to the partial model in this step.
        """
        if adds is None:
            adds = set()
        if falses is None:
            falses = set()
        if isinstance(model, list):
            model = set(model)
        self.model: Set = model
        self.step: int = step
        self.falses: Set = falses
        self.is_still_active: bool = is_still_active
        self.adds: Union[Set, None] = adds

    def __repr__(self):
        return f"{self.model}"

    def __eq__(self, other):
        if isinstance(other, SolverState):
            return self.model == other.model and self.step == other.step
        return False

    def __hash__(self):
        return id(self)

    def is_still_a_candidate(self) -> bool:
        return self.step == 0 or len(self.model) > 0


class SolveWorker:
    """
    Object that is concerned with creating new SolverStates to be added to the graph.
    """

    def __init__(self, main, ctl: clingo.Control, rule: ASTRuleSet, symbols_in_heads=None):
        """
        Should contain only information concerned with one rule set.
        :param main:
        :param ctl:
        :param rule:
        :param symbols_in_heads:
        :param fixed_assumptions:
        """
        if symbols_in_heads is None:
            symbols_in_heads = set()
        self.main = main
        self.ctl = ctl
        self.rule = rule
        self.singatures_in_heads = symbols_in_heads

    @staticmethod
    def _filter_relevant_assumptions(global_assumptions: Collection[Tuple[Symbol, bool]],
                                     signatures_of_heads: Collection[Tuple[str, int]]):
        """
        Returns a partial set of the global assumptions that only contains those assumptions that the head depends on
        :param global_assumptions: a collection of Symbols and whether they are globally considered true or false
        :param signatures_of_heads: a collection of signatures (name, arity)
        :return: a partial set of the global_assumptions that matched
        """
        result = [glob for glob in global_assumptions if
                  any(glob[0].match(name, arity) for name, arity in signatures_of_heads)]
        # print(f"{global_assumptions}, {signatures_of_heads} -> {result}")
        return result

    def run(self, i, global_assumptions):
        # analytically find recursive components and add them at once
        partial_models = self.main.find_active_nodes_at_time_step(i)
        log(f"{self.rule} with {len(partial_models)} previous partial models.")
        for partial_model in partial_models:
            # print(f"Continuing with {partial_model}")
            assumptions = self._create_true_symbols_from_solver_state(partial_model)
            relevant_assumptions = self._filter_relevant_assumptions(global_assumptions, self.singatures_in_heads)
            assumptions.extend(relevant_assumptions)
            log(f"Assumptions: {assumptions}, model: {partial_model}")
            new_partial_models = self._get_new_partial_models(assumptions, self.ctl, i)
            _consolidate_new_solver_states(assumptions, new_partial_models)
            self.main.update_graph(partial_model, self.rule, new_partial_models)

    def _get_new_partial_models(self, assumptions, ctl, i):
        solver_states_to_create = []
        with ctl.solve(assumptions=assumptions, yield_=True) as handle:
            hacky_counter = 0
            for m in handle:
                model = set(m.symbols(atoms=True))
                adds = model - get_all_trues_from_assumption(assumptions)
                syms = SolverState(m.symbols(atoms=True), True, i + 1, adds=adds)
                solver_states_to_create.append(syms)
                hacky_counter += 1
            if hacky_counter == 0:
                # HACK: This means the candidate model became conflicting.
                solver_states_to_create.append(SolverState(set(), False, i + 1))
            handle.wait()
            handle.get()
        return solver_states_to_create

    def _create_true_symbols_from_solver_state(self, s):
        syms = []
        for true in s.model:
            syms.append((true, True))
        for false in s.falses:
            if not any((false.match(s[0], s[1]) for s in self.singatures_in_heads)):
                syms.append((false, False))
        return syms


def _update_falses_in_solver_states(sss):
    all_possible = set()
    for s in sss:
        all_possible.update(s.model)
    for s in sss:
        falses = all_possible - set(s.model)
        s.falses = falses


def _assert_falses_from_assumptions(syms: List[SolverState], assumpts):
    for sym in syms:
        for atom, value in assumpts:
            if not value:
                sym.falses.add(atom)


def _consolidate_new_solver_states(assumpts, solver_states_to_create):
    _update_falses_in_solver_states(solver_states_to_create)
    _assert_falses_from_assumptions(solver_states_to_create, assumpts)


def _make_new_control_and_ground(current_prg: FlatASTProgram) -> Control:
    prg_as_str = " ".join([str(rule) for rule in current_prg])
    ctl = clingo.Control(["0"])
    ctl.add("base", [], prg_as_str)
    ctl.ground([("base", [])])
    return ctl


class SolveRunner:
    """
    The main solve runner that delegates each ruleset to a worker and collects the resulting graph.
    """
    def __init__(self, program: ASTProgram,
                 symbols_in_heads_map=None):
        if symbols_in_heads_map is None:
            symbols_in_heads_map = dict()
        self.prg: ASTProgram = program
        log(f"Created AnotherOne with {len(self.prg)} rules, {symbols_in_heads_map} signatures.")

        self._g: nx.Graph = nx.DiGraph()
        self._g.add_node(INITIAL_EMPTY_SET)
        self._solvers: List[SolveWorker] = []
        current_prg: FlatASTProgram = []
        self.symbols_in_heads_map = symbols_in_heads_map

        for rule_set in self.prg:
            current_prg.extend(rule_set)
            ctl = _make_new_control_and_ground(current_prg)
            signatures_of_heads = set()
            for rule in rule_set:
                signatures_of_heads.update(symbols_in_heads_map.get(str(rule), set()))
            self._solvers.append(SolveWorker(self, ctl, rule_set, signatures_of_heads))

    def reset_graph(self):
        self._g = nx.DiGraph()
        self._g.add_node(INITIAL_EMPTY_SET)

    def find_active_nodes_at_time_step(self, step: int) -> List[SolverState]:
        """
        Returns all solver states in the solving graph at a certain solving step that are still active.
        :param step: the step number.
        :return: a list of SolverStates
        """
        nodes: List[SolverState] = []
        for node in self._g:
            if node.step == step and node.is_still_active:
                nodes.append(node)
        return nodes

    def make_graph(self, assumption_sets=None):
        result_graph = nx.DiGraph()
        if assumption_sets is None or len(assumption_sets) == 0:
            for i, s in enumerate(self._solvers):
                s.run(i, set())
            result_graph = self._g
        else:
            for assumptions in assumption_sets:
                for i, s in enumerate(self._solvers):
                    s.run(i, assumptions)
                result_graph = nx.compose(result_graph, self._g)
                self.reset_graph()
        return result_graph

    def update_graph(self, previous: SolverState, rule: str, following_solver_states: Set[SolverState]) -> None:
        """
        Adds edges from the previous state to all following. connects them with a rule edge.
        :param previous: the previous SolverState
        :param rule: the rule that extended the previous into the following
        :param following_solver_states: all SolverStates generated by applying the rule to previous
        """
        for following in following_solver_states:
            self._g.add_edge(previous, following, rule=rule)



INITIAL_EMPTY_SET = SolverState(model=set(), is_still_active=True, step=0)

from typing import List, Set, Union, Tuple, Collection

import clingo
import networkx as nx

from clingo import Control, Symbol

from vizlo.types import ASTRuleSet, FlatASTProgram, ASTProgram
from vizlo.util import log

EMERGENCY_EXIT_COUNTER = 0


def get_all_trues_from_assumption(assumptions: Collection[Tuple[Symbol, bool]]) -> Set[Symbol]:
    return set(atom for atom, is_true in assumptions if is_true)


class SolverState:
    """
    Represents a single solver state that is created during execution.
    """

    def __init__(self, model: Set, is_still_active, step: int = -1, falses=None, adds=None):
        """
        Initializer
        :param model: a Set of clingo.Symbol that are know to be true in this step.
        :param is_still_active: whether this SolverState was generated by successful solving or a constraint
        :param step: the iterative solver step it was generated
        :param falses: a Set of clingo.Symbol that are know to be false in this step.
        :param adds: a Set of clingo.Symbol that were added to the partial model in this step.
        """
        if adds is None:
            adds = set()
        if falses is None:
            falses = set()
        if isinstance(model, list):
            model = set(model)
        self.model: Set = model
        self.step: int = step
        self.falses: Set = falses
        self.is_still_active: bool = is_still_active
        self.adds: Union[Set, None] = adds

    def __repr__(self):
        return f"{self.model}"

    def __eq__(self, other):
        if isinstance(other, SolverState):
            return self.model == other.model and self.step == other.step
        return False

    def __hash__(self):
        return id(self)

    def is_still_a_candidate(self) -> bool:
        return self.step == 0 or len(self.model) > 0


class SolveWorker:
    """
    Object that is concerned with creating new SolverStates to be added to the graph.
    """

    def __init__(self, main, ctl: clingo.Control, rule: ASTRuleSet, symbols_in_heads=None):
        """
        Should contain only information concerned with one rule set.
        :param main:
        :param ctl:
        :param rule:
        :param symbols_in_heads:
        :param fixed_assumptions:
        """
        if symbols_in_heads is None:
            symbols_in_heads = set()
        self.main = main
        self.ctl = ctl
        self.rule = rule
        self.singatures_in_heads = symbols_in_heads

    @staticmethod
    def _filter_relevant_assumptions(global_assumptions: Collection[Tuple[Symbol, bool]],
                                     signatures_of_heads: Collection[Tuple[str, int]]):
        """
        Returns a partial set of the global assumptions that only contains those assumptions that the head depends on
        :param global_assumptions: a collection of Symbols and whether they are globally considered true or false
        :param signatures_of_heads: a collection of signatures (name, arity)
        :return: a partial set of the global_assumptions that matched
        """
        result = [glob for glob in global_assumptions if
                  any(glob[0].match(name, arity) for name, arity in signatures_of_heads)]
        # print(f"{global_assumptions}, {signatures_of_heads} -> {result}")
        return result

    def run(self, i, global_assumptions):
        # analytically find recursive components and add them at once
        partial_models = self.main.find_active_nodes_at_time_step(i)
        log(f"{self.rule} with {len(partial_models)} previous partial models.")
        for partial_model in partial_models:
            # print(f"Continuing with {partial_model}")
            assumptions = self._create_true_symbols_from_solver_state(partial_model)
            relevant_assumptions = self._filter_relevant_assumptions(global_assumptions, self.singatures_in_heads)
            assumptions.extend(relevant_assumptions)
            log(f"Assumptions: {assumptions}, model: {partial_model}")
            new_partial_models = self._get_new_partial_models(assumptions, self.ctl, i)
            _consolidate_new_solver_states(assumptions, new_partial_models)
            self.main.update_graph(partial_model, self.rule, new_partial_models)

    def _get_new_partial_models(self, assumptions, ctl, i):
        solver_states_to_create = []
        with ctl.solve(assumptions=assumptions, yield_=True) as handle:
            hacky_counter = 0
            for m in handle:
                model = set(m.symbols(atoms=True))
                adds = model - get_all_trues_from_assumption(assumptions)
                syms = SolverState(m.symbols(atoms=True), True, i + 1, adds=adds)
                solver_states_to_create.append(syms)
                hacky_counter += 1
            if hacky_counter == 0:
                # HACK: This means the candidate model became conflicting.
                solver_states_to_create.append(SolverState(set(), False, i + 1))
            handle.wait()
            handle.get()
        return solver_states_to_create

    def _create_true_symbols_from_solver_state(self, s):
        syms = []
        for true in s.model:
            syms.append((true, True))
        for false in s.falses:
            if not any((false.match(s[0], s[1]) for s in self.singatures_in_heads)):
                syms.append((false, False))
        return syms


def _update_falses_in_solver_states(sss):
    all_possible = set()
    for s in sss:
        all_possible.update(s.model)
    for s in sss:
        falses = all_possible - set(s.model)
        s.falses = falses


def _assert_falses_from_assumptions(syms: List[SolverState], assumpts):
    for sym in syms:
        for atom, value in assumpts:
            if not value:
                sym.falses.add(atom)


def _consolidate_new_solver_states(assumpts, solver_states_to_create):
    _update_falses_in_solver_states(solver_states_to_create)
    _assert_falses_from_assumptions(solver_states_to_create, assumpts)


def _make_new_control_and_ground(current_prg: FlatASTProgram) -> Control:
    prg_as_str = " ".join([str(rule) for rule in current_prg])
    ctl = clingo.Control(["0"])
    ctl.add("base", [], prg_as_str)
    ctl.ground([("base", [])])
    return ctl


class SolveRunner:
    """
    The main solve runner that delegates each ruleset to a worker and collects the resulting graph.
    """
    def __init__(self, program: ASTProgram,
                 symbols_in_heads_map=None):
        if symbols_in_heads_map is None:
            symbols_in_heads_map = dict()
        self.prg: ASTProgram = program
        log(f"Created AnotherOne with {len(self.prg)} rules, {symbols_in_heads_map} signatures.")

        self._g: nx.Graph = nx.DiGraph()
        self._g.add_node(INITIAL_EMPTY_SET)
        self._solvers: List[SolveWorker] = []
        current_prg: FlatASTProgram = []
        self.symbols_in_heads_map = symbols_in_heads_map

        for rule_set in self.prg:
            current_prg.extend(rule_set)
            ctl = _make_new_control_and_ground(current_prg)
            signatures_of_heads = set()
            for rule in rule_set:
                signatures_of_heads.update(symbols_in_heads_map.get(str(rule), set()))
            self._solvers.append(SolveWorker(self, ctl, rule_set, signatures_of_heads))

    def reset_graph(self):
        self._g = nx.DiGraph()
        self._g.add_node(INITIAL_EMPTY_SET)

    def find_active_nodes_at_time_step(self, step: int) -> List[SolverState]:
        """
        Returns all solver states in the solving graph at a certain solving step that are still active.
        :param step: the step number.
        :return: a list of SolverStates
        """
        nodes: List[SolverState] = []
        for node in self._g:
            if node.step == step and node.is_still_active:
                nodes.append(node)
        return nodes

    def make_graph(self, assumption_sets=None):
        result_graph = nx.DiGraph()
        if assumption_sets is None or len(assumption_sets) == 0:
            for i, s in enumerate(self._solvers):
                s.run(i, set())
            result_graph = self._g
        else:
            for assumptions in assumption_sets:
                for i, s in enumerate(self._solvers):
                    s.run(i, assumptions)
                result_graph = nx.compose(result_graph, self._g)
                self.reset_graph()
        return result_graph

    def update_graph(self, previous: SolverState, rule: str, following_solver_states: Set[SolverState]) -> None:
        """
        Adds edges from the previous state to all following. connects them with a rule edge.
        :param previous: the previous SolverState
        :param rule: the rule that extended the previous into the following
        :param following_solver_states: all SolverStates generated by applying the rule to previous
        """
        for following in following_solver_states:
            self._g.add_edge(previous, following, rule=rule)



INITIAL_EMPTY_SET = SolverState(model=set(), is_still_active=True, step=0)

Represents a single solver state that is created during execution. Initializer         :param model: a Set of clingo.Symbol that are know to be true in this step.         :param is_still_active: whether this SolverState was generated by successful solving or a constraint         :param step: the iterative solver step it was generated         :param falses: a Set of clingo.Symbol that are know to be false in this step.         :param adds: a Set of clingo.Symbol that were added to the partial model in this step. Object that is concerned with creating new SolverStates to be added to the graph. Should contain only information concerned with one rule set.         :param main:         :param ctl:         :param rule:         :param symbols_in_heads:         :param fixed_assumptions: Returns a partial set of the global assumptions that only contains those assumptions that the head depends on         :param global_assumptions: a collection of Symbols and whether they are globally considered true or false         :param signatures_of_heads: a collection of signatures (name, arity)         :return: a partial set of the global_assumptions that matched # print(f"{global_assumptions}, {signatures_of_heads} -> {result}") # analytically find recursive components and add them at once # print(f"Continuing with {partial_model}") # HACK: This means the candidate model became conflicting. The main solve runner that delegates each ruleset to a worker and collects the resulting graph. Returns all solver states in the solving graph at a certain solving step that are still active.         :param step: the step number.         :return: a list of SolverStates Adds edges from the previous state to all following. connects them with a rule edge.         :param previous: the previous SolverState         :param rule: the rule that extended the previous into the following         :param following_solver_states: all SolverStates generated by applying the rule to previous