import numpy as np
from pathlib import Path
from A_Global import path_relation
from D_relation import get_path_file_adj
from tool_data import get_stock_filted

''' 获取每个关系类别对应的股票对，可用于 relation_pic 画实验一的图 '''
if __name__ == '__main__':
    market          = 'NASDAQ'
    path_file_adj   = path_relation / 'ADJ_NASDAQ_2018-01-02_2018-12-31_0.6_0.6.npy'
    adj = np.load(path_file_adj)

    idx2stock, stock2idx = get_stock_filted(market)
    print(adj.shape)


    for i_relation in range(8):
        res = []
        for i_stock_a in range(len(idx2stock)):
            for i_stock_b in range(len(idx2stock)):
                if adj[i_stock_a, i_stock_b, i_relation] == 1:
                    res.append((idx2stock[i_stock_a], idx2stock[i_stock_b]))
        np.savetxt(path_relation / (market + '_' + str(i_relation) + '.txt'), res, fmt='%s')