import os
import sys
import time
import pandas as pd
from struct import unpack

# 通达信原始数据地址
origin_data_path1 = 'C:/qh/tdx/vipdoc/sz/lday'     # sz
origin_data_path2 = 'C:/qh/tdx/vipdoc/sh/lday'     # sh
#  注：处理后的数据，存放在此脚本文件所在目录下data中
target = os.path.dirname(os.path.abspath(sys.argv[0])) + os.sep + 'data/day'   # 保存csv文件的目录 根据需要修改 'data/day'



#         ---------(以下为代码，不需要修改)------------
# 将通达信的日线文件转换成CSV格式
def day2csv(source_dir, file_name, target_dir):
    # 以二进制方式打开源文件
    source_file = open(source_dir + os.sep + file_name, 'rb')
    buf = source_file.read()
    source_file.close()

    # 打开目标文件，后缀名为CSV
    target_file = open(target_dir + os.sep + file_name[: file_name.rindex('.')] + '.csv', 'w')
    buf_size = len(buf)
    rec_count = int(buf_size / 32)
    begin = 0
    end = 32
    header = str('date') + ',' + str('open') + ',' + str('high') + ',' + str('low') + ',' \
             + str('close') + ',' + str('amount') + ',' + str('volume') + '\n'
    target_file.write(header)
    for i in range(rec_count):
        # 将字节流转换成Python数据格式
        # I: unsigned int
        # f: float
        a = unpack('IIIIIfII', buf[begin:end])
        # 处理date数据
        year = a[0] // 10000
        month = (a[0] % 10000) // 100
        day = (a[0] % 10000) % 100
        date = '{}-{:02d}-{:02d}'.format(year, month, day)

        line = date + ',' + str(a[1] / 100.0) + ',' + str(a[2] / 100.0) + ',' \
               + str(a[3] / 100.0) + ',' + str(a[4] / 100.0) + ',' + str(a[5]) + ',' \
               + str(a[6]) + '\n'
        target_file.write(line)
        begin += 32
        end += 32
    target_file.close()

def transform_data():
    if not os.path.exists(target):
        os.makedirs(target)
    code_list = []
    source_list = [origin_data_path1, origin_data_path2]
    for source in source_list:
        file_list = os.listdir(source)
        # 逐个文件进行解析
        for f in file_list:
            if (f>='sh600000')&(f<='sh605999')|(f>='sz300000')&(f<='sz300999')\
                    |(f>='sz000001')&(f<='sz003999'):         # 筛选股票提抽范围设置
                day2csv(source, f, target)
        # 获取所有股票/指数代码
                code_list.append(f[:f.rindex('.')])
                # code_list.extend(list(map(lambda x: x[:x.rindex('.')], f)))            # 数据转换的股票代码
    # 保存所有代码列表
    stock_list=pd.DataFrame(code_list, columns=['code'])
    stock_list.to_csv('./data/stock_list.csv')

if __name__ == '__main__':
    # 程序开始时的时间
    time_start = time.time()
    # 转换所有数据
    transform_data()
    # 程序结束时系统时间
    time_end = time.time()
    print('程序所耗时间：', time_end - time_start)