execution/risk.py

"""
execution/risk.py — Position sizing, drawdown checks, daily limits, circuit breakers.
All risk checks run in strict order before any order submission.
"""
from __future__ import annotations

import datetime
import logging
import threading
from zoneinfo import ZoneInfo

_ET = ZoneInfo("America/New_York")

import config
from contracts import OrderRequest, FinalScore
from execution import broker
from data import storage
from signals.technical import compute_atr

logger = logging.getLogger("trading_system.risk")

# ── Circuit Breaker State ────────────────────────────────────────────────────

_state_lock = threading.Lock()

daily_realized_pnl: float = 0.0
daily_unrealized_pnl: float = 0.0
daily_open_equity: float = 0.0
consecutive_losses: int = 0
paused_until: datetime.datetime | None = None
half_size_mode: bool = False
profit_locked: bool = False
_last_reset_date: datetime.date | None = None


def reset_daily_state(opening_equity: float):
    global daily_open_equity, daily_realized_pnl, daily_unrealized_pnl
    global consecutive_losses, half_size_mode, paused_until, profit_locked
    
    daily_open_equity = opening_equity
    
    # Load persisted state if restarting mid-day
    today = datetime.datetime.now(_ET).strftime("%Y-%m-%d")
    persisted_date = storage.get_kv("risk_state_date")
    if persisted_date == today:
        # Mid-day restart — restore state
        daily_realized_pnl = float(storage.get_kv("risk_daily_realized") or 0)
        consecutive_losses = int(storage.get_kv("risk_consec_losses") or 0)
        half_size_mode = (storage.get_kv("risk_half_size") == "true")
        profit_locked = (storage.get_kv("risk_profit_locked") == "true")
        logger.info(
            "Restored risk state from DB: realized=$%.2f, consec_losses=%d, half_size=%s",
            daily_realized_pnl, consecutive_losses, half_size_mode,
        )
    else:
        # New day — reset
        daily_realized_pnl = 0.0
        consecutive_losses = 0
        half_size_mode = False
        profit_locked = False
        storage.set_kv("risk_state_date", today)
    
    daily_unrealized_pnl = 0.0
    paused_until = None


def update_pnl(realized: float, unrealized: float):
    global daily_realized_pnl, daily_unrealized_pnl
    with _state_lock:
        daily_realized_pnl = realized
        daily_unrealized_pnl = unrealized
    storage.set_kv("risk_daily_realized", str(realized))


def record_trade_result(profit: float):
    """Record the result of a closed trade for streak tracking."""
    global consecutive_losses, half_size_mode
    with _state_lock:
        if profit < 0:
            consecutive_losses += 1
        else:
            consecutive_losses = 0
            half_size_mode = False
        storage.set_kv("risk_consec_losses", str(consecutive_losses))
        storage.set_kv("risk_half_size", str(half_size_mode).lower())


# ── Position Sizing ──────────────────────────────────────────────────────────

def compute_position_size(
    equity: float,
    atr_value: float,
    symbol: str,
    current_price: float = 0.0,
    allow_overnight: bool = False,
) -> tuple[float, float, bool]:
    """Compute position size based on risk parameters.

    Args:
        equity: Current account equity
        atr_value: Current ATR(14) on 5m bars
        symbol: Stock symbol (for fractionable check)
        current_price: Current price (for notional cap)
        allow_overnight: Whether position might be held overnight

    Returns: (qty, stop_distance, is_fractional)
    """
    # Guard: reject if risk % exceeds hard cap
    risk_pct = min(config.RISK_PER_TRADE_PCT, 3.0)

    dollar_risk = equity * (risk_pct / 100)
    stop_distance = 1.0 * atr_value  # for swing: atr_value already = 2x daily ATR

    if stop_distance <= 0:
        logger.warning("ATR is zero for %s, cannot size position", symbol)
        return 0.0, 0.0, False

    qty = dollar_risk / stop_distance

    # Cap notional value — distribute capital across MAX_OPEN_POSITIONS
    if current_price > 0:
        # Allow up to 1/MAX_OPEN_POSITIONS of equity per position.
        # A configurable floor keeps tiny accounts tradable without forcing
        # oversized positions when the user intentionally runs many slots.
        per_position_pct = max(
            1.0 / max(config.MAX_OPEN_POSITIONS, 1),
            config.MIN_POSITION_NOTIONAL_PCT,
        )
        max_notional = equity * per_position_pct
        max_qty = max_notional / current_price
        if qty > max_qty:
            logger.info(
                "%s: capping qty from %.2f to %.2f (notional cap $%.0f, %.0f%% of equity)",
                symbol, qty, max_qty, max_notional, per_position_pct * 100,
            )
            qty = max_qty

    # Half size after 3 consecutive losses (but NOT for overnight — swing needs full risk)
    if half_size_mode:
        qty *= 0.5

    is_fractional = qty != int(qty) or qty < 1

    if is_fractional and qty < 1:
        # Check if symbol supports fractional shares
        try:
            if not broker.is_fractionable(symbol):
                logger.warning(
                    "qty too small and symbol not fractionable: %s (qty=%.4f)",
                    symbol, qty,
                )
                return 0.0, stop_distance, False
        except Exception as e:
            logger.error("Failed to check fractionable for %s: %s", symbol, e)
            return 0.0, stop_distance, False

    return qty, stop_distance, is_fractional


# ── Risk Check Pipeline ─────────────────────────────────────────────────────

