mock_data.py

"""
模拟数据生成模块
用于生成测试数据，以便程序在没有账号的情况下也能进行模拟测试
"""

import logging
import os
from typing import Dict, List, Any
import pandas as pd
import numpy as np
from datetime import datetime, timedelta


logger = logging.getLogger("AIQuant.mock_data")
logger.setLevel(logging.INFO)


def generate_random_kline_data(
    symbol: str,
    start_date: str = "2023-01-01",
    end_date: str = "2023-12-31",
    frequency: str = "d",
    volatility: float = 0.01,
    trend: float = 0.0015
) -> pd.DataFrame:
    """
    生成随机 K 线数据

    Args:
        symbol: 股票代码
        start_date: 起始日期
        end_date: 结束日期
        frequency: 频率（d: 日线，h: 小时线，m: 分钟线）
        volatility: 波动率
        trend: 趋势强度

    Returns:
        生成的 K 线数据 DataFrame
    """
    logger.info(f"生成 {symbol} 的模拟 K 线数据")

    # 生成日期范围
    start = datetime.strptime(start_date, "%Y-%m-%d")
    end = datetime.strptime(end_date, "%Y-%m-%d")
    days = (end - start).days + 1

    if frequency == "d":
        dates = [start + timedelta(days=i) for i in range(days)]
    elif frequency == "h":
        dates = [start + timedelta(hours=i) for i in range(days * 24)]
    elif frequency == "m":
        dates = [start + timedelta(minutes=i) for i in range(days * 24 * 60)]
    else:
        raise ValueError(f"不支持的频率: {frequency}")

    # 生成随机价格序列
    np.random.seed(hash(symbol) % 4294967296)  # 使用股票代码生成随机种子
    price = 10.0 + np.random.uniform(-5, 5)  # 随机起始价格
    prices = [price]

    # 使用几何布朗运动模型生成价格序列
    for i in range(1, len(dates)):
        # 几何布朗运动公式：dS/S = μdt + σdW
        drift = trend  # 漂移项（预期收益率）
        volatility_term = volatility * np.random.normal()  # 波动率项
        # 价格变化
        dS = price * (drift + volatility_term)
        price = price + dS
        # 确保价格不会太离谱
        price = max(0.01, price)
        prices.append(price)

    # 生成 OHLCV 数据
    high = [p * (1 + np.random.uniform(0, volatility)) for p in prices]
    low = [p * (1 - np.random.uniform(0, volatility)) for p in prices]
    open_price = [
        p * (1 + np.random.uniform(-volatility/2, volatility/2))
        for p in prices
    ]
    close = prices

    volume = np.random.randint(1000000, 10000000, size=len(dates))

    # 创建 DataFrame
    data = {
        "datetime": dates,
        "open": open_price,
        "high": high,
        "low": low,
        "close": close,
        "volume": volume
    }

    df = pd.DataFrame(data)
    df.set_index("datetime", inplace=True)

    # 计算均线等基础指标
    df["ma5"] = df["close"].rolling(window=5).mean()
    df["ma10"] = df["close"].rolling(window=10).mean()
    df["ma20"] = df["close"].rolling(window=20).mean()
    df["ma60"] = df["close"].rolling(window=60).mean()

    logger.info(f"生成成功，共 {len(df)} 条记录")
    return df


def generate_stock_pool(n: int = 10) -> List[str]:
    """
    生成股票池

    Args:
        n: 股票数量

    Returns:
        股票池列表
    """
    logger.info(f"生成包含 {n} 只股票的股票池")

    stock_pool = []
    for i in range(n):
        # 随机生成股票代码（模拟沪深交易所）
        exchange = "SZ" if np.random.rand() > 0.5 else "SH"
        code = str(np.random.randint(100000, 999999)).zfill(6)
        stock_pool.append(f"{code}.{exchange}")

    logger.info(f"股票池生成成功: {stock_pool}")
    return stock_pool


def generate_mock_account_data() -> Dict[str, Any]:
    """
    生成模拟账户数据

    Returns:
        模拟账户数据
    """
    logger.info("生成模拟账户数据")

    account_data = {
        "account": "1234567890",
        "balance": 1000000.0,
        "available": 800000.0,
        "margin": 200000.0,
        "profit": 50000.0,
        "positions": {
            "000001.SZ": {
                "symbol": "000001.SZ",
                "name": "平安银行",
                "volume": 10000,
                "price": 12.50,
                "value": 125000.0
            },
            "600036.SH": {
                "symbol": "600036.SH",
                "name": "招商银行",
                "volume": 5000,
                "price": 35.20,
                "value": 176000.0
            }
        }
    }

    logger.info("账户数据生成成功")
    return account_data


def cache_data_to_local(data: pd.DataFrame,
                        symbol: str,
                        frequency: str = "d") -> str:
    """
    将数据缓存到本地

    Args:
        data: 数据
        symbol: 股票代码
        frequency: 频率

    Returns:
        缓存文件路径
    """
    logger.info(f"将 {symbol} 的数据缓存到本地")

    # 创建数据缓存目录
    cache_dir = "data_cache"
    if not os.path.exists(cache_dir):
        os.makedirs(cache_dir, exist_ok=True)

    # 生成文件名
    filename = f"{symbol}_{frequency}_{datetime.now().strftime('%Y%m%d')}.csv"
    filepath = os.path.join(cache_dir, filename)

    # 保存数据
    data.to_csv(filepath)

    logger.info(f"数据缓存成功: {filepath}")
    return filepath


def load_data_from_local(symbol: str, frequency: str = "d") -> pd.DataFrame:
    """
    从本地加载数据

    Args:
        symbol: 股票代码
        frequency: 频率

    Returns:
        加载的数据
    """
    logger.info(f"从本地加载 {symbol} 的数据")

    cache_dir = "data_cache"
    if not os.path.exists(cache_dir):
        logger.warning("数据缓存目录不存在")
        return pd.DataFrame()

    # 查找最新的缓存文件
    files = [f for f in os.listdir(cache_dir) if f.startswith(
        f"{symbol}_{frequency}") and f.endswith(".csv")]
    if not files:
        logger.warning(f"未找到 {symbol} 的缓存数据")
        return pd.DataFrame()

    # 按时间戳排序，获取最新的文件
    files.sort(key=lambda x: os.path.getmtime(
        os.path.join(cache_dir, x)), reverse=True)
    latest_file = files[0]
    filepath = os.path.join(cache_dir, latest_file)

    # 加载数据
    df = pd.read_csv(filepath, index_col=0, parse_dates=True)

    logger.info(f"数据加载成功: {filepath}, 共 {len(df)} 条记录")
    return df


# 测试函数
if __name__ == "__main__":
    import os

    # 测试生成 K 线数据
    symbol = "000001.SZ"
    kline_data = generate_random_kline_data(symbol)
    print(f"K 线数据形状: {kline_data.shape}")
    print(kline_data.head())

    # 测试生成股票池
    stock_pool = generate_stock_pool(5)
    print(f"股票池: {stock_pool}")

    # 测试生成账户数据
    account_data = generate_mock_account_data()
    print(f"账户数据: {account_data}")

    # 测试数据缓存和加载
    filepath = cache_data_to_local(kline_data, symbol)
    loaded_data = load_data_from_local(symbol)
    print(f"加载的数据形状: {loaded_data.shape}")

    # 清理测试文件
    if os.path.exists(filepath):
        os.remove(filepath)