ob_ultimate/engine.py

"""Deterministic base-frame bar replay with the spec's scaled-exit management.

Fill & exit rules (spec §Entry, §Stop, §Targets & Position Management):

* A resting limit at the M5 OB fills no earlier than the bar in which its plan is
  known (``created_time``); cancelled on expiry, session end, or if the OB distal
  edge is crossed before filling (``invalidate_before_fill``).
* On fill the position is split: **50% books at the fixed 5R first target**; the
  remaining **runner** targets the structural secondary level.
* **Break-even** for the runner follows ``breakeven_mode``: ``never`` keeps the
  original stop; ``after_first_target`` moves to BE once the first tranche books;
  ``after_further_bos`` (the faithful rule) moves to BE only once a *further* BOS in
  the trade direction prints — "no reason for price to return".
* Same-bar stop/target collisions resolve stop-first (``same_bar_policy``); MAE/MFE
  track the whole trade. Consumed OB zones are not re-used (``first_touch_only``).

One event-log row per plan (filled / unfilled / invalidated / rejected) — nothing is
silently dropped. Option execution is mapped afterward on a size-weighted exit spot.
"""
from __future__ import annotations

from math import floor
from typing import Dict, List, Optional

import numpy as np
import pandas as pd

from .config import Config
from .confirm import Plan


class _Trade:
    __slots__ = ("plan", "start", "expiry_pos", "entry_idx", "entry_price",
                 "units_first", "units_runner", "first_done", "runner_stop",
                 "be_moved", "mae", "mfe", "realized_pnl", "realized_units",
                 "exit_spot_num", "first_exit", "runner_exit", "reasons")

    def __init__(self, plan: Plan, start: int, expiry_pos: int):
        self.plan = plan
        self.start = start
        self.expiry_pos = expiry_pos
        self.entry_idx: Optional[int] = None
        self.entry_price = float("nan")
        self.units_first = 0
        self.units_runner = 0
        self.first_done = False
        self.runner_stop = plan.stop_price
        self.be_moved = False
        self.mae = 0.0
        self.mfe = 0.0
        self.realized_pnl = 0.0
        self.realized_units = 0
        self.exit_spot_num = 0.0            # sum(units*exit_price) for weighted exit
        self.first_exit: Optional[tuple] = None    # (idx, price, reason)
        self.runner_exit: Optional[tuple] = None
        self.reasons: List[str] = []


def _pos(index: pd.DatetimeIndex, ts) -> int:
    return int(index.searchsorted(ts, side="left"))


