engine.py

"""
engine.py — Feature Engineering + Heavy XGBoost Engine + signal mapping.

Features: RSI, MACD, ATR, VWAP distance, Bollinger position/width, EMA
crossovers, volume delta, candle momentum, support/resistance distance,
volatility compression, regime score.

Model: XGBoost classifier trained per-symbol on the candle history.
Label  = does the close 4 candles ahead beat fees+buffer (BUY-worthy)?
Outputs: probability_up, expected_return, downside_risk, confidence,
market_regime, 4-candle forecast path, action (BUY/HOLD/SELL).
"""

import logging

import numpy as np
import pandas as pd

log = logging.getLogger("engine")

try:
    from xgboost import XGBClassifier
    HAVE_XGB = True
except ImportError:  # graceful degrade to logistic-style heuristic
    HAVE_XGB = False

FEE_BUFFER_PCT = 0.5      # round-trip fees + slippage threshold, in %
FORECAST_CANDLES = 4

FEATURES = [
    "rsi", "macd_hist", "atr_pct", "vwap_dist", "bb_pos", "bb_width",
    "ema_cross", "vol_delta", "momentum", "sup_dist", "res_dist",
    "vol_compress", "regime_score",
]


# ------------------------------------------------------------------ features
def build_features(df: pd.DataFrame) -> pd.DataFrame:
    out = df.copy()
    c, h, l, v = out["Close"], out["High"], out["Low"], out["Volume"]

    # RSI(14)
    delta = c.diff()
    gain = delta.clip(lower=0).rolling(14).mean()
    loss = (-delta.clip(upper=0)).rolling(14).mean()
    out["rsi"] = 100 - 100 / (1 + gain / loss.replace(0, np.nan))

    # MACD histogram
    ema12, ema26 = c.ewm(span=12).mean(), c.ewm(span=26).mean()
    macd = ema12 - ema26
    out["macd_hist"] = (macd - macd.ewm(span=9).mean()) / c * 100

    # ATR% (14)
    tr = pd.concat(
        [h - l, (h - c.shift()).abs(), (l - c.shift()).abs()], axis=1
    ).max(axis=1)
    out["atr_pct"] = tr.rolling(14).mean() / c * 100

    # VWAP distance (rolling 96 candles)
    tp = (h + l + c) / 3
    vwap = (tp * v).rolling(96).sum() / v.rolling(96).sum()
    out["vwap"] = vwap
    out["vwap_dist"] = (c - vwap) / vwap * 100

    # Bollinger (20, 2)
    mid = c.rolling(20).mean()
    sd = c.rolling(20).std()
    out["bb_pos"] = (c - mid) / (2 * sd)
    out["bb_width"] = 4 * sd / mid * 100

    # EMA crossover state (9 vs 21)
    out["ema_cross"] = np.sign(c.ewm(span=9).mean() - c.ewm(span=21).mean())

    # volume delta: up-candle vs down-candle volume balance (12 candles)
    up_v = v.where(c > c.shift(), 0.0).rolling(12).sum()
    dn_v = v.where(c <= c.shift(), 0.0).rolling(12).sum()
    out["vol_delta"] = (up_v - dn_v) / (up_v + dn_v).replace(0, np.nan)

    # candle momentum (3-candle return %)
    out["momentum"] = c.pct_change(3) * 100

    # support / resistance distance (rolling 48-candle extremes)
    out["sup_dist"] = (c - l.rolling(48).min()) / c * 100
    out["res_dist"] = (h.rolling(48).max() - c) / c * 100

    # volatility compression: current bb_width vs its 96-candle mean
    out["vol_compress"] = out["bb_width"] / out["bb_width"].rolling(96).mean()

    # regime score: trend + flow composite in [-100, 100]
    trend = np.sign(c.ewm(span=21).mean().diff(8))
    out["regime_score"] = (
        40 * trend + 30 * out["ema_cross"] + 30 * out["vol_delta"].fillna(0)
    )

    return out


def classify_regime(row) -> str:
    s = row["regime_score"]
    if s >= 40:
        return "MARKUP"
    if s >= 10:
        return "ACCUMULATION"
    if s <= -40:
        return "MARKDOWN"
    if s <= -10:
        return "DISTRIBUTION"
    return "RANGE"


