import polars as pl

from ldnode.loader.files import INDEX_DATA
from ldnode.calculator.schemas import BenchmarkType


def get_index_data_lf(benchmark: BenchmarkType) -> pl.LazyFrame:
    """ 获取指数数据 """
    match benchmark:
        case BenchmarkType.EQUALWEIGHT:
            benchmark_column = 'JSL'
        case BenchmarkType.GZ2000:
            benchmark_column = '399303.SZ'
        case BenchmarkType.HS300:
            benchmark_column = '000300.SH'
        case BenchmarkType.ZZ1000:
            benchmark_column = '000852.SH'
        case BenchmarkType.STOCK:
            benchmark_column = 'JSL' # 可转债基准为正股尚未实现
        case _:
            benchmark_column = str(benchmark) # 自选基准，直接用原始字符串做列名
    try:
        return (
            INDEX_DATA.output_lazyframe()
            .select(
                pl.col('trade_date').cast(pl.Date),
                pl.col(benchmark_column).alias('time_return_benchmark') # 基准每周期收益率
            )
        )
    except Exception as e:
        raise ValueError(f"获取基准数据失败: {e}")
