class InvalidPortfolioError(Exception):
    """Raised when portfolio operation is invalid."""
    pass

class portfolio:
    """Portfolio class tracking positions, cash, equity, and total value."""
    positions: dict[str, int] #key is stock, value is num_shares
    cash: float
    equity: float
    total: float

    def __init__(self, positions: dict, cash: float, time_step: int):
        """Initialize portfolio with starting positions and cash.

        Args:
            positions: Dictionary mapping stock symbols to share counts.
            cash: Starting cash amount.
            time_step: Initial time step (unused, kept for consistency).
        """
        self.positions = positions
        self.cash = cash
        self.equity = 0
        self.total = cash + self.equity

    def set_portfolio(self, time_step: int, output: list, data_obj: DataAPI):
        """Update portfolio with trades and recalculate equity at current prices.

        Args:
            time_step: Current time step for price lookup.
            output: List of trade tuples (stock, shares, flag, price).
            data_obj: DataAPI instance for price retrieval.

        Raises:
            InvalidPortfolioError: If trade would result in invalid portfolio state.
        """

        for trade in output:
            stock = trade[0]
            num_shares = trade[1]
            flag = trade[2]
            stock_price = trade[3]

            transaction_cost = num_shares * stock_price

            if flag == 1: 
                if transaction_cost > self.cash:
                    raise InvalidPortfolioError("buy > cash")
            elif flag == -1:
                current_shares = self.positions.get(stock, 0)
                if num_shares > current_shares: 
                    raise InvalidPortfolioError("sell > shares")

            self.positions[stock] = self.positions.get(stock, 0) + (flag*num_shares) 
            if self.positions[stock] == 0:  #delete stocks from portfolio if no shares
                del self.positions[stock]
            elif self.positions[stock] < 0:
                raise InvalidPortfolioError("Num_shares cannot be < 0")

            self.cash -= flag * transaction_cost

        self.equity = 0 #have to recalculate equity at each time step
        for stock in self.positions:
            stock_price = data_obj.get_price(stock, time_step) 
            self.equity += self.positions[stock] * stock_price
            if self.positions[stock] < 0:
                raise InvalidPortfolioError("Num_shares cannot be < 0")

        if self.equity < 0 or self.cash < 0: 
            raise InvalidPortfolioError("cash or equity < 0")

        self.total = self.cash + self.equity