from vitamin_model_checker.models.costCGS.costCGS import *
from binarytree import Node
from vitamin_model_checker.logics.OATL import *
from timeit import default_timer as timer


# returns the states where the proposition holds
def get_states_prop_holds(prop):
    states = set()
    prop_matrix = cCGS.get_matrix_proposition()

    index = cCGS.get_atom_index(prop)
    if index is None:
        return None
    for state, source in enumerate(prop_matrix):
        if source[int(index)] == 1:
            states.add(state)
    return states



# set of states (ex. s1, s2) as input and returns a set of indices to identify them
def convert_state_set(state_set):
    states = set()
    for elem in state_set:
        position = cCGS.get_index_by_state_name(elem)
        states.add(int(position))
    return states


# converts a string into a set
def string_to_set(string):
    if string == 'set()':
        return set()
    set_list = string.strip("{}").split(", ")
    new_string = "{" + ", ".join(set_list) + "}"
    return eval(new_string)
    
    
def complement(state_set):
    result = set()
    graph = cCGS.get_graph()
    if (len(graph) - len(state_set) <= 0): 
        return result
    for i in range(0, len(graph)):
        if (i not in state_set):
            result.add(i)
    return result


#  function that builds a formula tree, used by the model checker
def build_tree(tpl):
    if isinstance(tpl, tuple):
        root = Node(tpl[0])
        if len(tpl) > 1:
            left_child = build_tree(tpl[1])
            if left_child is None:
                return None
            root.left = left_child
            if len(tpl) > 2:
                right_child = build_tree(tpl[2])
                if right_child is None:
                    return None
                root.right = right_child
    else:
        states = set()
        if (verify('FALSE', str(tpl))):
            return Node(str(states))
        elif (verify('TRUE', str(tpl))):
            return Node(str(set(cCGS.get_states())))
        states_proposition = get_states_prop_holds(str(tpl))
        if states_proposition is None:
            return None
        else:
            for element in states_proposition:
                states.add(cCGS.get_state_name_by_index(element))
            root = Node(str(states))
    return root



def pre(state_set):
    state_set = convert_state_set(state_set)  # returns a set of indexes
    graph = cCGS.get_graph()
    result = set()
    for state in state_set:
        for i in range(0, len(graph)):
                if graph[i][state] != 0:
                    result.add(cCGS.get_state_name_by_index(i))
    return result
    

def check_if_action_is_extension(action, extension_action):
    for index, letter in enumerate(action):
        if letter == '-':
            continue
        elif letter != extension_action[index]:
            return False
    return True
    
def get_action_extensions(action, actions_at_state):
    result = set()
    for extension_action in actions_at_state:
        if extension_action == action:
            continue
        if check_if_action_is_extension(action, extension_action):
            result.add(extension_action)
    return result        
    
def next(action, state):
    graph = cCGS.get_graph()
    result = set()
    for index, elem in enumerate(graph[state]):
        if graph[state][index] != 0:
            for elem_action in cCGS.build_list(elem):
                if action == elem_action or check_if_action_is_extension(action, elem_action):
                    result.add(index)
                    break
    return result         
        

def D(s, agents, state_set):
    agents = cCGS.get_agents_from_coalition(agents)
    tmp_set = set()
    tmp_set.add(s)
    state_index = convert_state_set(tmp_set).pop()
    state_set_index = convert_state_set(state_set)
    state_set_complement = complement(state_set_index)
    graph = cCGS.get_graph()
    result = set()
    for index, elem in enumerate(graph[state_index]):
        if graph[state_index][index] != 0:
            for action in cCGS.build_list(graph[state_index][index]):
                base_action = cCGS.get_base_action(action, agents)
                if (next(base_action, state_index) <= state_set_complement):
                    result.add(action)
    return result
    
def Cost(action_set, state):
    total = 0
    for action in action_set:
        total += cCGS.get_cost_for_action(action, state)
    return total
    
def cross(n, agents, states):
    result = set()
    for state in pre(states):
        actions = D(state, agents, states)
        if (Cost(actions, state) <= n):
            result.add(state)
    return result
    
def extract_coalition_and_cost(string):
    tmp = string[1:].split(">")
    coalition = tmp[0]
    cost = tmp[1][1:]
    return (coalition, cost)
    
    