def run(plans: List[Plan], base_df: pd.DataFrame, cfg: Config,
        option_map=None) -> List[dict]:
    index = base_df.index
    high = base_df["high"].to_numpy(dtype=float)
    low = base_df["low"].to_numpy(dtype=float)
    close = base_df["close"].to_numpy(dtype=float)
    times = index.time
    n = len(base_df)

    planned = [p for p in plans if p.planned]
    rejected = [p for p in plans if not p.planned]

    working: List[_Trade] = []
    by_start: Dict[int, List[_Trade]] = {}
    for p in sorted(planned, key=lambda x: (x.created_time, x.entry_price)):
        start = _pos(index, p.created_time)
        if start >= n:
            rejected.append(_as_unfilled(p, "no_base_bar"))
            continue
        expiry_pos = max(min(_pos(index, p.expiry_time), n - 1), start)
        by_start.setdefault(start, []).append(_Trade(p, start, expiry_pos))

    consumed: set = set()
    open_trades: List[_Trade] = []
    rows: List[dict] = []

    for t in range(n):
        # ---- 1) manage open trades ----
        still: List[_Trade] = []
        for tr in open_trades:
            _excursion(tr, high[t], low[t])
            _maybe_breakeven(tr, index[t], cfg)
            done = _manage_exit(tr, t, high[t], low[t], close[t], times[t], cfg)
            if done:
                rows.append(_row(tr, base_df, cfg, option_map))
            else:
                still.append(tr)
        open_trades = still

        # ---- 2) activate newly-tradeable orders ----
        for tr in by_start.get(t, []):
            working.append(tr)

        # ---- 3) work resting limits ----
        nxt: List[_Trade] = []
        for tr in working:
            p = tr.plan
            if t > tr.expiry_pos:
                rows.append(_unfilled(tr, t, "unfilled_expired", base_df)); continue
            if cfg.first_touch_only and p.zone_key in consumed:
                rows.append(_unfilled(tr, t, "zone_consumed", base_df)); continue
            if cfg.intraday_only and times[t] >= cfg.intraday_exit_time:
                rows.append(_unfilled(tr, t, "session_end", base_df)); continue
            # first target reached before entry -> stale, cancel
            if _tgt_pre_entry(p, high[t], low[t]):
                rows.append(_unfilled(tr, t, "target_pre_entry", base_df)); continue
            # fill?
            if low[t] <= p.entry_price <= high[t]:
                if len(open_trades) >= cfg.max_concurrent:
                    rows.append(_unfilled(tr, t, "skipped_concurrency", base_df)); continue
                _fill(tr, t, cfg)
                consumed.add(p.zone_key)
                _excursion(tr, high[t], low[t])
                if _manage_exit(tr, t, high[t], low[t], close[t], times[t], cfg):
                    rows.append(_row(tr, base_df, cfg, option_map))
                else:
                    open_trades.append(tr)
                continue
            if cfg.invalidate_before_fill and _invalidated(p, high[t], low[t]):
                rows.append(_unfilled(tr, t, "invalidated", base_df)); continue
            nxt.append(tr)
        working = nxt

    # flush leftovers at series end
    for tr in open_trades:
        _force_close(tr, n - 1, close[-1], "eod_series")
        rows.append(_row(tr, base_df, cfg, option_map))
    for tr in working:
        rows.append(_unfilled(tr, min(tr.expiry_pos, n - 1), "unfilled_expired", base_df))
    for p in rejected:
        rows.append(_reject_row(p))
    return rows


# ---- position lifecycle ----------------------------------------------------
def _fill(tr: _Trade, t: int, cfg: Config):
    p = tr.plan
    tr.entry_idx = t
    tr.entry_price = p.entry_price
    total = _size(p, cfg)
    tr.units_first = int(round(total * p.first_scale_fraction))
    tr.units_runner = total - tr.units_first
    tr.runner_stop = p.stop_price


def _size(p: Plan, cfg: Config) -> int:
    if p.risk <= 0:
        return 0
    units = floor(cfg.risk_pct * cfg.start_equity / p.risk)
    if cfg.position_cap_pct is not None and p.entry_price > 0:
        cap = floor(cfg.position_cap_pct * cfg.start_equity / p.entry_price)
        units = min(units, cap)
    return max(units, 2)


def _sign(p: Plan) -> float:
    return 1.0 if p.direction == "long" else -1.0


def _excursion(tr: _Trade, hi, lo):
    if tr.entry_idx is None:
        return
    if tr.plan.direction == "long":
        tr.mfe = max(tr.mfe, hi - tr.entry_price)
        tr.mae = min(tr.mae, lo - tr.entry_price)
    else:
        tr.mfe = max(tr.mfe, tr.entry_price - lo)
        tr.mae = min(tr.mae, tr.entry_price - hi)


def _maybe_breakeven(tr: _Trade, ts, cfg: Config):
    if tr.be_moved or cfg.breakeven_mode == "never":
        return
    p = tr.plan
    move = False
    if cfg.breakeven_mode == "after_first_target":
        move = tr.first_done
    elif cfg.breakeven_mode == "after_further_bos":
        move = any(ts >= bt for bt in p.further_bos_times)
    if move:
        sign = _sign(p)
        tr.runner_stop = tr.entry_price + sign * p.breakeven_buffer
        tr.be_moved = True


def _book(tr: _Trade, units: int, price: float, idx: int, reason: str):
    p = tr.plan
    sign = _sign(p)
    tr.realized_pnl += sign * (price - tr.entry_price) * units
    tr.realized_units += units
    tr.exit_spot_num += units * price
    tr.reasons.append(reason)


