MultiAssetFinalV2Strategy

from __future__ import annotations

import numpy as np
import pandas as pd
from pandas import DataFrame, Series

from freqtrade.strategy import IStrategy


TIMEFRAME = "1h"

EMA_FAST = 20
EMA_MID = 50
EMA_SLOW = 200

ATR_PERIOD = 14
RSI_PERIOD = 14

TREND_BREAKOUT_LOOKBACK = 20
TREND_MAX_ATR_PCT = 0.03
TREND_MOMENTUM_LOOKBACK = 24
MIN_TREND_MOMENTUM = 0.030

REGIME_BULL_MIN_SLOPE_EMA200 = 0.0
REGIME_CHOP_MAX_DISTANCE_EMA200 = 0.06
REGIME_CHOP_MAX_SLOPE_EMA200 = 0.0005

BULL_STRENGTH_MIN_SLOPE_EMA200 = 0.00002
BULL_STRENGTH_MIN_DISTANCE_EMA200 = 0.015
BULL_STRENGTH_MOMENTUM_LOOKBACK = 24
BULL_STRENGTH_MIN_EMA20_MOMENTUM = 0.002
BULL_STRENGTH_MIN_SCORE = 2

BULL_PERSISTENCE_LOOKBACK = 36
BULL_PERSISTENCE_MIN_RATIO = 0.50

FRESH_BULL_MAX_AGE_BARS = 72
BEAR_COOLDOWN_BARS = 12

BULL_TARGET_PCT = 1.0

STARTUP_CANDLE_COUNT = max(
    EMA_SLOW + 20,
    BULL_STRENGTH_MOMENTUM_LOOKBACK + BULL_PERSISTENCE_LOOKBACK + 5,
    FRESH_BULL_MAX_AGE_BARS + BEAR_COOLDOWN_BARS,
    TREND_BREAKOUT_LOOKBACK + 5,
)


def ensure_input_schema(df: pd.DataFrame) -> pd.DataFrame:
    required_cols = ["open", "high", "low", "close", "volume"]
    missing = [column for column in required_cols if column not in df.columns]
    if missing:
        raise ValueError(f"Missing required columns: {missing}")

    out = df.copy()
    if "date" in out.columns:
        out["date"] = pd.to_datetime(out["date"], utc=True)
        out = out.sort_values("date").reset_index(drop=True)
        out.index = pd.DatetimeIndex(out["date"].to_numpy())
        out.index.name = None
    else:
        out.index = pd.to_datetime(out.index, utc=True)
        out = out.sort_index()
        out.index.name = None

    out[required_cols] = out[required_cols].astype(float)
    return out


def _ma_from_frozen_research(close: Series, window: int) -> Series:
    # The frozen source uses vectorbt `MA.run(..., ewm=False)` defaults.
    return close.rolling(window=window, min_periods=window).mean()


def _rsi_from_frozen_research(close: Series, window: int) -> Series:
    delta = close.diff()
    gain = delta.clip(lower=0.0)
    loss = -delta.clip(upper=0.0)

    avg_gain = gain.rolling(window=window, min_periods=window).mean()
    avg_loss = loss.rolling(window=window, min_periods=window).mean()

    rs = avg_gain / avg_loss
    return 100.0 - (100.0 / (1.0 + rs))


def _atr_from_frozen_research(high: Series, low: Series, close: Series, window: int) -> Series:
    prev_close = close.shift(1)
    true_range = pd.concat(
        [
            high - low,
            (high - prev_close).abs(),
            (low - prev_close).abs(),
        ],
        axis=1,
    ).max(axis=1)

    atr = pd.Series(np.nan, index=close.index, dtype=float)
    if len(true_range) < window:
        return atr

    atr.iloc[window - 1] = true_range.iloc[:window].mean()
    for i in range(window, len(true_range)):
        atr.iloc[i] = ((atr.iloc[i - 1] * (window - 1)) + true_range.iloc[i]) / window
    return atr


