import numpy as np
import json
from A_Global import path_relation_tgc
from tool_data import get_stock_filted
from tool_pic import pic_heatmap
from tool_relation import relation_set_2_dict

'''
建立 2 维矩阵 [nstock, nstock]

遍历关系 json，对于每一个行业中的股票
选出 stock_filted 中的股票
在 2 维矩阵中两两建立关系

画热力图

'''

def list_2_combination(stock_list):
    comb_list = []
    for i in range(len(stock_list)):
        for j in range(i + 1, len(stock_list)):
            comb_list.append((stock_list[i], stock_list[j]))
    return comb_list

def gen_adj_tgc(market):
    stocks, stock2idx   = get_stock_filted(market)
    adj                 = np.zeros((len(stocks), len(stocks)))
    with open(path_relation_tgc / (market + '_industry_ticker.json'), 'r', encoding='utf8') as fp:
        json_data       = json.load(fp)

    relation_set = set()
    for (k, v_list) in json_data.items():
        stocks_need = []
        for stock in v_list:
            if market == 'NASDAQ': stock = stock[:-2] + '.O'
            if (stock in stocks):
                stocks_need.append(stock)

        comb_list = list_2_combination(stocks_need)
        # relation_list = relation_list + comb_list
        relation_set = relation_set.union(set(comb_list))
        for comb in comb_list:
            a = stock2idx[comb[0]]
            b = stock2idx[comb[1]]
            adj[a][b] = 1
            adj[b][a] = 1

    relation_dict = relation_set_2_dict(relation_set)
    return relation_set, relation_dict, adj

if __name__ == '__main__':
    market = 'NYSE'

    (relation_set, relation_dict, adj) = gen_adj_tgc(market)
    print(sum(adj))
    print(sum(sum(adj)))
    print(len(relation_set))
    exit()
    pic_heatmap(adj)

    '''
    检查 relation_dict 的正确性
    对自己生成的 adj 生成 relation_dict, adj - 使用取交集的方法来获取共同关系
    优化关系生成逻辑 - 使用长期关系和短期关系
    '''