import numpy as np
import scipy as sp

def calc_relation_intersec_rate(rsetbase, rsetcompare):
    rsetbase    = set(rsetbase)
    rsetcompare = set(rsetcompare)

    return round((len(rsetbase.intersection(rsetcompare))) / len(rsetbase), 3)

def relation_set_2_dict(relation_list):
    relation_dict = {}
    # print(relation_list)
    # print(len(relation_list))
    # exit()
    for pair in relation_list:
        # print(pair)
        # exit()
        if(pair[0] not in relation_dict): relation_dict[pair[0]] = []
        relation_dict[pair[0]].append(pair[1])
        if(pair[1] not in relation_dict): relation_dict[pair[1]] = []
        relation_dict[pair[1]].append(pair[0])
    return relation_dict


def relation_list_2_adj(relation_list):

    pass

def relation_diff_to_adj(nstock, relations_pn):

    if len(relations_pn) == 0:
        return np.zeros((nstock, nstock))

    adj = sp.coo_matrix((np.ones(relations_pn.shape[0]), (relations_pn[:, 0], relations_pn[:, 1])),
                        shape=(nstock, nstock),
                        dtype=np.float32)
    adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)
    # adj = _normalize(adj + sp.eye(adj.shape[0]))
    return adj.toarray()

def relation_same_to_adj(nstock, relations_pn):
    '''
    :param
        nstock:          nnode
        relations_pn:    edges
    :return:
        adj
    '''
    vertex_num = nstock

    if len(relations_pn) == 0:
        return np.eye(nstock)

    adj = sp.coo_matrix((np.ones(relations_pn.shape[0]), (relations_pn[:, 0], relations_pn[:, 1])),
                        shape=(vertex_num, vertex_num),
                        dtype=np.float32)
    adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)
    adj = adj + sp.eye(adj.shape[0])
    # adj = _normalize(adj + sp.eye(adj.shape[0]))
    return adj.toarray()

def gen_adjs_from_relations(relations, stock2idx):
    '''
    :param relations: List
        [[(stock_1, stock_2), ..], [], // 正关系，负关系
        ...]
    :param stock2idx: Dict
        {'stock_1': 0, ...}
    :return:
    '''
    adj = []
    for i in range(len(relations)):
        # stock => number
        # 偶数为同向关系；奇数为异向关系
        r_i = relations[i]
        r_i = np.array(list(map(lambda x: [stock2idx[x[0]], stock2idx[x[1]]], r_i)))
        if(i % 2 == 0):
            adj_i = relation_same_to_adj(len(stock2idx), r_i)
        else:
            adj_i = relation_diff_to_adj(len(stock2idx), r_i)
        adj.append(adj_i)
        # if i == 0:
        #     print(adj_i)
        #     exit()

    # (nrelation, nstock, nstock) → (nstock, nstock, nrelation)
    adj = np.array(adj)

    adj = adj.swapaxes(0, 1)
    adj = adj.swapaxes(1, 2)
    return adj