def compute_indicators(df: pd.DataFrame) -> pd.DataFrame:
    df = ensure_input_schema(df)

    close = df["close"]
    high = df["high"]
    low = df["low"]

    ema20 = _ma_from_frozen_research(close, EMA_FAST)
    ema50 = _ma_from_frozen_research(close, EMA_MID)
    ema200 = _ma_from_frozen_research(close, EMA_SLOW)
    rsi = _rsi_from_frozen_research(close, RSI_PERIOD)
    atr = _atr_from_frozen_research(high, low, close, ATR_PERIOD)

    atr_pct = atr / close

    out = df.copy()
    out["ema20"] = ema20
    out["ema50"] = ema50
    out["ema200"] = ema200
    out["rsi"] = rsi
    out["atr"] = atr
    out["atr_pct"] = atr_pct
    out["ema200_slope"] = ema200.pct_change()
    out["distance_ema200"] = (close - ema200) / ema200
    out["atr_change"] = atr_pct.pct_change().abs()
    out["ema20_momentum"] = ema20.pct_change(BULL_STRENGTH_MOMENTUM_LOOKBACK)

    return out


def compute_regimes(ind: pd.DataFrame) -> pd.DataFrame:
    out = ind.copy()

    close = out["close"]
    ema50 = out["ema50"]
    ema200 = out["ema200"]
    ema200_slope = out["ema200_slope"]
    distance_ema200 = out["distance_ema200"]

    bull_regime = (
        (close > ema200)
        & (ema50 > ema200)
        & (ema200_slope > REGIME_BULL_MIN_SLOPE_EMA200)
    )

    chop_candidate = (
        (distance_ema200.abs() < REGIME_CHOP_MAX_DISTANCE_EMA200)
        & (ema200_slope.abs() < REGIME_CHOP_MAX_SLOPE_EMA200)
    )

    bull_regime = bull_regime.fillna(False)
    chop_regime = (chop_candidate & (~bull_regime)).fillna(False)
    bear_regime = ((~bull_regime) & (~chop_regime)).fillna(False)

    out["bull_regime"] = bull_regime
    out["chop_regime"] = chop_regime
    out["bear_regime"] = bear_regime

    return out


def compute_bull_strength_filter_v1(ind: pd.DataFrame) -> pd.DataFrame:
    out = ind.copy()

    ema200_slope = out["ema200_slope"]
    distance_ema200 = out["distance_ema200"]
    ema20_momentum = out["ema20_momentum"]
    close = out["close"]
    ema20 = out["ema20"]
    ema50 = out["ema50"]
    ema200 = out["ema200"]

    cond_slope = ema200_slope > BULL_STRENGTH_MIN_SLOPE_EMA200
    cond_distance = distance_ema200 > BULL_STRENGTH_MIN_DISTANCE_EMA200
    cond_momentum = ema20_momentum > BULL_STRENGTH_MIN_EMA20_MOMENTUM

    score = (
        cond_slope.fillna(False).astype(int)
        + cond_distance.fillna(False).astype(int)
        + cond_momentum.fillna(False).astype(int)
    )

    structure_ok = (
        (close > ema200)
        & (ema50 > ema200)
        & (ema20 > ema50)
    ).fillna(False)

    strong_bull = (structure_ok & (score >= BULL_STRENGTH_MIN_SCORE)).fillna(False)

    out["cond_slope"] = cond_slope.fillna(False)
    out["cond_distance"] = cond_distance.fillna(False)
    out["cond_momentum"] = cond_momentum.fillna(False)
    out["bull_strength_score"] = score
    out["strong_bull"] = strong_bull

    return out


def compute_bull_persistence_filter_v1(ind: pd.DataFrame) -> pd.DataFrame:
    out = ind.copy()

    prior_strong_bull = out["strong_bull"].astype(bool).shift(1, fill_value=False)

    strong_bull_ratio = (
        prior_strong_bull
        .rolling(BULL_PERSISTENCE_LOOKBACK, min_periods=BULL_PERSISTENCE_LOOKBACK)
        .mean()
    )

    persistent_strong_bull = (
        out["strong_bull"]
        & (strong_bull_ratio >= BULL_PERSISTENCE_MIN_RATIO)
    ).fillna(False)

    out["prior_strong_bull"] = prior_strong_bull
    out["strong_bull_ratio"] = strong_bull_ratio
    out["persistent_strong_bull"] = persistent_strong_bull

    return out