def _manage_exit(tr: _Trade, t, hi, lo, cl, tm, cfg: Config) -> bool:
    """Advance exits on bar ``t``; return True when the whole position is closed."""
    p = tr.plan
    long = p.direction == "long"

    # ---- phase A: first tranche + runner share the original stop ----
    if not tr.first_done and tr.units_first > 0:
        hit_stop = (lo <= p.stop_price) if long else (hi >= p.stop_price)
        hit_first = (hi >= p.first_target_price) if long else (lo <= p.first_target_price)
        if hit_stop and hit_first:
            if cfg.same_bar_policy == "stop_first":
                _book(tr, tr.units_first + tr.units_runner, p.stop_price, t, "stop")
                tr.first_exit = (t, p.stop_price, "stop")
                tr.runner_exit = (t, p.stop_price, "stop")
                _force_close(tr, t, p.stop_price, "stop", already=True)
                return True
            # target-first
            _book(tr, tr.units_first, p.first_target_price, t, "first_target_5R")
            tr.first_exit = (t, p.first_target_price, "first_target_5R")
            tr.first_done = True
            _apply_be_on_first(tr, cfg)
        elif hit_stop:
            _book(tr, tr.units_first + tr.units_runner, p.stop_price, t, "stop")
            tr.first_exit = (t, p.stop_price, "stop")
            tr.runner_exit = (t, p.stop_price, "stop")
            return True
        elif hit_first:
            _book(tr, tr.units_first, p.first_target_price, t, "first_target_5R")
            tr.first_exit = (t, p.first_target_price, "first_target_5R")
            tr.first_done = True
            _apply_be_on_first(tr, cfg)

    # ---- phase B: runner only ----
    if tr.first_done and tr.units_runner > 0 and tr.runner_exit is None:
        rstop = tr.runner_stop
        hit_stop = (lo <= rstop) if long else (hi >= rstop)
        hit_sec = (hi >= p.secondary_target_price) if long else (lo <= p.secondary_target_price)
        if hit_stop and hit_sec:
            px = rstop if cfg.same_bar_policy == "stop_first" else p.secondary_target_price
            reason = "runner_stop" if cfg.same_bar_policy == "stop_first" else "secondary_target"
            _book(tr, tr.units_runner, px, t, reason)
            tr.runner_exit = (t, px, reason)
            return True
        if hit_stop:
            _book(tr, tr.units_runner, rstop, t, "runner_stop")
            tr.runner_exit = (t, rstop, "runner_stop")
            return True
        if hit_sec:
            _book(tr, tr.units_runner, p.secondary_target_price, t, "secondary_target")
            tr.runner_exit = (t, p.secondary_target_price, "secondary_target")
            return True

    # ---- intraday / session force-close of anything still open ----
    if cfg.intraday_only and tm >= cfg.intraday_exit_time:
        _force_close(tr, t, cl, "intraday_exit")
        return True

    # fully closed?
    if tr.realized_units >= tr.units_first + tr.units_runner and tr.units_first + tr.units_runner > 0:
        return True
    return False


def _apply_be_on_first(tr: _Trade, cfg: Config):
    if cfg.breakeven_mode == "after_first_target" and not tr.be_moved:
        sign = _sign(tr.plan)
        tr.runner_stop = tr.entry_price + sign * tr.plan.breakeven_buffer
        tr.be_moved = True


def _force_close(tr: _Trade, t, price, reason, already: bool = False):
    if already:
        return
    if not tr.first_done and tr.units_first > 0 and tr.first_exit is None:
        _book(tr, tr.units_first, price, t, reason)
        tr.first_exit = (t, price, reason)
        tr.first_done = True
    if tr.units_runner > 0 and tr.runner_exit is None:
        _book(tr, tr.units_runner, price, t, reason)
        tr.runner_exit = (t, price, reason)


def _tgt_pre_entry(p: Plan, hi, lo) -> bool:
    return (hi >= p.first_target_price) if p.direction == "long" else (lo <= p.first_target_price)


def _invalidated(p: Plan, hi, lo) -> bool:
    return (lo < p.invalidation_price) if p.direction == "long" else (hi > p.invalidation_price)


