# -*- coding: utf-8 -*-
"""
---------------------------------------------
Created on 2025/6/23 09:49
@author: ZhangYundi
@email: yundi.xxii@outlook.com
@description: 基础转换器仓库
---------------------------------------------
"""

import inspect
import math
from collections.abc import Sequence
from typing import Literal

import numpy as np
import polars as pl
import ygo
from pandas import Timedelta
from sklearn.base import BaseEstimator, TransformerMixin

SPECIAL_VALUES = {
    "null": None,
    "none": None,
    "nan": None,
    "inf": np.inf,
    "-inf": -np.inf,
}


class Function(BaseEstimator, TransformerMixin):
    """
    将任意函数封装为支持 set_params 的 BaseEstimator

    Parameters
    ----------
    fn: callable
        方法对象
    kwargs: dict
        fn的参数

    Examples
    --------
    >>> import quda
    >>> func = Function(lambda date, ds_path: quda.sql(f"select * from {ds_path} where date='{date}'"))
    >>> func.set_params(ds_path="mc/stock_kline_day")
    """

    def __init__(self, fn: callable, **kwargs):
        self.fn = ygo.delay(fn)(**kwargs)
        self.params = {k: inspect.Parameter.empty for k in ygo.fn_signature_params(self.fn)}
        self.params.update(self.fn.stored_kwargs)

    def set_params(self, **params):
        self.fn = ygo.delay(self.fn)(**params)
        self.params.update(self.fn.stored_kwargs)
        return self

    def get_params(self, deep=True):
        return self.params

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        if self.fn is None:
            raise ValueError("Function not provided.")
        params = {k: v for k, v in self.fn.stored_kwargs.items() if v is not inspect.Parameter.empty}
        return self.fn(X, **params)


class FilterNotNull(BaseEstimator, TransformerMixin):
    """
    Drop all rows that contain one or more null values

    Parameters
    ----------
    subset: Sequence[str] | str | None
        剔除空值的字段

    """

    def __init__(self, subset: Sequence[str] | str | None = None):
        self.subset = subset

    def fit(self, X=None, y=None):
        return self

    def transform(self, X: pl.LazyFrame):
        if self.subset:
            return X.drop_nulls(self.subset)
        return X


class Cast(BaseEstimator, TransformerMixin):
    """
    类型转换

    Parameters
    ----------
    old: str | polars.DataType | polars.Expr
        原本类型或者需要转换的列
    new: polars.DataType
        转换的目标类型

    Examples
    --------
    >>> Cast("asset", pl.UInt16)

    """

    def __init__(self, old: pl.Expr | str | pl.DataType = None, new=None):
        self.old = old
        self.new = new
        self.old_expr: pl.Expr = None

    def fit(self, X, y=None):
        if isinstance(self.old, str):
            if self.old in SPECIAL_VALUES:
                self.old = SPECIAL_VALUES[self.old]
        self.old_expr = self.old if isinstance(self.old, pl.Expr) else pl.col(self.old_expr)
        return self

    def transform(self, X):
        return X.with_columns(self.old_expr.cast(self.new))


class Expression(BaseEstimator, TransformerMixin):
    """
    添加新列

    Parameters
    ----------
    exprs: list[polars.Expr]
        需要添加的列(使用`polars.Expr`表示)

    """

    def __init__(self, exprs: list[pl.Expr] | None = None):
        self.exprs = exprs

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        for expr in self.exprs:
            X = X.with_columns(expr)
        return X


class Imputer(BaseEstimator, TransformerMixin):
    """
    插值器

    Parameters
    ----------
    strategy
        插值策略
    over_spec
        用于滚动插值
    columns
        需要处理的列
    """

    def __init__(self,
                 strategy: Literal[
                               "forward", "backward", "mean", "zero", "max", "min", "one"] | None = "forward",
                 over_spec: dict | None = None,
                 columns: Sequence[str] | None = None):
        self.strategy = strategy
        self.over_spec = over_spec
        self.columns = [columns, ] if isinstance(columns, str) else columns
        self._exprs: list[pl.Expr] = list()

    def fit(self, X=None, y=None):
        return self

    def transform(self, X: pl.LazyFrame):
        cols = X.collect_schema().names() if not self.columns else self.columns
        if not self.over_spec:
            self._exprs = [pl.col(col).fill_null(self.strategy) for col in cols]
        else:
            self._exprs = [pl.col(col).fill_null(self.strategy).over(**self.over_spec) for col in self.columns]
        return X.with_columns(self._exprs)


class Replace(BaseEstimator, TransformerMixin):
    """
    替换器

    Parameters
    ----------
    columns: list
        需要替换的目标列
    old
        旧值, 可以是list, 如[np.inf, -np.inf]
    new
        新值
    """

    def __init__(self, columns: list | None = None, old=None, new=None):

        self.columns = columns
        self.old = old
        self.new = new

    def fit(self, X, y=None):
        if isinstance(self.old, str):
            if self.old in SPECIAL_VALUES:
                self.old = SPECIAL_VALUES[self.old]
        elif isinstance(self.old, list):
            self.old = [SPECIAL_VALUES[old] if old in SPECIAL_VALUES else old for old in self.old]
        if self.columns is None:
            self.columns = list()
        return self

    def transform(self, X):
        return X.with_columns(pl.col(c).replace(self.old, self.new) for c in self.columns)


class Target(BaseEstimator, TransformerMixin):
    """
    添加目标收益列

    Parameters
    ----------
    price_tag
        价格列，支持sql语法，如 `if(ask_p1 > 0, bid_p1 > 0, (ask_p1+bid_p1)/2, null) as price`
    frequency
        数据集频率, 默认`3s`
    target
        目标收益，默认`5min`
    partition_by
        分区
    order_by
        排序
    """

    def __init__(self,
                 price_tag: str = "price",
                 frequency: str = "3s",
                 target: str = "5min",
                 partition_by: Sequence[str] | str = "asset",
                 order_by: Sequence[str] | str | None = "time", ):
        self.price_tag = price_tag
        self.frequency = frequency
        self.target = target
        self.partition_by = partition_by
        self.order_by = order_by

    def fit(self, X=None, y=None):
        return self

    def transform(self, X: pl.LazyFrame):
        over_spec = {"partition_by": self.partition_by}
        if self.order_by:
            over_spec["order_by"] = self.order_by
        freq_secs = int(Timedelta(self.frequency).seconds)
        expr_target = (
            (
                    pl.sql_expr(self.price_tag)
                    .shift(-math.ceil(Timedelta(self.target).seconds / freq_secs))
                    .over(**over_spec) / pl.sql_expr(self.price_tag) - 1
            ).alias(self.target)
        )
        expr_inf = pl.col(self.target).replace([float("inf"), float("-inf")], None)
        exprs = [expr_target, expr_inf]
        for expr in exprs:
            X = X.with_columns(expr)
        return X.fill_nan(None)