def compute_bull_age_bars(ind: pd.DataFrame) -> pd.DataFrame:
    out = ind.copy()

    bull_regime = out["bull_regime"].fillna(False)
    bull_age = pd.Series(0, index=out.index, dtype="int64")

    age = 0
    for i, is_bull in enumerate(bull_regime.values):
        if is_bull:
            age += 1
        else:
            age = 0
        bull_age.iloc[i] = age

    fresh_bull = ((bull_age > 0) & (bull_age <= FRESH_BULL_MAX_AGE_BARS)).fillna(False)

    out["bull_age_bars"] = bull_age
    out["fresh_bull"] = fresh_bull

    return out


def get_trend_signals(ind: pd.DataFrame) -> pd.DataFrame:
    out = ind.copy()

    close = out["close"]
    high = out["high"]
    ema20 = out["ema20"]
    ema50 = out["ema50"]
    ema200 = out["ema200"]
    atr_pct = out["atr_pct"]

    trend_filter = (
        (close > ema200)
        & (ema50 > ema200)
        & (ema20 > ema50)
    )

    rolling_high = high.shift(1).rolling(TREND_BREAKOUT_LOOKBACK).max()
    trend_momentum = close.pct_change(TREND_MOMENTUM_LOOKBACK)

    breakout = (
        (close > rolling_high)
        & (close.shift(1) > rolling_high.shift(1))
    )

    momentum_ok = trend_momentum >= MIN_TREND_MOMENTUM
    volatility_ok = atr_pct < TREND_MAX_ATR_PCT

    raw_entries = trend_filter & breakout & momentum_ok & volatility_ok
    entries = raw_entries & (~raw_entries.shift(1, fill_value=False))
    exits = close < ema50

    out["trend_filter"] = trend_filter.fillna(False)
    out["rolling_high"] = rolling_high
    out["trend_momentum"] = trend_momentum
    out["breakout"] = breakout.fillna(False)
    out["momentum_ok"] = momentum_ok.fillna(False)
    out["volatility_ok"] = volatility_ok.fillna(False)
    out["raw_entries"] = raw_entries.fillna(False)
    out["trend_entries"] = entries.fillna(False)
    out["trend_exits"] = exits.fillna(False)

    return out


def compute_bear_cooldown_mask(ind: pd.DataFrame) -> pd.DataFrame:
    out = ind.copy()

    bear_regime = out["bear_regime"].fillna(False)
    cooldown_mask = pd.Series(False, index=out.index)

    if BEAR_COOLDOWN_BARS > 0:
        bear_int = bear_regime.astype(int)
        for i in range(BEAR_COOLDOWN_BARS):
            cooldown_mask = cooldown_mask | bear_int.shift(i).fillna(0).astype(bool)

    out["bear_cooldown_mask"] = cooldown_mask.fillna(False)

    return out


