import os
import time
import pandas as pd
import tushare as ts
from datetime import datetime
from multiprocessing import Pool

pd.set_option('max_rows', None)
pd.set_option('max_columns', None)
pd.set_option('expand_frame_repr', False)
pd.set_option('display.unicode.ambiguous_as_wide', True)
pd.set_option('display.unicode.east_asian_width', True)

# 输入参数
start_date = '20230101'  # 数据获取开始日期
end_date = '20230615'  # 数据获取结束日期
adj = ''  # 复权类型：None未复权 qfq前复权 hfq后复权

# 输入your token,初始化pro接口
ts.set_token('7c5fd48b6b3b16df637935ed8f740438b3e57e72b02e52fc468f02d9')
pro = ts.pro_api()

# 查询历史所有上市交易的股票列表
data1 = pro.query('stock_basic', exchange='', list_status='L', fields='ts_code,name') # 上市
time.sleep(1)
data2 = pro.query('stock_basic', exchange='', list_status='D', fields='ts_code,name') # 退市
time.sleep(1)
data3 = pro.query('stock_basic', exchange='', list_status='P', fields='ts_code,name') # 暂停上市
time.sleep(1)
data = pd.concat([data1, data2, data3], ignore_index=True)
code_list = data[['ts_code', 'name']].values


def create_path(ts_code):
    date_str = str(pd.to_datetime(start_date).date())  # 日期转换成字符串
    path = os.path.join("..", "all_stock_candle", "stock", date_str)
    # 保存数据
    if not os.path.exists(path):
        # os.mkdir(path)  # 可以建一级文件夹
        os.makedirs(path)  # 可以建多级文件夹
    file_name = ts_code + ".csv"
    return os.path.join(path, file_name)


# 获取所有股票的历史数据
def do_load(ts_code, name, adj, start_date, end_date):
    print(ts_code)
    for i in range(5):
        try:
            df = ts.pro_bar(ts_code=ts_code, adj=adj, start_date=start_date, end_date=end_date)
            if df.empty:
                continue
            df['name']=name
            df=df[['ts_code','name','trade_date','open','high', 'low','close','pre_close','change', 'pct_chg','vol','amount']]
            df['ts_code'] = df["ts_code"].astype(str)
            df.loc[df["ts_code"].str.startswith('6'), 'ts_code'] = "sh" + df["ts_code"].str[:6]
            df.loc[df["ts_code"].str.startswith('4') | df['ts_code'].str.startswith('8'), 'ts_code'] = "bj" + df["ts_code"].str[:6]
            df.loc[df["ts_code"].str.startswith('3') | df['ts_code'].str.startswith('0'), 'ts_code'] = "sz" + df["ts_code"].str[:6]
            # df["股票名称"] = pd.to_numeric(df["股票名称"], errors="coerce")
            df['trade_date']=pd.to_datetime(df['trade_date'],format='%Y%m%d')
            df.sort_values(by=['trade_date'], ascending=True, inplace=True)
            df.reset_index(drop=True, inplace=True)
            path = create_path(ts_code)
            df.to_csv(path, index=False, mode='w', encoding='gbk')
            break
        except Exception as e:
            print(e)


if __name__ == '__main__':
    # 多进程获取股票数据
    start_time = datetime.now()
    if 1 == 0:
        pool = Pool(10)
        pool.starmap(do_load, [(code_list[i][0], code_list[i][1],adj, start_date, end_date) for i in
                               range(len(code_list))])
        pool.close()
        pool.join()
        # 计算获取所有数据需要的时间
        print(datetime.now() - start_time)
    else:
        for i in range(len(code_list)):
            ak_code = code_list[i][0]
            name = code_list[i][1]
            do_load(ak_code, name,adj,start_date, end_date)
        # 计算获取所有数据需要的时间
        print("获取数据时间：", datetime.now() - start_time)