data_and_features.py

def load_data() -> pd.DataFrame:
    df = yf.download(SYMBOL, start=START, end=END, interval="1d", progress=False, group_by='tickers')[SYMBOL]
    if df.empty:
        raise RuntimeError(f"No data returned for {SYMBOL}. Check symbol/date range.")
    df = df.rename(columns=str.title)
    if "Close" not in df.columns and "Adj Close" in df.columns:
        df["Close"] = df["Adj Close"]
    df = df.dropna(subset=["Close"])
    df.index = pd.to_datetime(df.index)
    return df


def add_features(df: pd.DataFrame) -> pd.DataFrame:
    out = df.copy()

    # log returns
    out["ret"] = np.log(out["Close"]).diff()

    # rolling vol (annualized)
    out["vol_20"] = out["ret"].rolling(20).std() * np.sqrt(252)

    # trend score = rolling mean / rolling std
    rmean = out["ret"].rolling(20).mean()
    rstd = out["ret"].rolling(20).std()
    out["trend_20"] = (rmean / (rstd + 1e-12)).clip(-5, 5)

    # ATR proxy (normalized)
    if {"High", "Low", "Close"}.issubset(out.columns):
        prev_close = out["Close"].shift(1)
        tr = pd.concat(
            [
                out["High"] - out["Low"],
                (out["High"] - prev_close).abs(),
                (out["Low"] - prev_close).abs(),
            ],
            axis=1,
        ).max(axis=1)
        atr_14 = tr.rolling(14).mean()
        out["atr_norm"] = (atr_14 / out["Close"]).clip(0, 1)
    else:
        out["atr_norm"] = out["ret"].abs().rolling(14).mean()

    # z-score for range mean reversion
    ma = out["Close"].rolling(20).mean()
    sd = out["Close"].rolling(20).std()
    out["z_20"] = ((out["Close"] - ma) / (sd + 1e-12)).clip(-6, 6)

    return out.dropna()


def label_dates(index: pd.DatetimeIndex) -> list[pd.Timestamp]:
    return [index[i] for i in range(LOOKBACK_DAYS, len(index), LABEL_STEP_DAYS)]


def window_summary(window_feat: pd.DataFrame) -> dict:
    """
    Tiny numeric summary of the last LOOKBACK_DAYS.
    This is what the LLM sees (not the full time series).
    """
    f = window_feat.dropna()
    eq = f["ret"].fillna(0).cumsum().values
    peak = np.maximum.accumulate(eq)
    max_dd = float((eq - peak).min())

    return {
        "mean_ret": float(f["ret"].mean()),
        "ann_vol": float(f["ret"].std() * math.sqrt(252)),
        "trend_score": float(f["trend_20"].iloc[-1]),
        "atr_norm": float(f["atr_norm"].iloc[-1]),
        "z_last": float(f["z_20"].iloc[-1]),
        "max_dd": max_dd,
    }