import random

from socialgraph.graph import Graph

class SocialGraph(Graph):
    def add_nodes(self, nodes):
        '''添加节点'''
        self.add_vertices_from(nodes)

    def gen_conn(self, type='null'):
        '''产生社会关系图
        type 为关系类型, 支持null star, weak_tie, mesh'''
        if type == 'null':
            return self.null_conn()
        elif type == 'star':
            return self.star_conn()
        elif type == 'weak_tie':
            return self.weak_tie_conn()
        elif type == 'mesh':
            return self.mesh_conn()
        else:
            print("[ERROR] gen_conn unsupport social type:", type)

    def null_conn(self):
        '''无连接关系'''
        pass

    def star_conn(self, center_v=None):
        '''星形连接
        center_v 指定星形中心的值，如果不指定则随机产生
        返回center的值'''
        nodes = self.vertices()

        if center_v and not center_v in nodes:
            print('[ERROR] star_conn SocialGraph does not have node: ', center_v)
            return
        count = len(nodes)
        center_v = center_v or nodes[random.randint(0, (count - 1))]
        print('start_conn center value: %s' % center_v)
        nodes.remove(center_v)

        for n in nodes:
            self.add_edge((n, center_v))
        return center_v

    def __make_tie_groups(self):
        nodes = self.vertices()
        random.shuffle(nodes)  # 随机进行分组
        count = len(nodes)

        center_pos = int(count/2)
        center_v = nodes[center_pos]
        left_part = nodes[:center_pos+1]  # 左边的group,含center
        right_part = nodes[center_pos:]  # 右边的group,含center
        return [left_part, center_v, right_part]

    def weak_tie_conn(self, left_group=None, right_group=None):
        '''weak tie连接
        left_group 左边的group成员
        right_group 右边的group成员
        如果未指定，则随机分成两个group
        返回center的值'''
        groups = []
        if left_group and right_group:
            centers = [v for v in left_group if v in right_group]
            if not centers:
                print("weak_tie_conn groups not have center")
                return
            groups = [left_group, centers[0], right_group]
        else:
            groups = self.__make_tie_groups()

        [left_part, center_v, right_part] = groups
        print("weak_tie center value: ", center_v)
        print("weak_tie left group: ", left_part)
        print("weak_tie right group: ", right_part)

        # 每个group是一个mesh连接
        self.__make_mesh_conn(left_part)
        self.__make_mesh_conn(right_part)
        return groups

    def __make_mesh_conn(self, nodes):
        for m in nodes:
            for n in nodes:
                self.add_edge((m, n))

    def mesh_conn(self):
        '''网状连接
        所有节点全部连通'''
        nodes = self.vertices()
        self.__make_mesh_conn(nodes)

    def is_direct_conn(self, n1, n2):
        '''测试两个节点是否直接连通'''
        graph = self.graph_dict
        return (n1 == n2) or (n2 in graph[n1]) or (n1 in graph[n2])

if __name__ == "__main__":
    # 测试代码
    g2 = SocialGraph()
    g2.add_nodes([1, 2, 3, 4, 5, 6, 7])
    assert g2.is_direct_conn(1, 3) == False
    assert g2.is_direct_conn(1, 1) == True
    assert g2.is_direct_conn(7, 7) == True

    g2.star_conn(4)
    assert g2.is_direct_conn(6, 1) == False
    assert g2.is_direct_conn(6, 4) == True
    assert g2.is_direct_conn(4, 2) == True
    assert g2.is_direct_conn(2, 2) == True

    g2 = SocialGraph()
    g2.add_nodes([1, 2, 3, 4, 5, 6, 7])
    g2.mesh_conn()
    for i in range(1, 8):
        for j in range(1, 8):
            assert g2.is_direct_conn(i, j) == True

    g2 = SocialGraph()
    g2.add_nodes([1, 2, 3, 4, 5, 6, 7])
    g2.weak_tie_conn(left_group=[1,2,3,4], right_group=[4,5,6,7])
    assert g2.is_direct_conn(1, 5) == False
    assert g2.is_direct_conn(3, 5) == False
    assert g2.is_direct_conn(1, 7) == False
    assert g2.is_direct_conn(1, 6) == False
    assert g2.is_direct_conn(1, 4) == True
    assert g2.is_direct_conn(5, 4) == True