# Copyright (c) ZhangYundi.
# Licensed under the MIT License. 
# Created on 2026/2/25 16:13
# Description:
from typing import Sequence, Literal

import polars as pl
import polars.selectors as cs
from silars.transformer import as_transformer
from sklearn.pipeline import Pipeline
import polars_ds as pds


class Preprocessor(Pipeline):

    def __init__(self, steps):
        steps = [
            (step.__class__.__name__, step)
            for step in steps
        ]
        super().__init__(steps)

    def transform(self, input):
        if isinstance(input, pl.DataFrame):
            input = input.lazy()
        return super().transform(input)


@as_transformer
def fillna_median(input: pl.LazyFrame, columns: Sequence[str], index_names: Sequence[str] = ("date", "time", "asset")):
    over = set(index_names) - {"asset"}
    return input.with_columns(cs.by_name(columns).fill_null(cs.by_name(columns).median().over(over)))

@as_transformer
def fillna_indust(input: pl.LazyFrame, columns: Sequence[str], index_names: Sequence[str] = ("date", "time", "asset")):
    """填充空值为行业中位数"""
    over = set(index_names) - {"asset"}
    over.add("group")
    return input.with_columns(cs.by_name(columns).fill_null(cs.by_name(columns).median().over(over)))

@as_transformer
def winsorize_mad(input: pl.LazyFrame, columns: Sequence[str], bias: float = 5.0, index_names: Sequence[str] = ("date", "time", "asset")):
    """
    中位数绝对偏差, 5 * 1.4826中位数绝对偏差
    """
    over = set(index_names) - {"asset"}
    selector = cs.by_name(columns)
    mid = selector.median().over(over)
    mad = (selector - mid).abs().median().over(over)
    lower = mid - bias * mad
    upper = mid + bias * mad
    return input.with_columns(selector.clip(lower_bound=lower, upper_bound=upper))

@as_transformer
def standardize(input: pl.LazyFrame, columns: Sequence[str], index_names: Sequence[str] = ("date", "time", "asset")):
    """流通市值标准化"""
    over = set(index_names) - {"asset"}
    selector = cs.by_name(columns)
    mu = selector.mean().over(over)
    std = selector.std().over(over)
    return input.with_columns((selector - mu) /std)

@as_transformer
def standardize_by_fcap(input: pl.LazyFrame, columns: Sequence[str], index_names: Sequence[str] = ("date", "time", "asset")):
    """流通市值标准化"""
    over = set(index_names) - {"asset"}
    selector = cs.by_name(columns)
    weights = pl.col("FloatCap").sqrt()
    weights = weights / weights.sum().over(over)
    mu = (selector * weights).sum().over(over)
    std = (weights * (selector - mu) ** 2).sum().over(over).sqrt()
    return input.with_columns((selector - mu) /std)

@as_transformer
def neutralize(input: pl.LazyFrame,
               columns: Sequence[str],
               index_names: Sequence[str] = ("date", "time", "asset"),
               null_policy: Literal["raise", "skip", "one", "zero", "ignore"] = "skip",):
    """
    中性化：缩尾->市值加权标准化->行业市场中性化
    Parameters
    ----------
    input: polars.LazyFrame
        需要包含 FloatCap , group 字段
    columns: Sequence[str]
    index_names: Sequence[str]
    null_policy: Literal["raise", "skip", "one", "zero", "ignore"] = "skip",

    Returns
    -------

    """
    cols = input.collect_schema().names()
    # raw_index: pl.LazyFrame = input.select(index_names)
    over = set(index_names) - {"asset"}
    # factors = cs.by_name(columns)
    cap = pl.col("FloatCap")
    weights = cap.sqrt()
    weights = weights/weights.sum().over(over)
    data: pl.DataFrame = input.drop_nulls(subset=["FloatCap", "group"]).with_columns(_w=weights).collect()

    # 中性化
    industs = data["group"].unique().to_list()
    data = data.to_dummies("group")
    res = (data
           .lazy()
           .group_by(over)
           .agg(pl.col("asset"),
                *[pds.lin_reg(*[f"group_{ind}" for ind in industs],
                              target=name,
                              add_bias=True,
                              return_pred=True,
                              weights="_w",
                              null_policy=null_policy,
                              ).alias(name)
                for name in columns])
           .explode("asset", *columns)
           .select(*index_names, *[pl.col(c).struct[1].alias(c) for c in columns])
           .join(input.drop(columns), on=index_names, how="right")
           .select(cols))
    return res

def new_preprocessor(columns: Sequence[str],
                     index_names: Sequence[str] = ("date", "time", "asset"),
                     steps: list | None = None):
    """
    创建处理器: 默认 填充行业中位数->填充中位数->5倍mad缩尾->市值标准化->行业中性化->标准化

    Parameters
    ----------
    columns: Sequence[str]
    index_names: Sequence[str]
    steps: list | None

    Returns
    -------

    """
    if not steps:
        steps = [
            fillna_indust.set_params(columns=columns, index_names=index_names),
            fillna_median.set_params(columns=columns, index_names=index_names),
            winsorize_mad.set_params(columns=columns, index_names=index_names),
            standardize_by_fcap.set_params(columns=columns, index_names=index_names),
            neutralize.set_params(columns=columns, index_names=index_names, null_policy='skip'),
            standardize.set_params(columns=columns, index_names=index_names),
        ]
    else:
        steps = [step.set_params(columns=columns, index_Names=index_names) for step in steps]
    return Preprocessor(steps=steps)
