from typing import Callable, Any
from enum import StrEnum, IntEnum

import polars as pl


class FactorType(IntEnum):
    """ 因子类型 """
    BUILDIN = 0 # 内置
    COMPUTED = 1 # 计算


class ExprType(IntEnum):
    """ 表达式类型 """
    SQL = 0 # sql表达式
    POLARS = 1 # polars表达式
    DSL = 2 # dsl表达式


class RankScope(IntEnum):
    """ 打分范围 """
    ALL = 0 # 全部范围
    SW_INDUSTRY_1 = 1 # 申万一级行业
    SW_INDUSTRY_2 = 2 # 申万二级行业
    SW_INDUSTRY_3 = 3 # 申万三级行业


class ExcludeOperator(StrEnum):
    """ 排除因子比较操作 """
    compare: Callable[[str, Any], pl.Expr]
    
    def __new__(cls, value: str, compare: Callable[[str, Any], pl.Expr]) -> 'ExcludeOperator':
        obj = str.__new__(cls, value)
        obj._value_ = value
        obj.compare = compare
        return obj

    # strict=False：String 列中非数值字符串转为 null，null 比较结果为 null（视为不满足条件），
    # 兼容 is_call 等原本为数值但被 cast 为 String 存储的因子
    GE = ('ge', lambda factor, value: pl.col(factor).cast(pl.Float64, strict=False) >= value) # 大于等于
    GT = ('gt', lambda factor, value: pl.col(factor).cast(pl.Float64, strict=False) > value) # 大于
    LE = ('le', lambda factor, value: pl.col(factor).cast(pl.Float64, strict=False) <= value) # 小于等于
    LT = ('lt', lambda factor, value: pl.col(factor).cast(pl.Float64, strict=False) < value) # 小于
    NE = ('ne', lambda factor, value: pl.col(factor).ne_missing(value)) # 不等于
    EQ = ('eq', lambda factor, value: pl.col(factor) == value) # 等于
    ISNULL = ('isnull', lambda factor, _: pl.col(factor).is_null()) # 空值
    GT_PCT = ('gt_pct', lambda factor, value: pl.col(f'{factor}_pct') > value) # 百分比大于
    LT_PCT = ('lt_pct', lambda factor, value: pl.col(f'{factor}_pct') < value) # 百分比小于
    CONTAINS = ('contains', lambda factor, value: pl.col(factor).str.contains(value)) # 包含
    NOT_CONTAINS = ('not_contains', lambda factor, value: ~pl.col(factor).str.contains(value)) # 不包含
    STARTS_WITH = ('starts_with', lambda factor, value: pl.col(factor).str.starts_with(value)) # 以...开头
    ENDS_WITH = ('ends_with', lambda factor, value: pl.col(factor).str.ends_with(value)) # 以...结尾
    NOT_STARTS_WITH = ('not_starts_with', lambda factor, value: ~pl.col(factor).str.starts_with(value)) # 不以...开头
    NOT_ENDS_WITH = ('not_ends_with', lambda factor, value: ~pl.col(factor).str.ends_with(value)) # 不以...结尾


class ProcessMissingDataMethod(StrEnum):
    """ 缺失值处理方式 """
    process: Callable[[str], pl.Expr]
    
    def __new__(cls, value: str, process: Callable[[str], pl.Expr]) -> 'ProcessMissingDataMethod':
        obj = str.__new__(cls, value)
        obj._value_ = value
        obj.process = process
        return obj


class NeutralizationMethod(StrEnum):
    """ 中性化方式 """
    process: Callable[[list[str]], list[pl.Expr]]
    is_discrete: bool
    
    def __new__(cls, value: str, process: Callable[[list[str]], list[pl.Expr]], is_discrete: bool) -> 'NeutralizationMethod':
        obj = str.__new__(cls, value)
        obj._value_ = value
        obj.process = process
        obj.is_discrete = is_discrete
        return obj
