import polars as pl

from .factor import ProcessMissingDataMethod, NeutralizationMethod
from .screen import BuiltinScope


class STKNeutralizationMethod(NeutralizationMethod):
    """ 股票中性化方法 """
    # log(x) 对 x<=0 产生 -inf/NaN，会污染 OLS 全部残差，故先过滤为 null
    BARRA_SIZE = ('barra_size', lambda _: [pl.when(pl.col('circ_mv') > 0).then(pl.col('circ_mv').log()).otherwise(None)], False) # BARRA市值因子
    SW_INDUSTRY_1 = ('industry_1', lambda names: [pl.col(name) for name in names if name.startswith(f'industry_1_')], True) # 申万一级行业
    SW_INDUSTRY_2 = ('industry_2', lambda names: [pl.col(name) for name in names if name.startswith(f'industry_2_')], True) # 申万二级行业
    

class STKProcessMissingDataMethod(ProcessMissingDataMethod):
    """ 股票缺失值处理方式 """
    MARKET_MEAN = ('market_mean', lambda factor: pl.col(factor).fill_null(pl.col(factor).mean().over('trade_date')))  # 全市场平均数
    MARKET_MEDIAN = ('market_median', lambda factor: pl.col(factor).fill_null(pl.col(factor).median().over('trade_date')))  # 全市场中位数
    FORWARD_FILL = ('forward_fill', lambda factor: pl.col(factor).fill_null(strategy='forward').over('code'))  # 向前填充
    SW_INDUSTRY_1_MEAN = ('sw_industry_1_mean', lambda factor: pl.col(factor).fill_null(pl.col(factor).mean().over('trade_date', 'industry_1')))  # 申万一级行业平均数
    SW_INDUSTRY_1_MEDIAN = ('sw_industry_1_median', lambda factor: pl.col(factor).fill_null(pl.col(factor).median().over('trade_date', 'industry_1')))  # 申万一级行业中位数
    SW_INDUSTRY_2_MEAN = ('sw_industry_2_mean', lambda factor: pl.col(factor).fill_null(pl.col(factor).mean().over('trade_date', 'industry_2')))  # 申万二级行业平均数
    SW_INDUSTRY_2_MEDIAN = ('sw_industry_2_median', lambda factor: pl.col(factor).fill_null(pl.col(factor).median().over('trade_date', 'industry_2')))  # 申万二级行业中位数


class STKBuiltinScope(BuiltinScope):
    """股票内置交易所板块"""
    SH_MAINBOARD = ('sh_mainboard', lambda: pl.col('code').str.starts_with('60'))  # 沪市主板：60开头
    SZ_MAINBOARD = ('sz_mainboard', lambda: pl.col('code').str.starts_with('00'))  # 深市主板：00开头
    SH_SECONDBOARD = ('sh_secondboard', lambda: pl.col('code').str.starts_with('68'))  # 沪市科创板：68开头
    SZ_SECONDBOARD = ('sz_secondboard', lambda: pl.col('code').str.starts_with('30'))  # 深市创业板：30开头
    BJ_THIRDBOARD = ('bj_thirdboard', lambda: pl.Expr.or_(*[pl.col('code').str.starts_with(prefix) for prefix in ['4', '8', '9']]))  # 北交所：4、8、9开头

    SSE50 = ('000016.SH', lambda: pl.col(f'isin_sse50') == True) # 上证50
    HS300 = ('000300.SH', lambda: pl.col(f'isin_hs300') == True) # 沪深300
    ZZ500 = ('000510.SH', lambda: pl.col(f'isin_zz500') == True) # 中证500
    ZZ1000 = ('000852.SH', lambda: pl.col(f'isin_zz1000') == True) # 中证1000
    GZ2000 = ('399303.SZ', lambda: pl.col(f'isin_gz2000') == True) # 国证2000
