# -*-coding:UTF-8 -*
from LTL import *
from GBA import *


def equals(formules1: list[LTL], formules2: list[LTL]):
	difference = formules1.copy()
	for form1 in formules2:
		equals = False
		for form2 in formules1:
			if form1.equals(form2):
				difference.remove(form2)
				equals = True
		if not equals:
			return False
	return not difference


def findall(val: list[tuple[list[LTL], int]], formule: list[LTL]):
	result = set()
	for couple in val:
		if equals(formule,couple[0]):
			result.add(couple[1])
	return result


def clean(formules: Iterable[LTL]):
	new = list[LTL]()
	for form1 in formules:
		toAdd = True
		for form2 in new:
			if form1.equals(form2):
				toAdd = False
		if toAdd:
			new.append(form1)
	formules = new
	return new


def copy(formules: Iterable[LTL]):
	result = list[LTL]()
	for formule in formules:
		result.append(formule.copy())
	return result


def is_in(formule, formules):
	for formule2 in formules:
		if formule.equals(formule2):
			return True
	return False


def duplicate(state: int, structure: GBA, todo: dict[int, list[LTL]], val: list[tuple[list[LTL], int]], toExpand: list[int]):
	statecp=structure.add_state()
	if state in structure.initial_states:
		structure.add_initial(statecp)
	update_val(state, statecp, val)
	todo[statecp]=copy(todo[state])
	todo[state]=copy(todo[state])
	toExpand.append(statecp)
	return statecp


def update_val(state, statecp, val):
	for couple in val:
		if couple[1]==state:
			val.append((couple[0],statecp))
	return val


def traitET(formule: LTL, i: int, structure: GBA, todo: dict[int, list[LTL]], state: int, val: list[tuple[list[LTL], int]], toExpand: list[int]):
	"""
	Applies a 'AND' node of the LTL formula to the val.
	"""
	child1 = formule.children[0]
	child2 = formule.children[1]
	formulecp = copy(todo[state])
	todo[state] = formulecp
	formulecp[i] = child1
	formulecp.append(child2)
	val.append((formulecp,state))
	return True


def traitG(formule: LTL, i: int, structure: GBA, todo: dict[int, list[LTL]], state: int, val: list[tuple[list[LTL], int]], toExpand: list[int]):
	child = formule.children[0]
	formulecp = todo[state].copy()
	todo[state] = formulecp
	formulecp.append(child)
	formulecp[i] = LTL("X", [formule])
	val.append((formulecp,state))
	return True


def traitOU(formule: LTL, i: int, structure: GBA, todo: dict[int, list[LTL]], state: int, val: list[tuple[list[LTL], int]], toExpand: list[int]):
	#on duplique l'état considéré pour avoir deux choix de transitions possibles
	state2 = duplicate(state, structure, todo, val, toExpand)
	#on récupère les formules à droite et à gauche du OU
	child1 = formule.children[0]
	child2 = formule.children[1]
	#on modifie les formules courantes pour avoir dans un état child1 et dans l'autre child2
	todo[state][i] =  child1.copy()
	todo[state2][i] =  child2.copy()
	#on ajoute le fait que state et state2 satisfont le nouvel ensemble de formules
	val.append((todo[state], state))
	val.append((todo[state2],state2))
	return True


def traitF(formule: LTL, i: int, structure: GBA, todo: dict[int, list[LTL]], state: int, val: list[tuple[list[LTL], int]], toExpand: list[int]):
	#on duplique l'état considéré pour avoir deux choix de transitions possibles
	state2 = duplicate(state, structure, todo, val, toExpand)
	#on récupère la formule sous le F
	child = formule.children[0]
	#on construit la formule avec XF
	nextF = LTL("X", [formule])
	#on modifie les formules courantes pour avoir dans un état child et dans l'autre XF(child)
	todo[state2][i] =  child
	todo[state][i] = nextF
	#on ajoute le fait que state et state2 satisfont le nouvel ensemble de formules
	val.append((todo[state], state))
	val.append((todo[state2],state2))
	return True


