'''
离散化 获取日实体
'''
# 输入 股票 日期 关联性字段 字段离散化阈值
# 输出 df - (日期, 日离散化事件清单)

import os
import numpy as np
import pandas as pd
from pandas import DataFrame
import os
from pathlib import Path
from A_Global import path_date, path_stock, path_basket, threshold_discrete_dict
from tool_data import *

cur_path = Path(os.path.dirname(__file__))
pd.set_option('display.max_columns', None) # 显示完整的列
# pd.set_option('display.max_rows', None) # 显示完整的行
pd.set_option('display.expand_frame_repr', False) # 设置不折叠数据
pd.set_option('display.max_colwidth', 100)

def r10d_Classifier(data: DataFrame, threshold_rate) -> str:
    value = data['r10d']

    if abs(value) < threshold_rate: return None
    if value > 0: return '+'
    if value < 0: return '-'


def r1d_Classifier(data: DataFrame, threshold_rate) -> str:
    value = data['r1d']

    if abs(value) < threshold_rate: return None
    if value > 0: return '+'
    if value < 0: return '-'


def r5d_Classifier(data: DataFrame, threshold_rate) -> str:
    value = data['r5d']

    if abs(value) < threshold_rate: return None
    if value > 0: return '+'
    if value < 0: return '-'


def r20d_Classifier(data: DataFrame, threshold_rate) -> str:
    value = data['r20d']

    if abs(value) < threshold_rate: return None
    if value > 0: return '+'
    if value < 0: return '-'


def r30d_Classifier(data: DataFrame, threshold_rate) -> str:
    value = data['r30d']

    if abs(value) < threshold_rate: return None
    if value > 0: return '+'
    if value < 0: return '-'

def r60d_Classifier(data: DataFrame, threshold_rate) -> str:
    value = data['r60d']

    if abs(value) < threshold_rate: return None
    if value > 0: return '+'
    if value < 0: return '-'

classifier_dict = {
    'r1d': r1d_Classifier,
    'r5d': r5d_Classifier,
    'r10d': r10d_Classifier,
    'r20d': r20d_Classifier,
    'r30d': r30d_Classifier,
    'r60d': r60d_Classifier,
}

def get_path_file_basket(market, threshold_discrete_dict):
    suffix_basket = ''
    for (k, v) in threshold_discrete_dict.items(): suffix_basket += (str(v) + '_')
    return "BASKET_" + market + "_" + suffix_basket + ".csv"

def get_basket(market, threshold_dict):
    ''' Prepare 'BASKET.csv' '''
    path_file_basket = path_basket / get_path_file_basket(market, threshold_dict)

    if(os.path.exists(path_file_basket)):
        basket = pd.read_csv(path_file_basket, index_col=0)
        print('LOAD', path_file_basket)
    else:
        days    = np.genfromtxt(path_date / ('TRADE_DAYs_' + market + '.txt'), dtype=str)
        stocks  = np.genfromtxt(path_stock / ('STOCKs_' + market + '.txt'), dtype=str)

        basket  = pd.DataFrame([[''] * len(threshold_dict.keys())] * len(days), index=days,
                              columns=list(threshold_dict.keys()))

        i, n_stock = 0, len(stocks)
        for stock in stocks:
            data = pd.read_csv(path_stock / ('stocks_' + market.lower()) / stock, index_col=0)
            for day in days:
                stock_day = data.loc[day]  # 1_stock_1_day


                for col_threshold in threshold_dict.items():
                    col, threshold = col_threshold[0], col_threshold[1]
                    clazz = classifier_dict[col](stock_day, threshold)
                    if clazz:
                        goods = stock + '_' + col + '_' + clazz
                        basket.loc[day, col] += ',' + goods

            i += 1
            print(i, 'of', n_stock, 'Prepare BAKSET of', stock, ' OVER>>')
        print(basket)

        basket.to_csv(path_basket / path_file_basket)
    return basket

def calc_threshold_rate(market):
    stocks, _   = get_stock_filted(market)
    days        = get_experiment_days_train(market)
    col_discrete= get_discrete_col()

    df_col_max  = DataFrame([], columns=col_discrete)
    df_col_min  = DataFrame([], columns=col_discrete)
    for stock in stocks:
        df_trade_data = get_stock_data(stock)
        df_trade_data = df_trade_data.loc[days, col_discrete]

        df_col_max.loc[stock] = df_trade_data.max()
        df_col_min.loc[stock] = df_trade_data.min()

    print(df_col_max.max())
    print(df_col_min.min())

if __name__ == '__main__':
    '''生成全时间段所有股票在所有特征上的事件集，返回DF，并保存为csv

    ../stocks/BASKET_[market]_[threshold_rate].csv
    对每个分类事件出现的次数进行统计 - 在关系周期上，即 date_train
    
    统计所有股票在所有事件列(变化率)上的最大最小值，然后取 90% 为阈值 - 考虑各自取阈值，然后离散化，用于生成关系
    对每个分类事件出现的次数进行统计 - 在关系周期上，即 date_train - 之前只是眼看
    basket 是生成 adj 的基础
    '''
    market  = 'NASDAQ'

    get_basket(market, threshold_discrete_dict)

    # calc_threshold_rate(market)