import numpy as np
import itertools

class Node:
    def __init__(self, x, i, w, parent):
        self.x = x
        self.i = i
        self.w = w
        self.parent = parent
        self.children = dict()

    def is_root(self):
        return self.parent is None

    def simplex(self):
        s = []
        node = self
        if node.is_root():
            return None

        while True:
            s.append(node.x)
            node = node.parent
            if node.is_root():
                break
            
        s.reverse()
        return s

    def boundary_nodes(self):
        if self.is_root():
            return None
        if self.parent.is_root():
            return []

        ret = []
        up = []
        skip = self

        while not skip.is_root():
            down = list(up)
            node = skip.parent
            while len(down) != 0:
                node = node.children[down.pop()]
            ret.append(node)
            up.append(skip.x)
            skip = skip.parent

        assert(len(ret) == self.dimension() + 1)
        return ret

    def boundary_simplices(self):
        ret = []
        for node in self.boundary_nodes():
            ret.append(node.simplex())
        return ret

    def dimension(self):
        ret = 0
        if self.is_root():
            return None

        node = self
        while not node.parent.is_root():
            ret += 1
            node = node.parent
        return ret

class ComplexIterator:
    def __init__(self, start):
        self.__to_go = start.children.values()

    def next(self):
        if len(self.__to_go) == 0:
            raise StopIteration
        else:
            ret = self.__to_go.pop(0)
            self.__to_go.extend(ret.children.values())
            return ret

class Complex:
    def __init__(self):
        self.__root = Node(None, None, None, None)
        self.__count = 0
        self.__top_dim = -1
        self.__ordered = False

    def size(self):
        return self.__count

    def top_dim(self):
        return self.__top_dim

    def is_ordered(self):
        return self.__ordered

    def order(self):
        i = 0
        for node in self:
            node.i = i
            i += 1
        self.__ordered = True

    def find(self, simplex): # The simplex must be sorted!
        p = len(simplex) - 1

        if p < -1:
            return None

        node = self.__root

        for v in simplex:
            if v in node.children:
                node = node.children[v]
            else:
                return None
        return node

    def add(self, simplex, weight):
        self.__ordered = False
        simplex = sorted(simplex)

        p = len(simplex) - 1
        if p < 0:
            return None

        last = simplex[-1]
        pre = simplex[0:p]

        parent = self.find(pre)

        node = Node(last, None, weight, parent)
        
        ret = parent.children.setdefault(last, node)
        if ret == node:
            self.__count += 1
            self.__top_dim = max(self.__top_dim, p)

        return ret

    def __contains__(self, simplex):
        return self.find(simplex) is not None

    def __iter__(self):
        return ComplexIterator(self.__root)

    def __len__(self):
        return self.__count

    def num_vertices(self):
        return len(self.__root.children)

# Very naive temporary VR implementations.
def naive_vr_tmp(graph, top_dim):
    assert(graph.shape[0] == graph.shape[1])

    vertex_set = np.arange(0, graph.shape[0])
    cplx = Complex()
    for p in range(0, top_dim):
        to_add = itertools.combinations(vertex_set, p+1)
        if p == 0:
            for v in vertex_set:
                cplx.add([v], 0.0)
        elif p == 1:
            for simplex in to_add:
                if not graph.mask[simplex[0], simplex[1]]:
                    cplx.add(simplex, graph[simplex[0], simplex[1]])
        else:
            for simplex in to_add:
                codim1faces = itertools.combinations(simplex, p)
                weight = 0.0
                all_faces_in = True
                for face in codim1faces:
                    if face in cplx:
                        weight = max(cplx.find(face).w, weight)
                    else:
                        all_faces_in = False
                        break
                if all_faces_in:
                    cplx.add(simplex, weight)
            
        
    return cplx