def traitU(formule: LTL, i: int, structure: GBA, todo: dict[int, list[LTL]], state: int, val: list[tuple[list[LTL], int]], toExpand: list[int]):
	#on duplique l'état considéré pour avoir deux choix de transitions possibles
	state2 = duplicate(state, structure, todo, val, toExpand)
	#on récupère les formules à droite et à gauche du U
	child1 = formule.children[0]
	child2 = formule.children[1]
	#on construit la formule avec X..U..
	nextU = LTL("X", [formule])
	#on modifie les formules courantes

	#à compléter
	todo[state2][i] = child2
	todo[state][i]  = nextU
	todo[state].append(child1.copy())

	#on ajoute le fait que state et state2 satisfont le nouvel ensemble de formules
	val.append((todo[state], state))
	val.append((todo[state2],state2))
	return True


def traitement(done: bool, i: int, structure: GBA, todo: dict[int, list[LTL]], state: int, val: list[tuple[list[LTL], int]], toExpand: list[int]):
	"""
	Applies one step of the todo collection of LTL formulas.
	"""
	formule = todo[state][i]
	if formule.root == "et":
		return traitET(formule, i, structure, todo, state, val, toExpand)
	elif formule.root == "ou":
		return traitOU(formule, i, structure, todo, state, val, toExpand)
	elif formule.root == "F":
		return traitF(formule, i, structure, todo, state, val, toExpand)
	elif formule.root == "G":
		return traitG(formule, i, structure, todo, state, val, toExpand)
	elif formule.root == "U":
		return traitU(formule, i, structure, todo, state, val, toExpand)
	else:
		return False


def buildBuchi(formule: LTL, propositions: list[str]):
	"""
	Builds a buchi state machine.

	From the LTL formula and the possible propositions.
	"""
	structure = GBA([0],[0],propositions,[],[])
	val=[([formule], 0)]
	todo = {0:[formule]}
	expandState(structure, todo, 0, val, [0])
	accept_conditions(structure, todo, val)
	return structure


def expandState(structure: GBA, todo: dict[int, list[LTL]], state: int, val: list[tuple[list[LTL], int]], toExpand: list[int]):
	"""
	Creates all states from a collection of LTL transitions formulas (in todo).
	"""
	done = False
	i = 0
	while not done and i < len(todo[state]):
		done = traitement(done, i, structure, todo, state, val, toExpand)
		i=i+1
	if not done:
		toExpand = toExpand[1:]
	if toExpand:
		expandState(structure, todo, toExpand[0], val, toExpand)
	else:
		computeSucc(structure, todo, val)
	return val


def computeSucc(structure: GBA, todo: dict[int, list[LTL]], val: list[tuple[list[LTL], int]]):
	states = structure.all_states.copy()
	for state in states:
		succ = findall(val, next(todo[state]))
		if not succ:
			new = structure.add_state()
			val.append((next(todo[state]),new))
			todo[new]=next(todo[state])
			expandState(structure, todo, new, val, [new])
			succ = findall(val, next(todo[state]))
		label = now(todo[state])
		if is_consistent(label):
			for s in succ:
				structure.add_edge(state,label, s)


def accept_conditions(structure: GBA, todo: dict[int, list[LTL]], val: list[tuple[list[LTL], int]]):
	cond = list[LTL]()
	for couple in val:
		for formule in couple[0]:
			if formule.root in {"F","U"} and not is_in(formule,cond):
				cond.append(formule)
	sets = list[list[int]]()
	for i in range(len(cond)):
		formule = cond[i]
		sets.append([])
		for state in structure.all_states:
			if not is_in(LTL("X", [formule]), todo[state]):
				sets[i].append(state)
	structure.change_accept(sets)
	return structure


def next(formules: Iterable[LTL]):
	retour = list[LTL]()
	for formule in formules:
		if formule.root == "X":
			retour.append(formule.children[0])
	return retour


def now(formules: Iterable[LTL]):
	retour = list[str]()
	for formule in formules:
		if formule.root == "non":
			retour.append("~"+formule.children[0].root)
		if not formule.children:
			assert type(formule.root) is str
			retour.append(formule.root)
	return retour


if __name__ == "__main__":
	formule = new_LTL("G(a).Fb")
	trad = buildBuchi(formule, ["a","b"])
	trad.export("trad")
	print(trad.accepting_states)

	formule = new_LTL("aUb")
	trad = buildBuchi(formule, ["a","b"])
	trad.export("trad_union")
	print(trad.accepting_states)
