ai-trade/modules/strategy/risk_gate.py

"""
포트폴리오 리스크 게이트 (v3.2)

매수 체결 직전 호출되어 포트폴리오 레벨 제약을 검증:
  1. 총 보유 종목 수 상한
  2. 테마당 동시 보유 종목 수 상한
  3. 테마당 노출 금액 비율 상한 (총자산 대비)

기존 매수 필터(예수금, 종목당 상한, 사이클당 매수 수)는 유지하고
이 게이트가 "같은 테마에 집중되는 포지션"을 차단한다.

순수 함수로 구현 — 의존성 없음 → 단위 테스트 가능.
"""
from dataclasses import dataclass
from typing import Dict, Iterable, List, Optional


@dataclass
class RiskDecision:
    allowed: bool
    reason: str = ""
    max_allowed_amount: int = 0     # 일부만 허용되는 경우 (테마 노출 상한)


@dataclass
class RiskConfig:
    max_total_holdings: int = 7
    max_tickers_per_theme: int = 2
    max_theme_exposure_ratio: float = 0.40


class PortfolioRiskGate:
    """
    사용 예:
      gate = PortfolioRiskGate(theme_map, RiskConfig())
      decision = gate.evaluate_buy(
          ticker="005930",
          candidate_amount=3_000_000,
          current_holdings=[{"ticker":"000660","eval_amount":2_500_000}, ...],
          total_capital=50_000_000,
      )
      if not decision.allowed: skip
      elif decision.max_allowed_amount < candidate_amount: partial buy
    """

    def __init__(self, theme_lookup, config: Optional[RiskConfig] = None):
        """
        Args:
            theme_lookup: callable(ticker:str) -> list[str]  (종목→테마 매핑 함수)
                          혹은 dict 형태도 허용.
            config: RiskConfig
        """
        if callable(theme_lookup):
            self._theme_of = theme_lookup
        elif isinstance(theme_lookup, dict):
            self._theme_of = lambda t: theme_lookup.get(t, [])
        else:
            raise TypeError("theme_lookup must be callable or dict")
        self.config = config or RiskConfig()

    # ──────────────────────────────────────────────
    # 내부: 테마별 현재 노출 집계
    # ──────────────────────────────────────────────
    def _aggregate_by_theme(self, holdings: Iterable[dict]) -> Dict[str, dict]:
        """
        Returns:
            {theme: {"tickers": set, "amount": int}}
        """
        agg: Dict[str, dict] = {}
        for h in holdings:
            tkr = h.get("ticker")
            amt = int(h.get("eval_amount", 0) or 0)
            if not tkr:
                continue
            themes = self._theme_of(tkr) or []
            for th in themes:
                bucket = agg.setdefault(th, {"tickers": set(), "amount": 0})
                bucket["tickers"].add(tkr)
                bucket["amount"] += amt
        return agg

    # ──────────────────────────────────────────────
    # 공개 API
    # ──────────────────────────────────────────────
    def evaluate_buy(self, ticker: str, candidate_amount: int,
                     current_holdings: List[dict],
                     total_capital: int) -> RiskDecision:
        """
        매수 허가 여부 판단.

        Returns:
            RiskDecision
              - allowed=False: 이유와 함께 차단
              - allowed=True : max_allowed_amount만큼 허용 (candidate_amount 이하)
        """
        if candidate_amount <= 0 or total_capital <= 0:
            return RiskDecision(False, "invalid_amount")

        cfg = self.config

        # 이미 보유 중이면 추가 매수는 이 게이트 대상 아님 (scale-in은 상위에서 처리)
        held_tickers = {h.get("ticker") for h in current_holdings}
        is_new_position = ticker not in held_tickers

        # 1. 총 보유 종목 수 상한
        if is_new_position and len(held_tickers) >= cfg.max_total_holdings:
            return RiskDecision(
                False,
                f"max_total_holdings: {len(held_tickers)}/{cfg.max_total_holdings}"
            )

        themes = self._theme_of(ticker) or []
        if not themes:
            # 테마 정보 없음 → 테마 제약은 건너뛰고 통과
            return RiskDecision(True, "no_theme_info", candidate_amount)

        by_theme = self._aggregate_by_theme(current_holdings)

        allowed_amount = candidate_amount
        blocking_reasons = []

        for th in themes:
            bucket = by_theme.get(th, {"tickers": set(), "amount": 0})

            # 2. 테마당 종목 수 상한 (신규 포지션일 때만)
            if is_new_position and len(bucket["tickers"]) >= cfg.max_tickers_per_theme:
                blocking_reasons.append(
                    f"theme[{th}] tickers {len(bucket['tickers'])}/{cfg.max_tickers_per_theme}"
                )
                continue

            # 3. 테마당 노출 금액 비율 상한
            max_theme_amount = int(total_capital * cfg.max_theme_exposure_ratio)
            remaining = max_theme_amount - bucket["amount"]
            if remaining <= 0:
                blocking_reasons.append(
                    f"theme[{th}] exposure {bucket['amount']:,}/{max_theme_amount:,}"
                )
                continue

            # 테마 잔여액이 candidate보다 작으면 부분 허용
            allowed_amount = min(allowed_amount, remaining)

        if blocking_reasons:
            return RiskDecision(False, "; ".join(blocking_reasons))

        if allowed_amount <= 0:
            return RiskDecision(False, "theme_exposure_full")

        return RiskDecision(True, "ok", allowed_amount)