# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: abu_env
#     language: python
#     name: abu_env
# ---

import baostock as bs
import pandas as pd

#### 登陆系统 ####
lg = bs.login()
# 显示登陆返回信息
print('login respond error_code:'+lg.error_code)
print('login respond  error_msg:'+lg.error_msg)

# +
#### 获取沪深A股历史K线数据 ####
# 详细指标参数，参见“历史行情指标参数”章节；“分钟线”参数与“日线”参数不同。
# 分钟线指标：date,time,code,open,high,low,close,volume,amount,adjustflag
# 周月线指标：date,code,open,high,low,close,volume,amount,adjustflag,turn,pctChg
rs = bs.query_history_k_data_plus("sz.300719",
    "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
    start_date='2020-08-01', end_date='2020-09-03',
    frequency="d", adjustflag="3")
print('query_history_k_data_plus respond error_code:'+rs.error_code)
print('query_history_k_data_plus respond  error_msg:'+rs.error_msg)

#### 打印结果集 ####
data_list = []
while (rs.error_code == '0') & rs.next():
    # 获取一条记录，将记录合并在一起
    data_list.append(rs.get_row_data())
result = pd.DataFrame(data_list, columns=rs.fields)

#### 结果集输出到csv文件 ####   
//result.to_csv("D:\\history_A_stock_k_data.csv", index=False)
print(result)
# -

#### 登出系统 ####
bs.logout()

# 如上所示，volume 为成交股数（非手数），amount 为成交金额（单位为元），turn 为换手率（单位为%，若为20，则换手为 20%）

# +
#### 获取沪深A股历史K线数据 ####
# 详细指标参数，参见“历史行情指标参数”章节；“分钟线”参数与“日线”参数不同。
# 分钟线指标：date,time,code,open,high,low,close,volume,amount,adjustflag
# 周月线指标：date,code,open,high,low,close,volume,amount,adjustflag,turn,pctChg
rs = bs.query_history_k_data_plus("sh.000001",
    "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
    start_date='2020-08-01', end_date='2020-09-03',
    frequency="d", adjustflag="3")
print('query_history_k_data_plus respond error_code:'+rs.error_code)
print('query_history_k_data_plus respond  error_msg:'+rs.error_msg)

#### 打印结果集 ####
data_list = []
while (rs.error_code == '0') & rs.next():
    # 获取一条记录，将记录合并在一起
    data_list.append(rs.get_row_data())
result = pd.DataFrame(data_list, columns=rs.fields)

print(result)

# -

type(data_list[0]),type(data_list[0][0])

a = float(result[-1:].close)
a+22008


def bao_query_history_k_data_plus(stock_code,start_date,end_date=None,freq="d",field="date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST"):
    if (stock_code[:2] != 'sh' and stock_code[:2] != 'sz') or (not (stock_code[2:]).isdigit()):
        logging.error("bad stock_code",stock_code)
        return None

    # 复权类型(只针对股票)：None未复权 qfq前复权 hfq后复权 , 默认None
    adj = "2"

    code = stock_code[:2] + '.' + stock_code[2:]

    rs = bs.query_history_k_data_plus(code,field,start_date=start_date, end_date=end_date,frequency=freq, adjustflag=adj)

    if rs.error_code != '0' :
        print(rs.error_code)
        return None

    data_list = []
    while rs.next():
        data_list.append(rs.get_row_data())
        
    if freq in "dwm":
        dates = [pd.to_datetime(info_list[0]) for info_list in data_list]
    else:
        dates = [pd.to_datetime(info_list[0][:-5],format="%Y%m%d%H%M") for info_list in data_list]
        
    result = pd.DataFrame(data_list, index=dates,columns=rs.fields)
    
    if 'open' in rs.fields:
        result.open = [float(f) for f in result.open]
    if 'close' in rs.fields:
        result.close = [float(f) for f in result.close]
    if 'high' in rs.fields:
        result.high = [float(f) for f in result.high]
    if 'low' in rs.fields:
        result.low = [float(f) for f in result.low]
    if 'volume' in rs.fields:
        result.volume = [float(f) for f in result.volume]
    if 'amount' in rs.fields:
        result.amount = [float(f) for f in result.amount]
        
    if 'date' in rs.fields:
        result = result.drop(columns=['date'])
    if 'time' in rs.fields:
        result = result.drop(columns=['time'])
        
    return result


def bao_query_history_k_data_plus2(stock_code,start_date,end_date=None,freq="d",field="date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST"):
    if (stock_code[:2] != 'sh' and stock_code[:2] != 'sz') or (not (stock_code[2:]).isdigit()):
        logging.error("bad stock_code",stock_code)
        return None

    # 复权类型(只针对股票)：None未复权 qfq前复权 hfq后复权 , 默认None
    adj = "2"

    code = stock_code[:2] + '.' + stock_code[2:]

    rs = bs.query_history_k_data_plus(code,field,start_date=start_date, end_date=end_date,frequency=freq, adjustflag=adj)

    if rs.error_code != '0' :
        print(rs.error_code)
        return None

    data_list = []
    while rs.next():
        data_list.append(rs.get_row_data())
    dates = [pd.to_datetime(info_list[0]) for info_list in data_list]
    result = pd.DataFrame(data_list, index=dates,columns=rs.fields)
    
        
    return result,data_list

