risk.py

"""
risk.py
-------
Risk management gate for the Momentum Flip agent.

Responsibilities:
- Position sizing based on account equity and risk %.
- Daily loss cap enforcement (halts trading if breached).
- Maximum open position count enforcement.
- Stop-loss and take-profit price validation.
- Pre-trade risk check that must pass before any order is placed.
"""

import logging
from dataclasses import dataclass, field
from datetime import date
from typing import Optional

from config import Config
from strategy_flip import Direction, SignalResult

logger = logging.getLogger(__name__)


@dataclass
class RiskCheckResult:
    approved: bool
    reason: str
    position_size_usd: float = 0.0
    position_size_contracts: float = 0.0
    adjusted_sl: float = 0.0
    adjusted_tp: float = 0.0


@dataclass
class DailyStats:
    date: date = field(default_factory=date.today)
    realized_pnl: float = 0.0
    trade_count: int = 0
    win_count: int = 0
    loss_count: int = 0

    def reset_if_new_day(self) -> None:
        today = date.today()
        if self.date != today:
            logger.info(f"[Risk] New trading day — resetting daily stats. Previous PnL: ${self.realized_pnl:.2f}")
            self.date = today
            self.realized_pnl = 0.0
            self.trade_count = 0
            self.win_count = 0
            self.loss_count = 0

    def record_trade(self, pnl: float) -> None:
        self.realized_pnl += pnl
        self.trade_count += 1
        if pnl > 0:
            self.win_count += 1
        else:
            self.loss_count += 1

    @property
    def win_rate(self) -> float:
        if self.trade_count == 0:
            return 0.0
        return self.win_count / self.trade_count


