import numpy as np
import pandas as pd
from A_Global import path_stock, path_date, threshold_discrete_dict



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 get_all_days(market):
    return np.genfromtxt(path_date / ('TRADE_DAYs_' + market + '.txt'), dtype=str)

def get_experiment_days(market):
    return np.genfromtxt(path_date / ('EXPERIMENT_DAYs_' + market + '.txt'), dtype=str)

def get_experiment_days_train(market):
    return np.genfromtxt(path_date / ('EXPERIMENT_DAYs_TRAIN_' + market + '.txt'), dtype=str)

def get_stock_data(stock):
    if(stock[-1] == 'N'):   market = 'NYSE'
    elif(stock[-1] == 'O'): market = 'NASDAQ'
    else: raise Exception(stock, 'cannot find')

    data = pd.read_csv(path_stock / ('stocks_' + market.lower()) / stock, index_col=0)
    return data

def get_basket_df(market):

    pass

def get_discrete_col():
    return list(threshold_discrete_dict.keys())

if __name__ == '__main__':

    # data = get_discrete_col()
    # print(data)

    idx2stock, stock2idx = get_stock_filted('NASDAQ')
    res = ''
    for s in idx2stock:
        res += (', ' + s)
    print(res)