ob_ultimate/data.py

"""Data loading, multi-timeframe resampling, and causal cross-timeframe mapping.

One base OHLCV CSV per symbol (``data/{SYMBOL}_{base}.csv``, left-labelled bars) is
resampled up to the three strategy timeframes:

  * HTF (bias + OB identification)   — ``cfg.htf_tf``      (default 15m)
  * M5 (entry refinement)            — ``cfg.refine_tf``   (default 5m)
  * M1 (confirmation)                — ``cfg.confirm_tf``  (default = base)

**Causality.** A left-labelled bar with label ``L`` aggregates ``[L, L+freq)`` and
is only *known* at its close ``L+freq``. ``bar_close_times`` returns those closes;
any consumer that acts on a higher-TF bar must gate on the first finer-TF bar whose
timestamp is ``>=`` that close. No finer bar dated after a coarser bar's decision
time is ever consumed before the coarser bar has closed (see ``map_times_to_index``).
"""
from __future__ import annotations

from pathlib import Path
from typing import List, Optional

import numpy as np
import pandas as pd

from .config import Config, DATA_DIR, UNIVERSE_CSV

# spot cache uses the current instrument name for demerged Tata Motors
ALIASES = {"TATAMOTORS": "TMPV"}

_OHLC = {"open": "first", "high": "max", "low": "min", "close": "last", "volume": "sum"}


def load_base(symbol: str, cfg: Config, data_dir: Path = DATA_DIR) -> pd.DataFrame:
    """Load the raw base-timeframe frame for ``symbol``.

    Returns lowercase OHLCV with a tz-aware DatetimeIndex, session-filtered to
    ``[session_start, session_end)`` and sorted. Looks for ``{SYMBOL}_{base}.csv``
    (e.g. ``SBIN_1min.csv`` or ``SBIN_5min.csv``), applying the Tata alias.
    """
    tag = cfg.base_tf.replace("min", "min")
    name = ALIASES.get(symbol, symbol)
    candidates = [data_dir / f"{name}_{tag}.csv", data_dir / f"{symbol}_{tag}.csv"]
    path = next((p for p in candidates if p.exists()), None)
    if path is None:
        raise FileNotFoundError(
            f"no base CSV for {symbol} at {cfg.base_tf} (looked for {candidates})")
    df = pd.read_csv(path, parse_dates=["timestamp"]).set_index("timestamp")
    df = df[["open", "high", "low", "close", "volume"]].sort_index()
    t = df.index.time
    df = df[(t >= cfg.session_start) & (t < cfg.session_end)]
    return df


def resample(base_df: pd.DataFrame, tf: str) -> pd.DataFrame:
    """Aggregate the base frame to ``tf`` (left-labelled), never straddling the
    overnight gap (resampled per trading day)."""
    if tf == "" or _same_freq(base_df, tf):
        return base_df.copy()
    days = base_df.index.normalize()
    parts: List[pd.DataFrame] = []
    for _, day_df in base_df.groupby(days):
        r = day_df.resample(tf, label="left", closed="left").agg(_OHLC)
        parts.append(r.dropna(subset=["open"]))
    out = pd.concat(parts).sort_index()
    return out[out["volume"] >= 0]


def _same_freq(base_df: pd.DataFrame, tf: str) -> bool:
    """True when ``tf`` matches the native base spacing (resampling is a no-op)."""
    if len(base_df) < 3:
        return False
    step = pd.Series(base_df.index).diff().dropna().mode()
    if step.empty:
        return False
    return step.iloc[0] == pd.tseries.frequencies.to_offset(tf)


def bar_close_times(df: pd.DataFrame, tf: str) -> pd.Series:
    """Close timestamp of each bar = the instant its information becomes known."""
    freq = pd.tseries.frequencies.to_offset(tf)
    return pd.Series(df.index + freq, index=df.index)


def map_times_to_index(times, target_index: pd.DatetimeIndex) -> List[int]:
    """First ``target_index`` position with timestamp ``>=`` each time; -1 if none.

    The causal join primitive: given decision times on one frame, find the earliest
    bar on another frame that is allowed to act on them.
    """
    ts = pd.DatetimeIndex(pd.to_datetime(list(times)))
    if target_index.tz is not None and ts.tz is None:
        ts = ts.tz_localize(target_index.tz)
    elif target_index.tz is not None and ts.tz is not None:
        ts = ts.tz_convert(target_index.tz)
    pos = target_index.searchsorted(ts, side="left")
    n = len(target_index)
    return [int(p) if p < n else -1 for p in pos]


def available_symbols(cfg: Config, only_universe: bool = True,
                      data_dir: Path = DATA_DIR) -> List[str]:
    tag = cfg.base_tf
    files = sorted(p.stem.replace(f"_{tag}", "") for p in data_dir.glob(f"*_{tag}.csv")
                   if "synthetic" not in p.stem)
    if not only_universe or not UNIVERSE_CSV.exists():
        return files
    uni = set(pd.read_csv(UNIVERSE_CSV)["symbol"].astype(str))
    out = []
    for s in files:
        if s in uni or ALIASES.get(s) in uni or _rev_alias(s) in uni:
            out.append(s)
    return out


def _rev_alias(name: str) -> Optional[str]:
    for k, v in ALIASES.items():
        if v == name:
            return k
    return None