class RiskManager:
    """
    Enforces all risk rules before any order is placed.
    Acts as a hard gate — if check fails, no order is sent.
    """

    def __init__(self, config: Config):
        self.config = config
        self.daily_stats = DailyStats()
        self._open_positions: int = 0
        self._trading_halted: bool = False
        self._halt_reason: str = ""

    # ------------------------------------------------------------------
    # State Management
    # ------------------------------------------------------------------

    def set_open_positions(self, count: int) -> None:
        """Update the current open position count (called by controller)."""
        self._open_positions = count

    def record_trade_result(self, pnl: float) -> None:
        """Record a closed trade's PnL and check daily cap."""
        self.daily_stats.reset_if_new_day()
        self.daily_stats.record_trade(pnl)
        logger.info(
            f"[Risk] Trade recorded: PnL=${pnl:.2f} | "
            f"Daily PnL=${self.daily_stats.realized_pnl:.2f} | "
            f"Trades today={self.daily_stats.trade_count}"
        )
        self._check_daily_cap()

    def _check_daily_cap(self) -> None:
        """Halt trading if daily loss cap is breached."""
        if self.daily_stats.realized_pnl <= -abs(self.config.DAILY_LOSS_CAP_USD):
            self._trading_halted = True
            self._halt_reason = (
                f"Daily loss cap breached: ${self.daily_stats.realized_pnl:.2f} "
                f"<= -${self.config.DAILY_LOSS_CAP_USD}"
            )
            logger.critical(f"[Risk] 🛑 TRADING HALTED — {self._halt_reason}")

    def resume_trading(self) -> None:
        """Manually resume trading after a halt (operator action)."""
        self._trading_halted = False
        self._halt_reason = ""
        logger.warning("[Risk] Trading manually resumed by operator.")

    def is_halted(self) -> bool:
        return self._trading_halted

    # ------------------------------------------------------------------
    # Position Sizing
    # ------------------------------------------------------------------

    def calculate_position_size(
        self,
        account_equity: float,
        entry_price: float,
        stop_loss_price: float,
    ) -> tuple[float, float]:
        """
        Calculate position size using fixed fractional risk.

        Risk amount = equity * RISK_PER_TRADE_PCT / 100
        Position size (USD) = risk_amount / (|entry - sl| / entry)
        Position size (contracts) = position_size_usd / entry_price

        Returns:
            (position_size_usd, position_size_contracts)
        """
        if entry_price <= 0 or stop_loss_price <= 0:
            return 0.0, 0.0

        risk_amount = account_equity * (self.config.RISK_PER_TRADE_PCT / 100)
        sl_distance_pct = abs(entry_price - stop_loss_price) / entry_price

        if sl_distance_pct == 0:
            logger.warning("[Risk] SL distance is zero — cannot size position.")
            return 0.0, 0.0

        position_size_usd = risk_amount / sl_distance_pct

        # Cap at max position size
        position_size_usd = min(position_size_usd, self.config.MAX_POSITION_SIZE_USD)

        # Apply leverage
        position_size_usd_leveraged = position_size_usd * self.config.LEVERAGE
        position_size_usd_leveraged = min(position_size_usd_leveraged, self.config.MAX_POSITION_SIZE_USD)

        position_size_contracts = position_size_usd_leveraged / entry_price

        logger.debug(
            f"[Risk] Sizing: equity=${account_equity:.2f} risk={self.config.RISK_PER_TRADE_PCT}% "
            f"risk_amt=${risk_amount:.2f} sl_dist={sl_distance_pct:.4f} "
            f"size_usd=${position_size_usd_leveraged:.2f} contracts={position_size_contracts:.6f}"
        )

        return round(position_size_usd_leveraged, 2), round(position_size_contracts, 6)

    # ------------------------------------------------------------------
    # Pre-Trade Risk Check
    # ------------------------------------------------------------------

    def check(
        self,
        signal: SignalResult,
        account_equity: float,
    ) -> RiskCheckResult:
        """
        Full pre-trade risk gate. Must return approved=True for order to proceed.

        Checks:
        1. Trading not halted.
        2. Signal is actionable (direction != HOLD, trend aligned, confidence >= threshold).
        3. Open position count below max.
        4. Position size is non-zero.
        5. Daily stats are current.
        """
        self.daily_stats.reset_if_new_day()

        # 1. Halt check
        if self._trading_halted:
            return RiskCheckResult(
                approved=False,
                reason=f"Trading halted: {self._halt_reason}"
            )

        # 2. Signal actionability
        if not signal.is_actionable():
            return RiskCheckResult(
                approved=False,
                reason=f"Signal not actionable: direction={signal.direction.value} "
                       f"trend_aligned={signal.trend_aligned} confidence={signal.confidence}"
            )

        # 3. Open position count
        if self._open_positions >= self.config.MAX_OPEN_POSITIONS:
            return RiskCheckResult(
                approved=False,
                reason=f"Max open positions reached: {self._open_positions}/{self.config.MAX_OPEN_POSITIONS}"
            )

        # 4. Account equity sanity
        if account_equity <= 0:
            return RiskCheckResult(
                approved=False,
                reason=f"Invalid account equity: ${account_equity}"
            )

        # 5. Position sizing
        size_usd, size_contracts = self.calculate_position_size(
            account_equity=account_equity,
            entry_price=signal.entry_price,
            stop_loss_price=signal.stop_loss_price,
        )

        if size_contracts <= 0:
            return RiskCheckResult(
                approved=False,
                reason="Position size calculated as zero — skipping trade."
            )

        logger.info(
            f"[Risk] ✅ Trade approved: {signal.direction.value} | "
            f"Size=${size_usd:.2f} ({size_contracts:.6f} contracts) | "
            f"SL={signal.stop_loss_price} TP={signal.take_profit_price}"
        )

        return RiskCheckResult(
            approved=True,
            reason="All risk checks passed.",
            position_size_usd=size_usd,
            position_size_contracts=size_contracts,
            adjusted_sl=signal.stop_loss_price,
            adjusted_tp=signal.take_profit_price,
        )

    # ------------------------------------------------------------------
    # Status
    # ------------------------------------------------------------------

    def get_status(self) -> dict:
        return {
            "halted": self._trading_halted,
            "halt_reason": self._halt_reason,
            "open_positions": self._open_positions,
            "daily_pnl": round(self.daily_stats.realized_pnl, 2),
            "daily_trades": self.daily_stats.trade_count,
            "win_rate": round(self.daily_stats.win_rate, 3),
            "daily_loss_cap": self.config.DAILY_LOSS_CAP_USD,
            "cap_remaining": round(
                self.config.DAILY_LOSS_CAP_USD + self.daily_stats.realized_pnl, 2
            ),
        }