from freqtrade.strategy.interface import IStrategy
from freqtrade.freqai.data_kitchen import FreqaiDataKitchen
from pandas import DataFrame
import numpy as np
import logging

logger = logging.getLogger(__name__)

class AIEnhancedStrategy(IStrategy):
    INTERFACE_VERSION = 3

    can_short = False  # 只做多
    timeframe = '5m'

    process_only_new_candles = True
    use_exit_signal = True
    exit_profit_only = False
    ignore_roi_if_entry_signal = False

    def feature_engineering_expand_all(self, dataframe: DataFrame, period: int,
                                      metadata: dict, **kwargs) -> DataFrame:
        # 自定义特征工程
        for col in ['open', 'high', 'low', 'close', 'volume']:
            dataframe[f'{col}_pct_change'] = dataframe[col].pct_change(periods=period)
        return dataframe

    def set_freqai_targets(self, dataframe: DataFrame, metadata: dict, **kwargs) -> DataFrame:
        # 目标变量：未来n根K线的收盘价变化
        dataframe['&s-close_pct'] = dataframe['close'].pct_change(periods=5).shift(-5)
        return dataframe

    def populate_indicators(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        dk = self.freqai_info.get("dk", None)
        if not dk:
            raise ValueError("FreqaiDataKitchen is not available.")

        dataframe = dk.feature_engineering_standard(dataframe, metadata, self)

        if self.config["runmode"].value in ("live", "dry_run"):
            dataframe = dk.live_models(self, dataframe, metadata=metadata)

        return dataframe

    def populate_entry_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        # 预测值大于阈值时开多仓
        dataframe.loc[
            (dataframe["&s-close_pct"] > 0.002),
            "enter_long"
        ] = 1
        return dataframe

    def populate_exit_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        # 预测值小于负阈值时平仓
        dataframe.loc[
            (dataframe["&s-close_pct"] < -0.001),
            "exit_long"
        ] = 1
        return dataframe