import akshare as ak
import pandas as pd
from datetime import datetime, timedelta
import os

def get_cyb50_60min_data():
    """获取创业板50指数最近3个月的60分钟K线数据"""
    try:
        print("正在获取创业板50指数(399673)的60分钟K线数据...")
        
        # 获取创业板50指数的60分钟K线数据
        # 使用东方财富接口获取分钟级数据
        data = ak.index_zh_a_hist_min_em(symbol="399673", period="60")
        
        if data.empty:
            print("未获取到数据，请检查网络连接或接口是否可用")
            return None
        
        # 筛选最近3个月的数据
        three_months_ago = datetime.now() - timedelta(days=90)
        data['时间'] = pd.to_datetime(data['时间'])
        data_filtered = data[data['时间'] >= three_months_ago].copy()
        
        print(f"成功获取数据，共 {len(data_filtered)} 条记录")
        print(f"数据时间范围: {data_filtered['时间'].min()} 至 {data_filtered['时间'].max()}")
        print(f"数据列: {list(data_filtered.columns)}")
        
        return data_filtered
        
    except Exception as e:
        print(f"获取数据时出错: {str(e)}")
        return None

def export_to_csv(data, filename=None):
    """将数据导出到CSV文件"""
    if data is None or data.empty:
        print("没有数据可导出")
        return False
    
    try:
        # 生成文件名
        if filename is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"cyb50_60min_kline_{timestamp}.csv"
        
        # 导出到CSV
        data.to_csv(filename, index=False, encoding='utf-8-sig')
        print(f"数据已成功导出到: {filename}")
        print(f"文件大小: {os.path.getsize(filename) / 1024:.2f} KB")
        
        return True
        
    except Exception as e:
        print(f"导出CSV时出错: {str(e)}")
        return False

def get_cyb50_30min_data():
    """获取创业板50指数最近3个月的30分钟K线数据"""
    try:
        print("正在获取创业板50指数(399673)的30分钟K线数据...")
        
        # 获取创业板50指数的30分钟K线数据
        data = ak.index_zh_a_hist_min_em(symbol="399673", period="30")
        
        if data.empty:
            print("未获取到数据，请检查网络连接或接口是否可用")
            return None
        
        # 筛选最近3个月的数据
        three_months_ago = datetime.now() - timedelta(days=90)
        data['时间'] = pd.to_datetime(data['时间'])
        data_filtered = data[data['时间'] >= three_months_ago].copy()
        
        print(f"成功获取数据，共 {len(data_filtered)} 条记录")
        print(f"数据时间范围: {data_filtered['时间'].min()} 至 {data_filtered['时间'].max()}")
        print(f"数据列: {list(data_filtered.columns)}")
        
        return data_filtered
        
    except Exception as e:
        print(f"获取数据时出错: {str(e)}")
        return None

def get_cyb50_15min_data():
    """获取创业板50指数最近1个月的15分钟K线数据"""
    try:
        print("正在获取创业板50指数(399673)的15分钟K线数据...")
        
        # 获取创业板50指数的15分钟K线数据
        data = ak.index_zh_a_hist_min_em(symbol="399673", period="15")
        
        if data.empty:
            print("未获取到数据，请检查网络连接或接口是否可用")
            return None
        
        # 筛选最近1个月的数据
        one_month_ago = datetime.now() - timedelta(days=30)
        data['时间'] = pd.to_datetime(data['时间'])
        data_filtered = data[data['时间'] >= one_month_ago].copy()
        
        print(f"成功获取数据，共 {len(data_filtered)} 条记录")
        print(f"数据时间范围: {data_filtered['时间'].min()} 至 {data_filtered['时间'].max()}")
        print(f"数据列: {list(data_filtered.columns)}")
        
        return data_filtered
        
    except Exception as e:
        print(f"获取数据时出错: {str(e)}")
        return None

def main():
    """主函数"""
    print("=" * 50)
    print("创业板50指数K线数据获取工具")
    print("=" * 50)
    
    # 获取60分钟数据
    print("\n【60分钟K线数据】")
    data_60 = get_cyb50_60min_data()
    
    if data_60 is not None:
        # 显示前5条数据预览
        print("\n数据预览:")
        print(data_60.head())
        
        # 导出到CSV
        export_to_csv(data_60)
        
        print("\n数据统计信息:")
        print(data_60.describe())
    else:
        print("60分钟数据获取失败")
    
    # 获取30分钟数据
    print("\n" + "=" * 50)
    print("【30分钟K线数据】")
    data_30 = get_cyb50_30min_data()
    
    if data_30 is not None:
        # 显示前5条数据预览
        print("\n数据预览:")
        print(data_30.head())
        
        # 导出到CSV（带30分钟标识）
        export_to_csv(data_30, filename=f"cyb50_30min_kline_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv")
        
        print("\n数据统计信息:")
        print(data_30.describe())
    else:
        print("30分钟数据获取失败")
    
    # 获取15分钟数据
    print("\n" + "=" * 50)
    print("【15分钟K线数据】")
    data_15 = get_cyb50_15min_data()
    
    if data_15 is not None:
        # 显示前5条数据预览
        print("\n数据预览:")
        print(data_15.head())
        
        # 导出到CSV（带15分钟标识）
        export_to_csv(data_15, filename=f"cyb50_15min_kline_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv")
        
        print("\n数据统计信息:")
        print(data_15.describe())
    else:
        print("15分钟数据获取失败")

if __name__ == "__main__":
    main()