# source: https://raw.githubusercontent.com/Kazu-takahashi1213/freqtrade/806c73bec8c269addc056797a4778c6fcd00be50/strategies/ShortTermStrategy.py
"""Freqtrade strategy using LightGBM predictions."""

from __future__ import annotations

from typing import Any, Dict, List
import os

import pandas as pd
from freqtrade.strategy.interface import IStrategy

from src.data_pipeline.feature_store import compute_features
from src.model.trainer import load_model, train_model, prepare_dataset


class Github_Kazu_takahashi1213_freqtrade__ShortTermStrategy__20250618_172905(IStrategy):
    minimal_roi = {"0": 0.01}
    stoploss = -0.03
    timeframe = "5m"

    model: Any = None
    model_path: str = os.path.join("models", "lightgbm.txt")
    feature_columns: List[str] | None = None

    def __init__(self, config: Dict) -> None:
        super().__init__(config)
        self.model = load_model(self.model_path)

    def populate_indicators(self, dataframe: pd.DataFrame, metadata: Dict) -> pd.DataFrame:
        """Add technical indicator features."""
        df = compute_features(dataframe)

        if self.feature_columns is None:
            self.feature_columns = [c for c in df.columns if c not in {"date", "direction"}]

        return df

    def populate_entry_trend(self, dataframe: pd.DataFrame, metadata: Dict) -> pd.DataFrame:
        """Generate entry signals using the trained model."""
        if self.model is None:
            self.model = load_model(self.model_path)

        if self.model is None and "direction" in dataframe.columns:
            self.model = train_model(dataframe.dropna(), model_path=self.model_path)
        elif self.model is None:
            train_df = prepare_dataset(dataframe)
            self.model = train_model(train_df, model_path=self.model_path)

        if self.model and self.feature_columns:
            features = dataframe[self.feature_columns].fillna(0)
            preds = self.model.predict(features)
            signal = preds.argmax(axis=1)
            dataframe.loc[signal == 2, "enter_long"] = 1
            dataframe.loc[signal == 0, "enter_short"] = 1

        return dataframe

    def populate_exit_trend(self, dataframe: pd.DataFrame, metadata: Dict) -> pd.DataFrame:
        """Simple exit conditions based on price move."""
        dataframe["exit_long"] = dataframe["close"] > dataframe["close"].shift(1)
        dataframe["exit_short"] = dataframe["close"] < dataframe["close"].shift(1)
        return dataframe