'''
两个过滤指标：有停牌 | 小股价
'''
import numpy as np
import pandas as pd
from pathlib import Path
from A_Global import path_stock, path_date
path_src_data = path_stock / 'A_stocks_before_filter'

# 加载股票集，遍历
# 获取股价，检查停牌和价格情况
# 分类打印：全量 n_all + 停牌 n_suspend + 小股价 n_low_price + 合格 n_ok
# 输入：股票集 np.array, 实验日期范围, 股票日开盘价文件
# 输出：合格股票集

threshold_price = 15

def get_stock_filted(market):
    idx2stock = np.genfromtxt(path_stock / ('STOCKs_' + market + '.txt'), dtype=str)
    stock2idx = {j: i for i, j in enumerate(idx2stock)}
    return idx2stock, stock2idx

def filter_stock(stocks, days, path_open_price_file):
    n_all, n_suspend, n_low, n_ok = 0, 0, 0, 0
    stocks_ok = []

    for stock in stocks:
        n_all += 1
        try:
            df = pd.read_csv(path_open_price_file / stock, index_col=0)
        except FileNotFoundError as e:
            print(n_suspend, e)
            n_suspend += 1
            continue

        try:
            df = df.loc[days]
        except KeyError:
            print('OUT', stock, 'for nan---')
            n_suspend += 1
            continue

        # if len(df[df['OPEN'].isnull().T.any()]) != 0:
        if df['OPEN'].isnull().any():
            print('OUT', stock, 'for nan')
            n_suspend += 1
            continue

        if min(df['OPEN']) < threshold_price:
            print('OUT', stock, 'for <', threshold_price)
            n_low += 1
            continue

        # print('OK ', stock)
        stocks_ok.append(stock)
        n_ok += 1

    print(n_all, 'n_suspend:', n_suspend, 'n_low:', n_low, 'n_ok:', n_ok)
    return stocks_ok


def call_filter_stock(args):
    market          = args['market']

    stocks          = np.genfromtxt(path_src_data / ('STOCKs_Index_' + market.upper() + '.txt'), dtype=str)
    print('market ', market, '- stocks ', len(stocks))

    days_experiment = np.genfromtxt(path_date / ('TRADE_DAYs_' + market.upper() + '.txt'), dtype=str)

    stocks_ok       = filter_stock(stocks, days_experiment, path_src_data / ('stocks_' + market.lower()))
    return stocks_ok

if __name__ == '__main__':
    market = 'NYSE'
    threshold_price = 15
    args = {
        'market'    : market,
    }

    stocks = call_filter_stock(args)
    print(stocks)

    ''' SAVE '''
    # np.savetxt(path_stock / ('STOCKs_' + market + '.txt'), stocks, fmt='%s')

    ''' CHECK '''
    idx2stocks, stock2idx = get_stock_filted(market)
    print(idx2stocks == stocks)