# function that solves the formula tree. The result is the model checking result.
# It solves every node depending on the operator.
def solve_tree(node):
    if node.left is not None:
        solve_tree(node.left)
    if node.right is not None:
        solve_tree(node.right)

    if node.right is None:   # UNARY OPERATORS: not, globally, next, eventually
        if verify('NOT', node.value):  # e.g. ¬φ
            states = string_to_set(node.left.value)
            all_states = set(cCGS.get_states())
            ris = all_states - states
            node.value = str(ris)

        elif verify('COALITION_DEMONIC', node.value) and verify('GLOBALLY', node.value):  # e.g. <Jn>Gφ
            coalition_and_cost = extract_coalition_and_cost(node.value)
            coalition = coalition_and_cost[0]
            n = int(coalition_and_cost[1])
            states1 = set()
            states2 = string_to_set(node.left.value)
            p = set(cCGS.get_states())
            t = states2
            while t != p: 
                p = t
                t = states2 & (states1 | cross(n, coalition, p))
            node.value = str(p)

        elif verify('COALITION_DEMONIC', node.value) and verify('NEXT', node.value):  # e.g. <i,j><n>Xφ
            coalition_and_cost = extract_coalition_and_cost(node.value)
            coalition = coalition_and_cost[0]
            n = int(coalition_and_cost[1])
            states = string_to_set(node.left.value)
            ris = cross(n, states)
            node.value = str(ris)

        elif verify('COALITION_DEMONIC', node.value) and verify('EVENTUALLY', node.value):  # e.g. <i,j><n>Fφ
            coalition_and_cost = extract_coalition_and_cost(node.value)
            coalition = coalition_and_cost[0]
            n = int(coalition_and_cost[1])
            states1 = set(cCGS.get_states())
            states2 = string_to_set(node.left.value)
            p = set()
            t = states2
            while t != p:
                p = t
                t = states2 | (states1 & cross(n, coalition, p))
            node.value = str(p)

    if node.left is not None and node.right is not None:  # BINARY OPERATORS: or, and, until, implies
        if verify('OR', node.value): # e.g. φ || θ
            states1 = string_to_set(node.left.value)
            states2 = string_to_set(node.right.value)
            ris = states1.union(states2)
            node.value = str(ris)

        elif verify('COALITION_DEMONIC', node.value) and verify('UNTIL', node.value):  # e.g. <i,j><n>φUθ
            coalition_and_cost = extract_coalition_and_cost(node.value)
            coalition = coalition_and_cost[0]
            n = int(coalition_and_cost[1])
            states1 = string_to_set(node.left.value)
            states2 = string_to_set(node.right.value)
            p = set()
            t = states2
            while t != p: 
                p = t
                t = states2 | (states1 & cross(n, coalition, p))
            node.value = str(p)
        elif verify('COALITION_DEMONIC', node.value) and verify('RELEASE', node.value): #e.g. <i,j><n>φRθ
            coalition_and_cost = extract_coalition_and_cost(node.value)
            coalition = coalition_and_cost[0]
            n = int(coalition_and_cost[1])
            states1 = string_to_set(node.left.value)
            states2 = string_to_set(node.right.value)
            p = set(cCGS.get_states())
            t = states2
            while t != p: 
                p = t
                t = states2 & (states1 | cross(n, coalition, p))
            node.value = str(p)
        elif verify('COALITION_DEMONIC', node.value) and verify('WEAK', node.value): #e.g. <i,j><n>φWθ
            coalition_and_cost = extract_coalition_and_cost(node.value)
            coalition = coalition_and_cost[0]
            n = int(coalition_and_cost[1])
            states1 = string_to_set(node.right.value)
            states2 = string_to_set(node.left.value) | states1
            p = set(cCGS.get_states())
            t = states2
            while t != p: 
                p = t
                t = states2 & (states1 | cross(n, coalition, p))
            node.value = str(p)
        
        elif verify('AND', node.value):  # e.g. φ && θ
            states1 = string_to_set(node.left.value)
            states2 = string_to_set(node.right.value)
            ris = states1.intersection(states2)
            node.value = str(ris)

        elif verify('IMPLIES', node.value):  # e.g. φ -> θ
            # p -> q ≡ ¬p ∨ q
            states1 = string_to_set(node.left.value)
            states2 = string_to_set(node.right.value)
            not_states1 = set(cCGS.get_states()).difference(states1)
            ris = not_states1.union(states2)
            node.value = str(ris)

# returns whether the result of model checking is true or false in the initial state
def verify_initial_state(initial_state, string):
    if initial_state in string:
        return True
    return False


# does the parsing of the model, the formula, builds a tree and then it returns the result of model checking
# function called by front_end_CS
def model_checking(formula, filename):
    global cCGS

    if not formula.strip():
        result = {'res': 'Error: formula not entered', 'initial_state': ''}
        return result

    # model parsing
    cCGS = costCGS()
    cCGS.read_file(filename)

    # formula parsing
    res_parsing = do_parsing(formula, cCGS.get_number_of_agents())
    if res_parsing is None:
        result = {'res': "Syntax Error", 'initial_state': ''}
        return result
    root = build_tree(res_parsing)
    if root is None:
        result = {'res': "Syntax Error: the atom does not exist", 'initial_state': ''}
        return result

    # model checking
    solve_tree(root)

    # solution
    initial_state = cCGS.get_initial_state()
    bool_res = verify_initial_state(initial_state, root.value)
    result = {'res': 'Result: ' + str(root.value), 'initial_state': 'Initial state '+ str(initial_state) + ": " + str(bool_res)}
    return result
    
    
def model_checking_test(formula, filename):
    global cCGS

    if not formula.strip():
        result = {'res': 'Error: formula not entered', 'initial_state': ''}
        return result

    # model parsing
    cCGS = costCGS()
    cCGS.read_file(filename)

    # formula parsing
    res_parsing = do_parsing(formula, cCGS.get_number_of_agents())
    if res_parsing is None:
        result = {'res': "Syntax Error", 'initial_state': ''}
        return result
    root = build_tree(res_parsing)
    if root is None:
        result = {'res': "Syntax Error: the atom does not exist", 'initial_state': ''}
        return result

    # model checking
    time1 = timer()
    solve_tree(root)
    model_checking_time = timer() - time1
    # solution
    initial_state = cCGS.get_initial_state()
    bool_res = verify_initial_state(initial_state, root.value)
    result = {'res': 'Result: ' + str(root.value), 'initial_state': 'Initial state '+ str(initial_state) + ": " + str(bool_res), 'time': model_checking_time}
    return result


