# Copyright (c) ZhangYundi.
# Licensed under the MIT License. 
# Created on 2026/1/11 01:05
# Description:

from enum import Enum

import lidb
from sklearn.pipeline import Pipeline
import polars as pl

from .backtest import backtest, bt
from .weight_fns import qcut
import polars.selectors as cs


# ======================== 策略基类 ========================
class StrategyType(Enum):
    LONG_ONLY = "long_only"  # 纯多
    LONG_SHORT = "long_short"  # 多空
    HEDGE_INDEX = "hedge_index"  # 对冲指数


class Strategy(Pipeline):

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

    def run(self, input):
        return self.transform(input)


class FactorStrategy(Strategy):
    """因子策略"""

    def __init__(self, weight_tr=None):
        if weight_tr is None:
            weight_tr = qcut
        self.weight_tr = weight_tr
        steps = [weight_tr, bt]
        super().__init__(steps)

    def run(self, input: pl.DataFrame | pl.LazyFrame):
        index = ("date",  "time", "asset")
        weight_df = self[0].transform(input)
        # 对齐
        weight_df = lidb.from_polars(weight_df, index=index, align=True).data
        weight_df = (input
                     .join(weight_df, on=index, how="right", suffix="_right")
                     .with_columns(target_weight=pl.col("target_weight").fill_null(0.0))
                     .drop(cs.ends_with("_right"))
                     .select(*index, cs.exclude(index)))

        return self[1].transform(weight_df)