# ------------------------------------------------------------------ model
def _label(df: pd.DataFrame) -> pd.Series:
    fwd = df["Close"].shift(-FORECAST_CANDLES) / df["Close"] - 1
    return (fwd * 100 > FEE_BUFFER_PCT).astype(int)


def train_model(feat_df: pd.DataFrame):
    """Train XGBoost on this symbol's candle history. Returns (model, acc)."""
    df = feat_df.dropna(subset=FEATURES).copy()
    df["y"] = _label(df)
    df = df.iloc[:-FORECAST_CANDLES]
    if len(df) < 150:
        return None, 0.0
    split = int(len(df) * 0.8)
    Xtr, ytr = df[FEATURES].iloc[:split], df["y"].iloc[:split]
    Xte, yte = df[FEATURES].iloc[split:], df["y"].iloc[split:]
    if not HAVE_XGB:
        return None, 0.0
    model = XGBClassifier(
        n_estimators=200, max_depth=4, learning_rate=0.08,
        subsample=0.9, colsample_bytree=0.8, eval_metric="logloss",
        n_jobs=2,
    )
    model.fit(Xtr, ytr)
    acc = float((model.predict(Xte) == yte).mean()) if len(yte) else 0.0
    return model, acc


# ------------------------------------------------------------------ analysis
def _nz(x, default=0.0):
    try:
        x = float(x)
        return x if np.isfinite(x) else default
    except Exception:
        return default


def analyze(df: pd.DataFrame, model=None, model_acc: float = 0.0) -> dict:
    feat = build_features(df)
    row = feat.iloc[-1]
    price = float(row["Close"])

    # probability_up
    if model is not None:
        X = feat[FEATURES].iloc[[-1]].fillna(0)
        prob_up = min(99.0, max(1.0, float(model.predict_proba(X)[0, 1]) * 100))
    else:  # heuristic fallback if xgboost missing
        prob_up = 50 + 25 * float(np.tanh(row["regime_score"] / 60))

    atr = _nz(row["atr_pct"], 1.0) or 1.0
    expected_return = (prob_up / 100 - 0.5) * 2 * atr * FORECAST_CANDLES ** 0.5
    downside_risk = atr * 1.65  # ~95% one-candle stress
    confidence = max(0.0, min(100.0, 100 * abs(prob_up - 50) / 50 * 0.6
                              + model_acc * 100 * 0.4))
    regime = classify_regime(row)

    # 4-candle forecast path (drift from expected_return, ATR cone)
    drift = expected_return / FORECAST_CANDLES / 100
    path = [price * (1 + drift) ** (i + 1) for i in range(FORECAST_CANDLES)]
    upper = [p * (1 + atr / 100 * (i + 1) ** 0.5) for i, p in enumerate(path)]
    lower = [p * (1 - atr / 100 * (i + 1) ** 0.5) for i, p in enumerate(path)]

    # action mapping
    bullish = (
        prob_up >= 60
        and expected_return > FEE_BUFFER_PCT
        and regime not in ("MARKDOWN", "DISTRIBUTION")
        and _nz(row["vol_delta"]) > 0
    )
    bearish = prob_up <= 42 or regime == "MARKDOWN"
    action = "BUY" if bullish else "SELL" if bearish else "HOLD"

    return {
        "current_price": price,
        "action": action,
        "probability_up": round(prob_up, 1),
        "expected_return": round(expected_return, 2),
        "downside_risk": round(downside_risk, 2),
        "confidence": round(confidence, 1),
        "market_regime": regime,
        "model_acc": round(model_acc * 100, 1),
        "rsi": round(_nz(row["rsi"], 50.0), 1),
        "macd_hist": round(_nz(row["macd_hist"]), 3),
        "atr_pct": round(atr, 2),
        "vwap": _nz(row["vwap"], price),
        "vol_delta": round(_nz(row["vol_delta"]), 2),
        "ema_cross": int(_nz(row["ema_cross"])),
        "sup_dist": round(_nz(row["sup_dist"]), 2),
        "res_dist": round(_nz(row["res_dist"]), 2),
        "regime_score": round(_nz(row["regime_score"]), 0),
        "forecast_path": path,
        "forecast_upper": upper,
        "forecast_lower": lower,
        "features": {f: float(np.nan_to_num(row[f])) for f in FEATURES},
    }