#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
创业板50指数数据获取脚本
数据来源：baostock (证券宝) - http://www.baostock.com
"""

import baostock as bs
import pandas as pd
from datetime import datetime

def fetch_cyb50_data():
    """
    获取创业板50指数历史数据
    指数代码：sz.399673
    """
    # 登录baostock
    lg = bs.login()
    if lg.error_code != '0':
        print(f"登录失败: {lg.error_msg}")
        return None
    
    print("登录成功，开始获取数据...")
    
    # 获取创业板50指数日线数据
    # 参数说明：
    # - code: 指数代码 sz.399673 (深圳.创业板50)
    # - fields: 获取的字段 (日期, 开盘价, 最高价, 最低价, 收盘价, 成交量)
    # - start_date: 开始日期
    # - end_date: 结束日期
    rs = bs.query_history_k_data_plus(
        "sz.399673",
        "date,open,high,low,close,volume",
        start_date='2017-01-01',
        end_date='2025-12-31',
        frequency='d',  # 日频
        adjustflag='3'  # 不复权 (指数不需要复权)
    )
    
    if rs.error_code != '0':
        print(f"数据获取失败: {rs.error_msg}")
        bs.logout()
        return None
    
    # 解析数据
    data_list = []
    while rs.next():
        data_list.append(rs.get_row_data())
    
    # 转换为DataFrame
    df = pd.DataFrame(data_list, columns=['date', 'open', 'high', 'low', 'close', 'volume'])
    
    # 数据类型转换
    for col in ['open', 'high', 'low', 'close', 'volume']:
        df[col] = pd.to_numeric(df[col], errors='coerce')
    
    # 保存为CSV
    df.to_csv('cyb50_baostock.csv', index=False)
    
    print(f"数据获取成功!")
    print(f"  数据区间: {df['date'].iloc[0]} ~ {df['date'].iloc[-1]}")
    print(f"  总交易日: {len(df)} 天")
    print(f"  保存文件: cyb50_baostock.csv")
    
    # 登出
    bs.logout()
    return df

if __name__ == "__main__":
    df = fetch_cyb50_data()
    
    if df is not None:
        print("\n数据预览:")
        print(df.head())
        print("...")
        print(df.tail())