K = [(0,), (1,), (2,), (0,1), (0,2), (1, 2), (0, 1, 2)] 

"""
Explanation: Morse Complex
Tast of today's lab exercises is to implement an algorithm that computes a Morse Complex of a given simplicial complex K.
To make things simple the simplicial complex K is given as a list of simplices and each simplex is just a tuple of its vertices. 
Example: bellow is the definition of a complex representing a full triangle spanned on the vertices $0, 1$ and $2$.
End of explanation
"""


from itertools import combinations, chain

def simplex_closure(a):    
    """Returns the generator that iterating over all subsimplices (of all dimensions) in the closure
    of the simplex a. The simplex a is also included.
    """
    return chain.from_iterable([combinations(a, l) for l in range(1, len(a) + 1)])
        
def closure(K):
    """Add all missing subsimplices to K in order to make it a simplicial complex."""
    return list({s for a in K for s in simplex_closure(a)})

"""
Explanation: Just to make life easier we will implement a function closure, which allows us to enter only top-dimensional simplices of the complex.
End of explanation
"""


closure([(0,1,2)])

"""
Explanation: Using the above function the full triangle spanned on vertices $0, 1$ and $2$ can be given as bellow.
End of explanation
"""


from random import choice
from collections import defaultdict


def contained(a, b):
    """Returns True is a is a subsimplex of b, False otherwise."""
    return all((v in b for v in a))

def free_face(K):
    """Finds one of the free faces in the complex K.
    Returns the pair (a, b) where a is a free face of b.
    If K has no free faces, tuple (None, None) is returned.  
    """
    for s1 in K:
        cofaces = [s for s in K if len(s) == len(s1) + 1 and contained(s1, s)]
        if len(cofaces) == 1: 
            return s1, cofaces[0]
    return None, None

def top_simplex(K):
    """Returns one of the top dimensional simplices in K.  
    """
    max_length = max([len(s) for s in K])
    return choice([s for s in K if len(s) == max_length])
                
def generate_field(K):
    """Computes a discrete gradient vector field on the complex K.
    Returns a pair (V, C), where V representing pairs in the
    computed gradient vector field and C is a list of critical 
    cimplices.  
    """
    K1, C, V = set(K), [], dict()
    while K1:
        s1, s2 = free_face(K1)
        if s1:
            V[s1] = s2
            K1.remove(s1)
            K1.remove(s2)            
        else:
            s = top_simplex(K1)
            C.append(s)
            K1.remove(s)
    return V, C

"""
Explanation: The algorithm is implemented in 3 steps:
1. Compute a random discrete gradient vector field on K.
2. Cancel unnecessary critical simplices.
3. Compute the Morse boundary of each critical simplex.
Compute a random discrete gradient vector field on K
Use the procedure described in the article by Bruno Benedetti and Frank
Lutz. The procedure is implemented bellow.
End of explanation
"""


V, C = generate_field(K)
print(V, C)

"""
Explanation: Let us try the code on the complex K (full triangle).
End of explanation
"""


def boundary(a):
    """Return the boundary (codimension one faces) of the simplex a.
    """
    return combinations(a, len(a)-1)

"""
Explanation: Cancel unnecessary critical simplices
Two critical simplices $\alpha^{(p+1)}$ and $\beta^{(p)}$ are connected when there exists a V-path starting in simplices in $\partial \alpha$ and ending in $\beta$. Pairs of critical cells that are connected with exactly one gradient path can be canceled.
Implementation is tricky, is it requires searching through all V-paths connecting critical simplices. First we need some helper methods.
End of explanation
"""


from collections import defaultdict
from random import choice

def expand_path(V, ends, path):
    """Expands a path with the given prefix which is stored in the
    list path. The last simplex in the list path in the end of the 
    prefix. The path stops at the simplices in the set/list ends.
    Returns the list of all paths with the given prefix which end in
    critical cells.
    """
    if path[-1] in ends:      # path already ends at the critical simplex
        yield path
    elif path[-1] in V:       # can continue the path
        children = (s for s in boundary(V[path[-1]]) if s != path[-1])
        paths = (expand_path(V, ends, path + [V[path[-1]], c]) for c in children)
        for p in chain.from_iterable(paths):
            yield p

def find_paths_between_critical_simplices(V, C):
    """Find all paths connecting pairs of critical simplices a and b.
    (starting in the boundary of a and ending in b).
    Returns dictionary whose keys are pairs of connected critical simplices
    and value for the given key is a list of paths connecting those simplices.
    """
    paths = defaultdict(list)    
    for candidate in (s for s in C if len(s) > 1):
        for start in boundary(candidate):
            for path in expand_path(V, C, [start]):
                paths[(candidate, path[-1])].append(path)
    return paths

"""
Explanation: Our search for paths will be recursive. On each step a method is given already competed path which is extended in all possible directions. The function returns the iterator that eventualy returns all V-paths connecting pairs of critical cells.
To search for all paths connecting critical simplices we start in the boundary of each critical simplex and use the path extension function.
End of explanation
"""



def cancel_along_path(V, C, start, path):
    """Cancel two critical cells connected with the given path."""
    reversed_path = list(reversed(path)) + [start]
    for i in range(0, len(reversed_path), 2):
        V[reversed_path[i]] = reversed_path[i+1]
    C.remove(start)
    C.remove(path[-1])
        
def cancel_all(V, C):
    """Cancel all possible critical pairs."""
    while True:
        paths = find_paths_between_critical_simplices(V, C)
        candidates = [(pair, path[0]) for pair, path in paths.items() if len(path) == 1]
        if len(candidates) == 0:
            break
        pair, path = choice(candidates)
        cancel_along_path(V, C, pair[0], path)

"""
Explanation: Once all paths have been found we can cancel out critical pairs. The first function cancels out two critical simplices connected with a given path and the second one cancels out as many pairs of critical simplices as possible.
End of explanation
"""


K = closure([(0, 1, 2)])
V = dict()
C = list(K)
cancel_all(V, C)
print(V)
print(C)

"""
Explanation: Let us try the code on a simple example: full triangle where all cells are critical. Our expectation is that all simplices but one vertex will be added into V.
End of explanation
"""


def generate_simplified_field(K):
    V, C = generate_field(K)
    cancel_all(V, C)
    return V, C

"""
Explanation: For the convinience we implement the function generate_simplified_field which generates and simplifies the field on the complex.
End of explanation
"""


def morse_complex(V, C):
    boundary = defaultdict(list)
    paths = find_paths_between_critical_simplices(V, C)
    for pair in paths:
        if pair[1] in boundary[pair[0]]:
            boundary[pair[0]].remove(pair[1])
        else:
            boundary[pair[0]].append(pair[1])
    return boundary

K = closure([(0, 1, 2)])
V = dict()
C = list(K)
morse_complex(V, C)

K = closure([(0, 1), (1, 2), (0, 2)])
V, C = generate_simplified_field(K)
morse_complex(V, C)

"""
Explanation: Compute the Morse boundary
Boundary of a critical simplex $\alpha^{(p)}$ in the Morse complex consists of all such simplices $\beta^{(p-1)}$ such that the number of V-paths connecting them is odd. Computing it is easy since all necessary functions are already implemented.
End of explanation
"""


