trex-core/trex_core/domain/capital_allocation.py

"""Pure account-capital allocation policy shared by TREX strategies.

The domain owns only deterministic allocation math. Broker adapters own cash
and quote retrieval, while Portfolio Hub remains the portfolio-ledger owner.
No network, persistence, environment, or strategy dependency is allowed here.
"""
from __future__ import annotations

from dataclasses import dataclass
from decimal import Decimal, ROUND_CEILING


@dataclass(frozen=True)
class CapitalCandidate:
    """One strategy candidate priced at a protected one-share cost."""

    portfolio_id: str
    symbol: str
    protected_share_cost: int


@dataclass(frozen=True)
class CapitalAllocation:
    """One candidate's deterministic allocation decision."""

    portfolio_id: str
    symbol: str
    cap: int
    reason: str

    @property
    def eligible(self) -> bool:
        return self.cap > 0


@dataclass(frozen=True)
class CapitalAllocationPlan:
    """Account-level result that can be consumed by any TREX strategy."""

    available_cash: int
    slot_count: int
    allocations: tuple[CapitalAllocation, ...]

    @property
    def caps(self) -> dict[str, int]:
        return {
            allocation.portfolio_id: allocation.cap
            for allocation in self.allocations
        }


def protected_share_cost(price: Decimal, slippage_pct: Decimal) -> int:
    """Return the ceiling KRW cost of one share including execution margin."""
    if (
        not price.is_finite()
        or not slippage_pct.is_finite()
        or price <= 0
        or slippage_pct < 0
    ):
        raise ValueError(
            "price must be finite and positive and slippage must be finite "
            "and non-negative"
        )
    multiplier = Decimal("1") + slippage_pct / Decimal("100")
    return int(
        (price * multiplier).to_integral_value(rounding=ROUND_CEILING)
    )


def plan_liquid_capital(
    available_cash: int,
    *,
    slot_count: int,
    candidates: tuple[CapitalCandidate, ...],
    flex_share_count: int = 1,
) -> CapitalAllocationPlan:
    """Plan candidate caps whose worst active-slot combination fits cash.

    The starting point is an equal cash budget per slot. A limited amount,
    measured in units of the cheapest protected share, may move between slots
    so that a slightly more expensive candidate remains eligible. A candidate
    that needs more than that margin is excluded.

    After the first screen, the highest protected-share cost is removed until
    every possible active-slot combination can reserve at least one share
    without exceeding real account cash. Remaining candidates receive the
    largest common cap that preserves that invariant.
    """
    _validate_inputs(
        available_cash,
        slot_count=slot_count,
        candidates=candidates,
        flex_share_count=flex_share_count,
    )
    if available_cash == 0:
        return CapitalAllocationPlan(
            available_cash=available_cash,
            slot_count=slot_count,
            allocations=tuple(
                CapitalAllocation(
                    candidate.portfolio_id,
                    candidate.symbol,
                    0,
                    "no_available_cash",
                )
                for candidate in candidates
            ),
        )

    equal_budget = available_cash // slot_count
    flex_unit = min(
        candidate.protected_share_cost for candidate in candidates
    )
    maximum_one_share_cost = (
        equal_budget + flex_unit * flex_share_count
    )
    eligible_indexes = [
        index
        for index, candidate in enumerate(candidates)
        if candidate.protected_share_cost <= maximum_one_share_cost
    ]
    reasons = {
        index: "protected_share_above_flex_limit"
        for index, candidate in enumerate(candidates)
        if candidate.protected_share_cost > maximum_one_share_cost
    }

    while eligible_indexes:
        active_count = min(slot_count, len(eligible_indexes))
        largest_costs = sorted(
            (
                candidates[index].protected_share_cost
                for index in eligible_indexes
            ),
            reverse=True,
        )[:active_count]
        if sum(largest_costs) <= available_cash:
            break
        highest_index = max(
            eligible_indexes,
            key=lambda index: (
                candidates[index].protected_share_cost,
                index,
            ),
        )
        eligible_indexes.remove(highest_index)
        reasons[highest_index] = "worst_case_cash_overcommit"

    if not eligible_indexes:
        return CapitalAllocationPlan(
            available_cash=available_cash,
            slot_count=slot_count,
            allocations=tuple(
                CapitalAllocation(
                    candidate.portfolio_id,
                    candidate.symbol,
                    0,
                    reasons[index],
                )
                for index, candidate in enumerate(candidates)
            ),
        )

    active_count = min(slot_count, len(eligible_indexes))

    def worst_active_total(common_budget: int) -> int:
        caps = sorted(
            (
                max(
                    candidates[index].protected_share_cost,
                    common_budget,
                )
                for index in eligible_indexes
            ),
            reverse=True,
        )
        return sum(caps[:active_count])

    low, high = 0, available_cash
    while low < high:
        middle = (low + high + 1) // 2
        if worst_active_total(middle) <= available_cash:
            low = middle
        else:
            high = middle - 1

    eligible_set = set(eligible_indexes)
    allocations = tuple(
        CapitalAllocation(
            candidate.portfolio_id,
            candidate.symbol,
            (
                max(candidate.protected_share_cost, low)
                if index in eligible_set
                else 0
            ),
            "eligible" if index in eligible_set else reasons[index],
        )
        for index, candidate in enumerate(candidates)
    )
    plan = CapitalAllocationPlan(
        available_cash=available_cash,
        slot_count=slot_count,
        allocations=allocations,
    )
    _assert_plan_invariants(plan)
    return plan


def allocate_liquid_slots(
    available_cash: int,
    *,
    slot_count: int,
    candidates: tuple[CapitalCandidate, ...],
    flex_share_count: int = 1,
) -> dict[str, int]:
    """Compatibility surface returning only ``portfolio_id -> cap``."""
    return plan_liquid_capital(
        available_cash,
        slot_count=slot_count,
        candidates=candidates,
        flex_share_count=flex_share_count,
    ).caps


def _validate_inputs(
    available_cash: int,
    *,
    slot_count: int,
    candidates: tuple[CapitalCandidate, ...],
    flex_share_count: int,
) -> None:
    if available_cash < 0 or slot_count <= 0 or not candidates:
        raise ValueError(
            "cash, positive slot count, and candidates are required"
        )
    if flex_share_count < 0:
        raise ValueError("flex share count must be non-negative")
    portfolio_ids = [
        candidate.portfolio_id for candidate in candidates
    ]
    symbols = [candidate.symbol for candidate in candidates]
    if (
        len(portfolio_ids) != len(set(portfolio_ids))
        or len(symbols) != len(set(symbols))
    ):
        raise ValueError(
            "candidate portfolio ids and symbols must be unique"
        )
    if any(
        candidate.protected_share_cost <= 0
        for candidate in candidates
    ):
        raise ValueError(
            "candidate protected share costs must be positive"
        )


def _assert_plan_invariants(plan: CapitalAllocationPlan) -> None:
    positive = sorted(
        (
            allocation.cap
            for allocation in plan.allocations
            if allocation.cap > 0
        ),
        reverse=True,
    )
    if sum(positive[: plan.slot_count]) > plan.available_cash:
        raise AssertionError("capital plan exceeds available cash")