# -*- coding: utf-8 -*-
"""
因子数据功能测试

测试 QuoteFactor 类的 get_factors, get_factor_data, get_factor_static 方法
"""

import sys
import os

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from pyqdt import QuoteFactor

# 数据路径
DATA_PATH = '/data/QuantData/pyqdt_csv'


def test_get_factors():
    """测试获取所有因子列表"""
    print("\n=== test_get_factors ===")

    factor = QuoteFactor(data_path=DATA_PATH)

    # 获取所有因子
    df = factor.get_factors()
    print(f"所有因子数量: {len(df)}")
    print(df.head())

    # 按类别筛选
    df_quality = factor.get_factors(category='quality')
    print(f"\nquality 类因子数量: {len(df_quality)}")
    print(df_quality.head())

    print("✅ test_get_factors 通过")


def test_get_factor_data():
    """测试获取股票因子时序数据"""
    print("\n=== test_get_factor_data ===")

    factor = QuoteFactor(data_path=DATA_PATH)

    # 单只股票，使用 SH/SZ 后缀
    df = factor.get_factor_data('600000.SH', factors=['roe_ttm', 'roa_ttm', 'market_cap']
                                , start_date='2005-01-01', end_date='2005-03-31')
    print(f"单只股票因子数据 (SH后缀): {len(df)} 行")
    print(df.head())

    # 也支持 XSHG/XSHE 后缀
    df = factor.get_factor_data('000001.XSHE', factors=['roe_ttm', 'market_cap']
                                , start_date='2005-01-01', end_date='2005-03-31')
    print(f"\n单只股票因子数据 (XSHE后缀): {len(df)} 行")
    print(df.head())

    # 多只股票
    df = factor.get_factor_data(['600000.SH', '000001.SZ']
                                , factors=['roe_ttm', 'market_cap']
                                , start_date='2005-01-01', end_date='2005-03-31')
    print(f"\n多只股票因子数据: {len(df)} 行")
    print(df.head())

    print("✅ test_get_factor_data 通过")


def test_get_factor_static():
    """测试获取全市场因子快照"""
    print("\n=== test_get_factor_static ===")

    factor = QuoteFactor(data_path=DATA_PATH)

    # 获取快照
    df = factor.get_factor_static(date='2005-03-15', factors=['roe_ttm', 'market_cap'])
    print(f"因子快照数据: {len(df)} 行")
    print(df.head())

    # 指定交易所
    df_sh = factor.get_factor_static(date='2005-03-15', factors=['roe_ttm'], exch=['SH'])
    print(f"\n仅上海交易所: {len(df_sh)} 行")
    print(df_sh.head())

    print("✅ test_get_factor_static 通过")


if __name__ == '__main__':
    test_get_factors()
    test_get_factor_data()
    test_get_factor_static()
    print("\n=== 所有测试通过 ===")