# ---- event-log rows --------------------------------------------------------
def _base_row(p: Plan) -> dict:
    return {
        "symbol": p.symbol, "direction": p.direction, "trend": p.trend,
        "htf_break_kind": p.htf_break_kind, "qualified_by": p.qualified_by,
        "swept": p.swept, "sweep_level_type": p.sweep_level_type,
        "sweep_level_price": p.sweep_level_price, "sweep_time": p.sweep_time,
        "htf_known_time": p.htf_known_time, "htf_ob_low": p.htf_ob_low,
        "htf_ob_high": p.htf_ob_high, "htf_ob_time": p.htf_ob_time,
        "htf_touch_time": p.htf_touch_time, "m1_shift_time": p.m1_shift_time,
        "m1_shift_kind": p.m1_shift_kind, "m5_ob_low": p.m5_ob_low,
        "m5_ob_high": p.m5_ob_high, "m5_ob_time": p.m5_ob_time,
        "entry_price": p.entry_price, "stop_price": p.stop_price,
        "first_target_price": p.first_target_price,
        "secondary_target_price": p.secondary_target_price,
        "secondary_target_mode": p.secondary_target_mode,
        "planned_first_R": p.planned_first_R, "planned_secondary_R": p.planned_secondary_R,
        "created_time": p.created_time,
        "entry_timestamp": None, "first_exit_time": None, "first_exit_price": np.nan,
        "runner_exit_time": None, "runner_exit_price": np.nan, "exit_reason": None,
        "units_total": np.nan, "gross_pnl": np.nan, "costs": np.nan, "net_pnl": np.nan,
        "realized_R": np.nan, "mae": np.nan, "mfe": np.nan, "be_moved": False,
        "rejection_reason": p.rejection_reason,
        "option_symbol": None, "option_entry_price": np.nan, "option_exit_price": np.nan,
        "option_gross_pnl": np.nan, "option_costs": np.nan, "option_net_pnl": np.nan,
    }


def _row(tr: _Trade, base_df: pd.DataFrame, cfg: Config, option_map) -> dict:
    p = tr.plan
    r = _base_row(p)
    idx = base_df.index
    units = tr.units_first + tr.units_runner
    gross = tr.realized_pnl
    wexit = tr.exit_spot_num / tr.realized_units if tr.realized_units else np.nan
    notional = (tr.entry_price + wexit) * units
    costs = notional * (cfg.commission_pct + cfg.slippage_pct)
    risk = p.risk
    exit_time = None
    if tr.runner_exit is not None:
        exit_time = idx[tr.runner_exit[0]]
    elif tr.first_exit is not None:
        exit_time = idx[tr.first_exit[0]]
    r.update(
        entry_timestamp=idx[tr.entry_idx], units_total=units,
        first_exit_time=(idx[tr.first_exit[0]] if tr.first_exit else None),
        first_exit_price=(tr.first_exit[1] if tr.first_exit else np.nan),
        runner_exit_time=(idx[tr.runner_exit[0]] if tr.runner_exit else None),
        runner_exit_price=(tr.runner_exit[1] if tr.runner_exit else np.nan),
        exit_reason="+".join(dict.fromkeys(tr.reasons)) or "open",
        gross_pnl=gross, costs=costs, net_pnl=gross - costs,
        realized_R=(gross / (risk * units)) if (risk > 0 and units) else np.nan,
        mae=tr.mae, mfe=tr.mfe, be_moved=tr.be_moved,
    )
    if option_map is not None and cfg.options_mode and not np.isnan(wexit):
        r.update(option_map(p, idx[tr.entry_idx], exit_time, tr.entry_price, wexit))
    return r


def _unfilled(tr: _Trade, at_idx: int, reason: str, base_df: pd.DataFrame) -> dict:
    r = _base_row(tr.plan)
    r.update(exit_reason=reason, first_exit_time=base_df.index[at_idx])
    return r


def _as_unfilled(p: Plan, reason: str) -> Plan:
    p.meta["_unfilled_reason"] = reason
    return p


def _reject_row(p: Plan) -> dict:
    r = _base_row(p)
    r.update(exit_reason=(p.rejection_reason or p.meta.get("_unfilled_reason", "unfilled")))
    return r