def compute_v1_trend_only_signals(ind: pd.DataFrame) -> pd.DataFrame:
    out = ind.copy()

    raw_trend_entries = (
        out["bull_regime"]
        & out["trend_entries"]
        & (~out["bear_cooldown_mask"])
    ).fillna(False)

    trend_entries_fresh_bull = (raw_trend_entries & out["fresh_bull"]).fillna(False)
    trend_entries_mature_bull = (raw_trend_entries & (~out["fresh_bull"])).fillna(False)

    gated_fresh_bull_entries = (
        trend_entries_fresh_bull
        & out["persistent_strong_bull"]
    ).fillna(False)

    gated_mature_bull_entries = trend_entries_mature_bull.copy()

    filtered_trend_entries_fresh_bull = (
        trend_entries_fresh_bull
        & (~out["persistent_strong_bull"])
    ).fillna(False)

    filtered_trend_entries_mature_bull = pd.Series(False, index=out.index)

    trend_entries_active = (
        gated_fresh_bull_entries
        | gated_mature_bull_entries
    ).fillna(False)

    filtered_trend_entries = (
        filtered_trend_entries_fresh_bull
        | filtered_trend_entries_mature_bull
    ).fillna(False)

    trend_exits_active = (out["bull_regime"] & out["trend_exits"]).fillna(False)
    bear_exits_active = out["bear_regime"].fillna(False)

    final_entries = trend_entries_active.fillna(False)
    final_exits = (trend_exits_active | bear_exits_active).fillna(False)

    regime_label = pd.Series("bear", index=out.index)
    regime_label.loc[out["chop_regime"]] = "chop"
    regime_label.loc[out["bull_regime"]] = "bull"

    target_size = pd.Series(0.0, index=out.index)
    target_size.loc[trend_entries_active] = BULL_TARGET_PCT

    entry_score = pd.Series(np.nan, index=out.index, dtype=float)
    entry_score.loc[trend_entries_active] = (
        2000
        + out.loc[trend_entries_active, "ema20_momentum"].fillna(0).astype(float) * 100000
        - out.loc[trend_entries_active, "atr_pct"].fillna(0).astype(float) * 10000
    )

    out["raw_trend_entries_after_regime_and_cooldown"] = raw_trend_entries
    out["trend_entries_fresh_bull"] = trend_entries_fresh_bull
    out["trend_entries_mature_bull"] = trend_entries_mature_bull
    out["gated_fresh_bull_entries"] = gated_fresh_bull_entries
    out["gated_mature_bull_entries"] = gated_mature_bull_entries
    out["filtered_trend_entries_fresh_bull"] = filtered_trend_entries_fresh_bull
    out["filtered_trend_entries_mature_bull"] = filtered_trend_entries_mature_bull
    out["filtered_trend_entries"] = filtered_trend_entries
    out["trend_entries_active"] = trend_entries_active
    out["trend_exits_active"] = trend_exits_active
    out["bear_exits_active"] = bear_exits_active
    out["final_entries"] = final_entries
    out["final_exits"] = final_exits
    out["target_size"] = target_size
    out["regime_label"] = regime_label
    out["entry_score"] = entry_score

    return out


def build_v2_research_core(df: pd.DataFrame) -> pd.DataFrame:
    out = ensure_input_schema(df)
    out = compute_indicators(out)
    out = compute_regimes(out)
    out = compute_bull_strength_filter_v1(out)
    out = compute_bull_persistence_filter_v1(out)
    out = compute_bull_age_bars(out)
    out = compute_bear_cooldown_mask(out)
    out = get_trend_signals(out)
    out = compute_v1_trend_only_signals(out)
    return out


class MultiAssetFinalV2Strategy(IStrategy):
    INTERFACE_VERSION = 3
    can_short: bool = False

    timeframe = TIMEFRAME
    startup_candle_count: int = STARTUP_CANDLE_COUNT
    process_only_new_candles = True

    minimal_roi = {"0": 100}
    stoploss = -0.99

    use_exit_signal = True
    exit_profit_only = False
    ignore_roi_if_entry_signal = False

    order_types = {
        "entry": "limit",
        "exit": "limit",
        "stoploss": "market",
        "stoploss_on_exchange": False,
    }
    order_time_in_force = {"entry": "GTC", "exit": "GTC"}

    def version(self) -> str:
        return "2.0.0-research-v2-momentum"

    def populate_indicators(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        return build_v2_research_core(dataframe)

    def populate_entry_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        dataframe["enter_long"] = 0
        dataframe["enter_tag"] = ""

        fresh_entry_mask = dataframe["gated_fresh_bull_entries"].fillna(False)
        mature_entry_mask = dataframe["gated_mature_bull_entries"].fillna(False)

        dataframe.loc[fresh_entry_mask, "enter_long"] = 1
        dataframe.loc[fresh_entry_mask, "enter_tag"] = "trend_fresh_bull_persistent"

        dataframe.loc[mature_entry_mask, "enter_long"] = 1
        dataframe.loc[mature_entry_mask, "enter_tag"] = "trend_mature_bull"

        return dataframe

    def populate_exit_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        dataframe["exit_long"] = 0
        dataframe["exit_tag"] = ""

        trend_exit_mask = dataframe["trend_exits_active"].fillna(False)
        bear_exit_mask = dataframe["bear_exits_active"].fillna(False)

        dataframe.loc[trend_exit_mask, "exit_long"] = 1
        dataframe.loc[trend_exit_mask, "exit_tag"] = "trend_exit_ema50"

        dataframe.loc[bear_exit_mask, "exit_long"] = 1
        dataframe.loc[bear_exit_mask, "exit_tag"] = "bear_regime_exit"

        return dataframe