# Extended persistence
w = -1
cone = [Simplex([w] + [v for v in s.vertices]) for s in elephant_complex]
cone.append(Simplex([w]))

def projection(points, axis = 1):
    def value(v):
        return points[v][axis]

    return value
value = projection(elephant_points, 1)

def ep_compare(values):
    def max_vertex(s):
        return max(values(v) for v in s.vertices if v != w)

    def min_vertex(s):
        return min(values(v) for v in s.vertices if v != w)

    def compare(s1, s2):
        if s1.dimension() == 0 and w in s1.vertices:
            return -1
        if s2.dimension() == 0 and w in s2.vertices:
            return 1

        if s1.dimension() != s2.dimension():
            return cmp(s1.dimension(), s2.dimension())

        if (w in s1.vertices) ^ (w in s2.vertices):       # only one cone simplex
            return 1 if w in s1.vertices else -1
        elif w in s1.vertices:
            return -cmp(min_vertex(s1), min_vertex(s2))
        else:
            return cmp(max_vertex(s1), max_vertex(s2))

    return compare

f = Filtration(elephant_complex + cone)
f.sort(ep_compare(value))

persistence = StaticPersistence(f)
persistence.pair_simplices()

def eval_ep(values):
    def eval(s):
        if s.dimension() == 0 and w in s.vertices:
            return float('inf')
        if w in s.vertices:
            return min(values(v) for v in s.vertices if v != w)
        else:
            return max(values(v) for v in s.vertices)
    return eval

dgms = init_diagrams(persistence, f, eval_ep(value), lambda i: i)

print show_diagram(dgms)
#print dgms
#print len(dgms)