df = bao_query_history_k_data_plus("sh000001",'2020-08-20')
df

df = bao_query_history_k_data_plus("sz300719",'2020-08-20')
df

type(df.index),df.index

type(df.date),type(df.date[-1]),df.date[-1]

# ### 测试交易日管理器的正确性

# +
import time
import os
import os.path as path
g_cur_day = time.strftime("%Y-%m-%d", time.localtime())

# 运行该程序用户的 home 目录
root_drive = path.expanduser('~')

"""abu数据缓存主目录文件夹"""
g_project_root = path.join(root_drive, 'py_quant_data')
g_trade_day_mgr_dir = path.join(g_project_root, 'trade_day')
g_cn_trade_day_path = path.join(g_trade_day_mgr_dir, 'cn_trade_day')

g_MarketInfoManagerStartDate = "2017-08-30"
print(g_cn_trade_day_path)


# +

def getNextDate(dateFrom):
    parts = [ int(i) for i in dateFrom.split("-")]
    year = parts[0]
    month = parts[1]
    day = parts[2]
    if day < 28 :
        parts[2] = day + 1
    elif day == 28:
        if month == 2:
            if (year % 100 != 0 and year % 4 == 0) or year % 400 == 0:
                parts[2] = 29
            else:
                parts[1] = 3
                parts[2] = 1
        else:
            parts[2] = day + 1
    elif day == 29:
        if month == 2:
            parts[1] = 3
            parts[2] = 1
        else:
            parts[2] = day + 1
    elif day == 30:
        if month == 4 or month == 6 or month == 9 or month == 11:
            parts[1] = month + 1
            parts[2] = 1
        else:
            parts[2] = 31
    elif day == 31:
        if month == 12:
            parts[0] = year + 1
            parts[1] = 1
            parts[2] = 1
        else:
            parts[1] = month + 1
            parts[2] = 1
    else:
        return None
    return "-".join(str(i) if i >=10 else ("0" + str(i)) for i in parts )

def isValidDateFormat(input):
    if len(input) != 10:
        return False
    if not (input[4] == "-" and input[7] == "-"):
        return False
    try:
        temp = int(input.replace("-",""))
    except:
        return False
    return True

def getDateList(dateFrom,dateTo):
    if not (isValidDateFormat(dateFrom) and isValidDateFormat(dateTo)):
        return None
    if int(dateFrom.replace("-","")) >= int(dateTo.replace("-","")):
        return None
    ret = []
    ret.append(dateFrom)
    if dateFrom == dateTo:
        return ret
    nextDate = getNextDate(dateFrom)
    while nextDate != dateTo:
        ret.append(nextDate)
        nextDate = getNextDate(nextDate)
    ret.append(dateTo)
    return ret


# -

start_get_date = g_MarketInfoManagerStartDate
write_mode = 'a+'
if os.path.isfile(g_cn_trade_day_path):
    df = pd.read_csv(g_cn_trade_day_path, index_col=0)
    start_get_date = df.date[-1]
    start_get_date = getNextDate(start_get_date)
else:
    write_mode = 'w'
    os.system("mkdir -p " + g_trade_day_mgr_dir)
    print("mkdir -p " + g_trade_day_mgr_dir)

start_get_date_int = int(start_get_date.replace('-',''))
g_cur_day_int = int(g_cur_day.replace('-',''))
print(start_get_date_int,g_cur_day_int)
if start_get_date_int <= g_cur_day_int:
    df = bao_query_history_k_data_plus("sh000001",start_get_date)
    if df is not None and len(df) > 0:
        print("write to csv")
        header = True
        if write_mode == 'a+':
            header = False
        df.to_csv(g_cn_trade_day_path, columns=df.columns, index=True,header = header, mode=write_mode,encoding='utf-8')
    else:
        print("df is zero")
else:
    print("need not get")


df

df = pd.read_csv(g_cn_trade_day_path, index_col=0)
df.index = pd.to_datetime(df.index)


# +
def get5MinKline(stock_code,from_date):
    df = bao_query_history_k_data_plus(stock_code,from_date,freq='5',field="time,open,high,low,close,volume,amount")
    return df

def get5MinKline2(stock_code,from_date):
    df,rs = bao_query_history_k_data_plus2(stock_code,from_date,freq='5',field="date,open,high,low,close,volume,amount")
    return df,rs


# -

df.index

dt = pd.to_datetime('2017-08-30 09:35:00')
a = df.loc['2017-08-30 09:35:00']
type(a),a

a.open



StarcoderdataJupyterScriptsDedupFiltered