import os
import numpy as np
import pandas as pd
import torch
from pandas import Series, DataFrame
from A_Global import path_basket, path_stock, path_date, args, path_relation
from D_basket import get_basket
from B_date_tool import gen_date
import scipy.sparse as sp
from tool_data import get_stock_filted
from tool_relation import relation_set_2_dict
from D_basket import get_path_file_basket
# 输入 df_bakset, threshold_dict_discrete, end_date
# 输出 [nstock, nstock, nrelation] - 可转换为 2 维矩阵

def freq2relations(freq_pn) -> (list, set):
    ''' 把频繁2项集转成关系（去掉频繁2项集中元素的"维度"和"分类"，只留"股票名"），不做任何处理
    :param
        freq_pn:    [['stock1_dimension_classification', 'stock2_dimension_classification'], ..]
            exp. [['600030.XSHG_r1d_-', '600056.XSHG_r1d_-'], ..]
            note. dimension must be identical inside the function
    :return:
        relations   [(stock1, stock2), ..]
            exp. [('600030.XSHG', '600056.XSHG'), ..]
        stocks:     {stock1, ..}
            exp. {'600085.XSHG', '603000.XSHG', ..}
    '''
    relations = []
    stocks = set()
    for e in freq_pn:
        relation = (e[0].split('_r')[0], e[1].split('_r')[0])
        relations.append(relation)
        stocks.add(relation[0])
        stocks.add(relation[1])
    # print('relations: ', len(relations), 'stocks:', len(stocks))
    return relations, stocks

def _count_bits(x):
    i = 0
    while x != 0:
        x = x & (x-1)
        i += 1
    return i

def pick_stocks_from_basket_cols_2(col: Series, min_support=0.5):
    ''' 一天之内同涨或同跌，异涨或异跌
    :param
        col: Series(index=date, value=[stock1_dimension_classification, ..])
            exp.    2019-07-16    [600004.XSHG_r1d_-, 600009.XSHG_r1d_-, 600030....
                    2019-07-17    [600007.XSHG_r1d_-, 600009.XSHG_r1d_+, 600015....
    :return
        adj     [nstock, nstock, 2] - 2 <= relation_same, relation_diff

        relations:  [(stock1, stock2), ..] <= relations.append(relations_same) .append(relations_diff)
            exp. [('603258.XSHG', '300526.XSHE'), ..]
        stocks      [..] <= stocks_same | stocks_diff
    :intermediate
        freq_same:  [['stock1_dimension_classification1', 'stock2_dimension_classification1'], ..]
            exp. [['600030.XSHG_r1d_-', '600056.XSHG_r1d_-'], ..]
            note. dimension must be identical inside the function
        fred_diff   [['stock1_dimension_classification1', 'stock2_dimension_classification2'], ..]

        relations_same, stocks_same
        relations_diff, stocks_diff

    '''
    ''' freq_all => freq_same, freq_diff '''
    transactions = col.values
    bound = min_support * len(transactions)

    '''统计 "股票_特征_涨跌"事件 出现的次数、时间点
    in:     transactions    [['600004.XSHG_r1d_-', ..], ..] # 每日出现的所有股票涨跌事件
    out:    goods           {"600004.XSHG_r1d_-": [2, 0b001001], ..}
    '''
    goods, flag = {}, 0b1
    for each_day in transactions: # each day
        for t in each_day:
            if t == '': continue
            if t not in goods: goods[t] = [0, 0b0]
            goods[t][0] += 1
            goods[t][1] ^= flag

        flag <<= 1

    '''频繁2项集
    in:     goods_items   [('600004.XSHG_r1d_-', [8, 11010]), ..]    
    out:    freq_2set_simple
    '''
    goods_items = list(goods.items())
    n_goods = len(goods.items())
    freq_2set_simple = {}
    for i in range(n_goods):
        for j in range(i + 1, n_goods): # 遍历所有的 goods_items 对
            conj = goods_items[i][1][1] & goods_items[j][1][1]
            freq_2set_simple[(goods_items[i][0], goods_items[j][0])] = [_count_bits(conj), conj, 0]

    '''合并频繁2项集 同号或异号 
    freq_2set_simple    {('600851.XSHG_r1d_-', '603123.XSHG_r1d_-'): [7, 11011, 0], ..}
        value: [出现次数, 出现时间二进制标记, 访问标记]
    freq_2set_merged    {('600129.XSHG_r1d_+', '600329.XSHG_r1d_+'): (18, '0b11111111101111011111'), ..}
    '''
    freq_2set_merged, freq_2set_same, freq_2set_diff = {}, [], []
    for key in freq_2set_simple:
        if(freq_2set_simple[key][2] == 1): continue
        freq_2set_simple[key][2] = 1

        num, bits = \
            freq_2set_simple[key][0], \
            freq_2set_simple[key][1]
        s_0, s_1 = \
            key[0][-1], \
            key[1][-1]
        new_key, has_new_key = '', False
        if(s_0 == '-' and s_1 == '-'):
            new_key = (key[0][:-1] + '+', key[1][:-1] + '+')
            if(new_key in freq_2set_simple): has_new_key = True

        elif(s_0 == '+' and s_1 == '+'):
            new_key = (key[0][:-1] + '-', key[1][:-1] + '-')
            if (new_key in freq_2set_simple): has_new_key = True

        elif(s_0 == '-' and s_1 == '+'):
            new_key = (key[0][:-1] + '+', key[1][:-1] + '-')
            if(new_key in freq_2set_simple): has_new_key = True

        elif(s_0 == '+' and s_1 == '-'):
            new_key = (key[0][:-1] + '-', key[1][:-1] + '+')
            if(new_key in freq_2set_simple): has_new_key = True

        if has_new_key:
            freq_2set_simple[new_key][2] = 1
            num = freq_2set_simple[key][0] + freq_2set_simple[new_key][0]
            bits = freq_2set_simple[key][1] | freq_2set_simple[new_key][1]

        if num < bound: continue

        if key in freq_2set_merged: raise Exception(key, 'has been existed in freq_2set_merged')
        freq_2set_merged[key] = (num, str(bin(bits)))

        if s_0 == s_1:
            freq_2set_same.append(key)
        if s_0 != s_1:
            freq_2set_diff.append(key)

    '''freq_2set => stock_relations
    freq_2set_same / freq_2set_diff     [('600004.XSHG_r1d_-', '600009.XSHG_r1d_-'), ..]
    relations_same / relations_diff     [('600004.XSHG', '600007.XSHG'), ..]
    '''
    relations_same, stocks_same = freq2relations(freq_2set_same)
    relations_diff, stocks_diff = freq2relations(freq_2set_diff)

    relations, stocks = [], []
    stocks.append(stocks_same)
    stocks.append(stocks_diff)
    relations.append(relations_same)
    relations.append(relations_diff)

    return relations, stocks


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

