from A_Global import path_stock, path_src_data, args, Path, path_date
import numpy as np
import pandas as pd
from WindPy import w

w.start()
w.isconnected()

path_a_stocks_before_filter = Path('D:\Code\python\paper_liuxh_21_1202\Data\dat\stocks\A_stocks_before_filter')

def download_days_in_market(market, start_date, end_date):
    (flag, df) = w.tdays(start_date, end_date, "TradingCalendar=" + market.upper(), usedf=True)
    print(df)
    return df

def download_():
    # res = w.wset("tradesuspend", "startdate=2014-01-01;enddate=2019-12-31", usedf=True)
    (flag, df) = w.wset("sectorconstituent", "date=2021-12-15;windcode=HSI.HI", usedf=True)
    print(df)
    print(len(df))

    stocks = df['wind_code'].values
    print(stocks)
    # print(df)
    np.savetxt(path_src_data / 'STOCKs_HSI_.txt', stocks, fmt='%s')
    # print(type(df))


def download_suspend():
    (flag, df) = w.wsd("0865.HK", "open", "2015-01-01", "2017-01-01", "TradingCalendar=HKEX", usedf=True)
    df.to_csv(Path('./suspend.csv'))

def download_index_daily():
    idx_name = 'IXIC.GI'
    (flag, df) = w.wsd(idx_name, "close,open", "2017-01-01", "2020-01-01", "TradingCalendar=NASDAQ", usedf=True)
    df.to_csv(path_stock / 'stocks_index' / idx_name)
    print(df)


def download_daily_trading(stocks, start_date, end_date, path):
    ''' store in path_src_data '''
    i = 1
    for stock in stocks:
        (flag, df) = w.wsd(stock,
              "open", #  ,high,low,close,turn,volume,RSI
              start_date, end_date,
              "RSI_N=6;TradingCalendar=" + market.upper(), usedf=True)
        df.to_csv(path / stock)
        print('OK', i, stock)
        i += 1

def download_stock_open(stocks, start_date, end_date, path):
    i = 1
    for stock in stocks:
        (flag, df) = w.wsd(stock,
              "open",
              start_date, end_date,
              "RSI_N=6;TradingCalendar=" + market.upper(), usedf=True)
        df.to_csv(path / stock)
        print('OK', i, stock)
        i += 1
    pass

def get_stocks_in_market(market):
    market = market.upper()

    df = pd.read_csv(path_src_data / ('STOCKs_' + market + '_ALL.csv'))
    print('%4d stocks in market %s' % (len(df), market))

    return list(df.iloc[:, 0].values)

def get_stocks_in_index_market(indic, market):
    indic  = indic.upper()
    market = market.upper()

    return list(np.genfromtxt(path_src_data / ('STOCKs_' + indic + '_' + market + '.txt'), dtype=str))

def get_stocks_in_txt(filename):
    return list(np.genfromtxt(path_src_data / filename, dtype=str))

def convert_singletxt_2_list(filename):
    file = open(path_src_data / filename)
    line = file.readline()
    file.close()
    return line.split(',')



def get_stocks_in_file(market):
    return list(np.genfromtxt(path_stock / ('STOCKs_' + market.upper() + '.txt'), dtype=str))

if __name__ == '__main__':
    ''' 下载股票交易数据 '''
    market      = 'NASDAQ'
    start_date  = '2017-01-01'
    end_date    = '2020-01-01'

    # stocks = get_stocks_in_file(market)
    # print(stocks)
    # print(len(stocks))
    # stocks = ['GRMN.O']

    stocks = np.genfromtxt(path_src_data / ('STOCKs_Index_' + market.upper() + '.txt'), dtype=str)
    download_daily_trading(stocks, start_date, end_date, path_a_stocks_before_filter / ('stocks_' + market))

    ''' 下载指数数据 '''
    # download_index_daily()

    ''' 下载指数成分股 '''
    # download_()

    ''' 下载市场交易日期 '''
    # df = download_days_in_market(market, start_date, end_date)
    # days = np.array(list(df.index))
    # print(days)
    # np.savetxt(path_date / ('TRADE_DAYs_' + market.upper()), days, fmt='%s')

    '''  '''
    # import os
    # files = os.listdir(path_stock / 'stocks_nyse')
    # print(len(files))