import polars as pl

from ldnode.calculator.schemas.models.screen import (
    CBScreenParameters, 
    STKScreenParameters, 
    ETFScreenParameters,
)


EPSILON = 1e-8 # 精度


def gen_cb_filter_exprs(parameters: CBScreenParameters) -> list[pl.Expr]:
    """ 生成选债快捷排除设置公式 """
    filter_exprs = []
    # 标记处于赎回状态相关的标的
    if exclude_call := parameters.get('exclude_call'):
        if '公告实施强赎' in exclude_call:
            parameters['exclude_call'] = exclude_call + ['已公告强赎']
        filter_exprs.append(pl.col('is_call').is_in(exclude_call))
    # 标记上市前N个自然日标的（兜底 int 转换，防止前端传入字符串）
    if parameters.get('exclude_list_days'):
        filter_exprs.append(pl.col('list_days') <= int(parameters['exclude_list_days']))
    # 标记ST
    if parameters.get('exclude_st'):
        filter_exprs.append(pl.col('name_stk').str.to_uppercase().str.contains('ST'))
    # 标记市场
    if exclude_market := parameters.get('exclude_market'):
        _pattern = f"^({'|'.join(exclude_market)})"
        filter_exprs.append(pl.col('code_stk').str.contains(_pattern))      #(标的池股票排除参考（不用再取反）)
    # 标记行业
    if exclude_first_industries := parameters.get('exclude_first_industries'):
        filter_exprs.append(pl.col('industry_1').is_in(exclude_first_industries))
    if exclude_second_industries := parameters.get('exclude_second_industries'):
        filter_exprs.append(pl.col('industry_2').is_in(exclude_second_industries))
    if exclude_third_industries := parameters.get('exclude_third_industries'):
        filter_exprs.append(pl.col('industry_3').is_in(exclude_third_industries))
    # 标记企业类型
    if exclude_orgform := parameters.get('exclude_orgform'):
        filter_exprs.append(pl.col('orgform').is_in(exclude_orgform))
    # 标记地域
    if exclude_area := parameters.get('exclude_area'):
        filter_exprs.append(pl.col('area').is_in(exclude_area))
    # 标记外部评级
    if exclude_rating := parameters.get('exclude_rating'):
        filter_exprs.append(pl.col('rating').is_in(exclude_rating))
    # 标记yy评级
    if exclude_yy_rating := parameters.get('exclude_yy_rating'):
        filter_exprs.append(pl.col('yy_rating').is_in(exclude_yy_rating))
    # 标记指定债券
    if exclude_ignore_list := parameters.get('exclude_ignore_list'):
        filter_exprs.append(pl.col('code').is_in(exclude_ignore_list))
    return filter_exprs


def gen_stk_filter_exprs(parameters: STKScreenParameters) -> list[pl.Expr]:
    """ 生成选股快捷排除设置公式 """
    filter_exprs = []
    # 标记上市前N个自然日标的（兜底 int 转换，防止前端传入字符串）
    if parameters.get('exclude_list_days'):
        filter_exprs.append(pl.col('list_days') <= int(parameters['exclude_list_days']))
    # 标记ST
    if parameters.get('exclude_st'):
        filter_exprs.append(pl.col('name').str.to_uppercase().str.contains('ST'))
    # 标记退整理期股票
    if parameters.get('exclude_delisting'):
        filter_exprs.append(pl.col('name').str.contains('退'))
    # 标记市场
    if exclude_market := parameters.get('exclude_market'):
        _pattern = f"^({'|'.join(exclude_market)})"
        filter_exprs.append(pl.col('code').str.contains(_pattern))
    # 标记行业
    if exclude_first_industries := parameters.get('exclude_first_industries'):
        filter_exprs.append(pl.col('industry_1').is_in(exclude_first_industries))
    if exclude_second_industries := parameters.get('exclude_second_industries'):
        filter_exprs.append(pl.col('industry_2').is_in(exclude_second_industries))
    if exclude_third_industries := parameters.get('exclude_third_industries'):
        filter_exprs.append(pl.col('industry_3').is_in(exclude_third_industries))
    # 标记企业类型
    # if exclude_orgform := strategy_config.get('exclude_orgform'):
    #     filter_exprs.append(col('orgform').is_in(exclude_orgform))
    # 标记地域
    if exclude_area := parameters.get('exclude_area'):
        filter_exprs.append(pl.col('area').is_in(exclude_area))
    # 标记指定股票
    if exclude_ignore_list := parameters.get('exclude_ignore_list'):
        filter_exprs.append(pl.col('code').is_in(exclude_ignore_list))
    return filter_exprs


def gen_etf_filter_exprs(parameters: ETFScreenParameters) -> list[pl.Expr]:
    """ 生成选基快捷排除设置公式 """
    filter_exprs = []
     # 标记市场
    if exclude_market := parameters.get('exclude_market'):
        _pattern = f"^({'|'.join(exclude_market)})"
        filter_exprs.append(pl.col('code').str.contains(_pattern))
    # 标记指定基金
    if exclude_ignore_list := parameters.get('exclude_ignore_list'):
        filter_exprs.append(pl.col('code').is_in(exclude_ignore_list))
    return filter_exprs
