galaxis-po/backend/app/services/backtest/portfolio.py

"""
Virtual portfolio simulation for backtesting.
"""
from decimal import Decimal
from typing import Dict, List, Optional
from dataclasses import dataclass, field


@dataclass
class Transaction:
    """A single buy/sell transaction."""
    ticker: str
    action: str  # 'buy' or 'sell'
    shares: int
    price: Decimal
    commission: Decimal


@dataclass
class HoldingInfo:
    """Information about a single holding."""
    ticker: str
    name: str
    shares: int
    price: Decimal
    weight: Decimal


class VirtualPortfolio:
    """
    Simulates a portfolio for backtesting.
    Handles cash management, buying/selling with transaction costs.
    """

    def __init__(self, initial_capital: Decimal):
        self.cash: Decimal = initial_capital
        self.holdings: Dict[str, int] = {}  # ticker -> shares
        self.initial_capital = initial_capital

    def get_value(self, prices: Dict[str, Decimal]) -> Decimal:
        """Calculate total portfolio value."""
        holdings_value = sum(
            Decimal(str(shares)) * prices.get(ticker, Decimal("0"))
            for ticker, shares in self.holdings.items()
        )
        return self.cash + holdings_value

    def get_holdings_with_weights(
        self,
        prices: Dict[str, Decimal],
        names: Dict[str, str],
    ) -> List[HoldingInfo]:
        """Get current holdings with weights."""
        total_value = self.get_value(prices)
        if total_value == 0:
            return []

        result = []
        for ticker, shares in self.holdings.items():
            if shares <= 0:
                continue
            price = prices.get(ticker, Decimal("0"))
            value = Decimal(str(shares)) * price
            weight = value / total_value * 100

            result.append(HoldingInfo(
                ticker=ticker,
                name=names.get(ticker, ticker),
                shares=shares,
                price=price,
                weight=weight,
            ))

        # Sort by weight descending
        result.sort(key=lambda x: x.weight, reverse=True)
        return result

    def rebalance(
        self,
        target_tickers: List[str],
        prices: Dict[str, Decimal],
        names: Dict[str, str],
        commission_rate: Decimal,
        slippage_rate: Decimal,
    ) -> List[Transaction]:
        """
        Rebalance portfolio to equal-weight target tickers.
        Returns list of transactions executed.
        """
        transactions: List[Transaction] = []
        target_set = set(target_tickers)

        # Step 1: Sell holdings not in target
        for ticker in list(self.holdings.keys()):
            if ticker not in target_set and self.holdings[ticker] > 0:
                txn = self._sell_all(ticker, prices, commission_rate, slippage_rate)
                if txn:
                    transactions.append(txn)

        # Step 2: Calculate target value per stock (equal weight)
        total_value = self.get_value(prices)
        if len(target_tickers) == 0:
            return transactions

        target_value_per_stock = total_value / Decimal(str(len(target_tickers)))

        # Step 3: Sell excess from current holdings in target
        for ticker in target_tickers:
            if ticker in self.holdings:
                current_shares = self.holdings[ticker]
                price = prices.get(ticker, Decimal("0"))
                if price <= 0:
                    continue

                current_value = Decimal(str(current_shares)) * price

                if current_value > target_value_per_stock * Decimal("1.05"):
                    # Sell excess
                    excess_value = current_value - target_value_per_stock
                    sell_price = price * (1 - slippage_rate)
                    shares_to_sell = int(excess_value / sell_price)

                    if shares_to_sell > 0:
                        txn = self._sell(ticker, shares_to_sell, sell_price, commission_rate)
                        if txn:
                            transactions.append(txn)

        # Step 4: Buy to reach target weight
        for ticker in target_tickers:
            price = prices.get(ticker, Decimal("0"))
            if price <= 0:
                continue

            current_shares = self.holdings.get(ticker, 0)
            current_value = Decimal(str(current_shares)) * price

            if current_value < target_value_per_stock * Decimal("0.95"):
                # Buy more
                buy_value = target_value_per_stock - current_value
                buy_price = price * (1 + slippage_rate)

                # Account for commission in available cash
                max_buy_value = self.cash / (1 + commission_rate)
                actual_buy_value = min(buy_value, max_buy_value)

                shares_to_buy = int(actual_buy_value / buy_price)

                if shares_to_buy > 0:
                    txn = self._buy(ticker, shares_to_buy, buy_price, commission_rate)
                    if txn:
                        transactions.append(txn)

        return transactions

    def _buy(
        self,
        ticker: str,
        shares: int,
        price: Decimal,
        commission_rate: Decimal,
    ) -> Optional[Transaction]:
        """Execute a buy order."""
        cost = Decimal(str(shares)) * price
        commission = cost * commission_rate
        total_cost = cost + commission

        if total_cost > self.cash:
            # Reduce shares to fit budget
            available = self.cash / (1 + commission_rate)
            shares = int(available / price)
            if shares <= 0:
                return None
            cost = Decimal(str(shares)) * price
            commission = cost * commission_rate
            total_cost = cost + commission

        self.cash -= total_cost
        self.holdings[ticker] = self.holdings.get(ticker, 0) + shares

        return Transaction(
            ticker=ticker,
            action='buy',
            shares=shares,
            price=price,
            commission=commission,
        )

    def _sell(
        self,
        ticker: str,
        shares: int,
        price: Decimal,
        commission_rate: Decimal,
    ) -> Optional[Transaction]:
        """Execute a sell order."""
        current_shares = self.holdings.get(ticker, 0)
        shares = min(shares, current_shares)

        if shares <= 0:
            return None

        proceeds = Decimal(str(shares)) * price
        commission = proceeds * commission_rate
        net_proceeds = proceeds - commission

        self.cash += net_proceeds
        self.holdings[ticker] -= shares

        if self.holdings[ticker] <= 0:
            del self.holdings[ticker]

        return Transaction(
            ticker=ticker,
            action='sell',
            shares=shares,
            price=price,
            commission=commission,
        )

    def _sell_all(
        self,
        ticker: str,
        prices: Dict[str, Decimal],
        commission_rate: Decimal,
        slippage_rate: Decimal,
    ) -> Optional[Transaction]:
        """Sell all shares of a ticker."""
        shares = self.holdings.get(ticker, 0)
        if shares <= 0:
            return None

        price = prices.get(ticker, Decimal("0"))
        if price <= 0:
            return None

        sell_price = price * (1 - slippage_rate)
        return self._sell(ticker, shares, sell_price, commission_rate)