def run_risk_checks(
    symbol: str,
    final_score: FinalScore,
    open_position_count: int,
    open_positions: dict[str, dict],
    alert_callback=None,
) -> tuple[bool, str | None]:
    """Run all risk checks in order. Returns (allowed, skip_reason).

    Check order:
      1. SAFE MODE
      2. Daily loss circuit breaker
      3. Equity drop 5% circuit breaker
      4. 3 consecutive losses → pause
      5. Daily profit lock
      6. Max open positions
      7. PDT check (delegated to pdt_tracker)
      8. Overnight time check
    """
    now = datetime.datetime.now(_ET)  # ET

    # 1. SAFE MODE
    if broker.safe_mode_active:
        return False, "safe_mode_active"

    # 1b. Short selling guard — independent of signal filter
    if not config.ALLOW_SHORT and final_score.decision == "SELL":
        # Check if we actually hold this symbol
        if symbol not in open_positions or not open_positions.get(symbol):
            return False, "short_selling_blocked"

    # 2. Daily loss circuit breaker
    total_pnl = daily_realized_pnl + daily_unrealized_pnl
    if total_pnl <= -config.MAX_DAILY_LOSS_USD:
        msg = f"Daily loss limit hit: ${total_pnl:.2f} <= -${config.MAX_DAILY_LOSS_USD}"
        logger.critical(msg)
        if alert_callback:
            alert_callback(f"🛑 CIRCUIT BREAKER: {msg}")
        try:
            broker.cancel_all_orders()
            broker.close_all_positions()
        except Exception as e:
            logger.error("Failed to close on circuit breaker: %s", e)
        return False, "daily_loss_limit"

    # 3. Equity drop 5%
    if daily_open_equity > 0:
        current_equity = daily_open_equity + total_pnl
        drop_pct = (daily_open_equity - current_equity) / daily_open_equity * 100
        if drop_pct >= 5.0:
            msg = f"Equity dropped {drop_pct:.1f}% from daily open"
            logger.critical(msg)
            if alert_callback:
                alert_callback(f"🛑 CIRCUIT BREAKER: {msg}")
            try:
                broker.cancel_all_orders()
                broker.close_all_positions()
            except Exception as e:
                logger.error("Failed to close on equity drop: %s", e)
            return False, "equity_drop_5pct"

    # 4. Consecutive losses → 30 min pause at 50% size
    global paused_until, half_size_mode
    if consecutive_losses >= 3:
        if paused_until is None:
            paused_until = datetime.datetime.now(
                datetime.timezone.utc
            ) + datetime.timedelta(minutes=30)
            half_size_mode = True
            msg = f"3 consecutive losses. Pausing 30 min, resuming at 50% size."
            logger.warning(msg)
            if alert_callback:
                alert_callback(f"⚠️ {msg}")
        if datetime.datetime.now(datetime.timezone.utc) < paused_until:
            return False, "consecutive_loss_pause"

    # 5. Daily profit lock
    if (config.LOCK_PROFITS_AT_TARGET and
            daily_realized_pnl >= config.TARGET_DAILY_PROFIT_USD):
        global profit_locked
        profit_locked = True
        return False, "profit_target_locked"

    # 6. Max open positions
    if open_position_count >= config.MAX_OPEN_POSITIONS:
        return False, "max_positions_reached"

    # 6b. Max daily trades
    today_str = now.strftime("%Y-%m-%d")
    daily_trade_count = storage.count_trades_today(today_str)
    if daily_trade_count >= config.MAX_DAILY_TRADES:
        return False, f"max_daily_trades_reached ({daily_trade_count}/{config.MAX_DAILY_TRADES})"

    # 7. PDT check
    from execution.pdt_tracker import can_day_trade
    side = "buy" if final_score.decision == "BUY" else "sell"
    pdt_ok, pdt_reason = can_day_trade(symbol, side, open_positions, alert_callback)
    if not pdt_ok:
        return False, pdt_reason

    # 8. Overnight check
    if not config.ALLOW_OVERNIGHT_POSITIONS:
        hour, minute = now.hour, now.minute
        if hour >= 15 and minute >= 45:
            return False, "after_1545_no_new_orders"

    return True, None


# ── Per-Position Loss Guardian ───────────────────────────────────────────────

def check_position_loss_limits(
    portfolio,
    alert_callback=None,
) -> list[str]:
    """Close any position whose unrealized loss exceeds MAX_LOSS_PER_POSITION_USD.

    Returns list of symbols that were force-closed.
    """
    max_loss = config.MAX_LOSS_PER_POSITION_USD
    closed = []

    for sym, pos in list(portfolio.positions.items()):
        unrealized = pos.get("unrealized_pl", 0)
        if unrealized <= -max_loss:
            logger.warning(
                "🛑 %s unrealized P&L $%.2f exceeds -$%.2f limit → closing",
                sym, unrealized, max_loss,
            )
            try:
                if config.DRY_RUN:
                    from execution.order_manager import close_virtual_position
                    close_virtual_position(sym)
                else:
                    broker.close_position(sym)
                closed.append(sym)
                record_trade_result(unrealized)
                portfolio.record_close(
                    symbol=sym,
                    realized_pnl=unrealized,
                    entry_price=pos.get("entry_price", 0),
                    exit_price=pos.get("current_price", 0),
                    qty=pos.get("qty", 0),
                )
                if alert_callback:
                    alert_callback(
                        f"🛑 LOSS CUT: {sym} closed at ${unrealized:.2f} "
                        f"(limit: -${max_loss:.2f})"
                    )
            except Exception as e:
                logger.error("Failed to close losing position %s: %s", sym, e)

    return closed