import polars as pl

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


class CBNeutralizationMethod(NeutralizationMethod):
    """ 转债中性化方法 """
    # log(x) 对 x<=0 产生 -inf/NaN，会污染 OLS 全部残差，故先过滤为 null
    LOG_REMAIN_SIZE = ('remain_size', lambda _: [pl.when(pl.col('remain_size') > 0).then(pl.col('remain_size').log()).otherwise(None)], False) # 剩余规模
    LOG_REMAIN_CAP = ('remain_cap', lambda _: [pl.when(pl.col('remain_cap') > 0).then(pl.col('remain_cap').log()).otherwise(None)], False) # 剩余市值


class CBProcessMissingDataMethod(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'))  # 向前填充
    

class CBBuiltinScope(BuiltinScope):
    """ 转债内置交易所板块 """    
    SH_MAINBOARD = ('sh_mainboard', lambda: pl.col('code').str.starts_with('11'))  # 沪市转债：11开头
    SZ_MAINBOARD = ('sz_mainboard', lambda: pl.col('code').str.starts_with('12'))  # 深市转债：12开头