def get_path_file_adj(args):
    from tool_data import get_experiment_days_train
    market              = args['market']
    days_train          = get_experiment_days_train(market)
    start_date          = days_train[0]
    end_date            = days_train[-1]
    min_support_same    = args['min_support_same']
    min_support_diff    = args['min_support_diff']
    return 'ADJ_' + market + '_' + start_date + '_' + end_date + '_' + str(min_support_same) + '_' + str(min_support_diff) + '.npy'

def basket_2_adj(stock2idx, basket_all, start_date, end_date, min_support=0.5, min_support_diff=0.5):
    # 1. 确定生成关系的时间段 date_relation
    date_relation = list(basket_all.loc[start_date: end_date].index)
    print(' GEN RELATIONs', len(date_relation), 'days before', start_date, 'to', end_date)

    # 2. 时间段内对所有特征列生成关系 每列生成两种关系：同向、异向
    basket = basket_all.loc[date_relation]
    stocks, relation_list = set(), []
    for col in basket.columns:  # 一列对应一种关系
        col_series = basket[col].apply(lambda x: x.split(',')[1:])

        relations_col, stocks_col = \
            pick_stocks_from_basket_cols_2(col_series, min_support)
        relation_list += relations_col
        # print(relations_col[0])
        # print(relations_col[1])
        # np.savetxt('./r1d.txt', relations_col[0], fmt='%s')
        # '''
        # 画
        # '''
        # exit()
        print('%5s: s %3d r %5d - same(%3d,%5d) diff(%3d,%5d)' % (
            col,
            len(stocks_col[0] | stocks_col[1]), len(relations_col[0]) + len(relations_col[1]),
            len(stocks_col[0]), len(relations_col[0]),
            len(stocks_col[1]), len(relations_col[1])
        ))

    # 3. 关系编码成邻接矩阵 adj <= relations + stocks
    adj = gen_adjs_from_relations(relation_list, stock2idx)
    # np.save(path_relation / ('ADJ_' + start_date + '_' + end_date + '_' + str(min_support) + '_' + str(min_support_diff) + '.npy'), adj)
    np.save(path_relation / get_path_file_adj(args), adj)
    return relation_list, adj


def calc_adj(args):
    ''' return more info and called by gen_adj
    calculate adj of 3d arrays, and return some intermediate data

    :param args: market, min_support_same/diff
    :return: relation_list, relation_set, relation_dict, adj
    '''
    market              = args['market']
    min_support_same    = args['min_support_same']
    min_support_diff    = args['min_support_diff']

    date_all            = gen_date(args)
    date_train          = date_all['date_train']
    train_start_date    = date_train[0]
    train_end_date      = date_train[-1]

    idx2stock, stock2idx= get_stock_filted(market)
    basket_all          = pd.read_csv(path_basket / get_path_file_basket(market, args['threshold_discrete_dict']), index_col=0, keep_default_na=False)
    relation_list, adj  = basket_2_adj(stock2idx, basket_all, train_start_date, train_end_date, min_support_same, min_support_diff)

    relation_set        = set()
    for relations in relation_list:
        relation_set    = relation_set.union(set(relations))

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

def gen_adj(args):
    ''' only return adj & load ADJ.npy or call calc_adj()
    generate adj of 3d arrays from the whole train date period

    :param args: see 'calc_adj(args)'
    :return: adj
    '''
    path_file_adj = path_relation / get_path_file_adj(args)
    if os.path.exists(path_file_adj):
        print('Load RELATIONs from', path_file_adj)
        adj = np.load(path_file_adj)
    else:
        print('Make RELATIONs')
        relation_list, relation_set, relation_dict, adj = calc_adj(args)

    adj = torch.Tensor(adj)
    return adj



if __name__ == '__main__':
    args['market'] = 'NASDAQ'

    relation_list, relation_set, relation_dict, adj = calc_adj(args)
    print(len(relation_set))

    res = gen_adj(args)
    print(sum(adj))
    print(sum(res))
    print((adj == res.numpy().tolist()).all())

    '''
    一步步捋生成关系的过程，把它写在论文上，可画图表示流程
    '''