Stock-Algorithm-Back-Tester/Account_Information.py

from copy import deepcopy

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from Orders import GenericLimitOrder, BuyOrder, SellOrder, OrderFilledResponse
    from datetime import datetime

class AccountInformation():
    def __init__(self):
        self._open_orders: list['GenericLimitOrder'] = []
        self._filled_orders: list['GenericLimitOrder'] = []
        self._cancelled_orders: list['GenericLimitOrder'] = []

        # For each symbol, have a StockPosition() object 
        self._cumulative_stock_positions: dict[str, CumulativeStockPosition] = {}

        self._total_cash: float = 0.0
        self._settled_cash: float = 0.0
        self._unsettled_cash: float = 0.0
        
        self._needed_cash: float = 0.0

    # General functions
    def __str__(self) -> str:
        return (f"Total Cash: {self._total_cash}\n"
                f"Settled Cash: {self._settled_cash}\n"
                f"Unsettled Cash: {self._unsettled_cash}\n"
                f"Open Orders: {len(self._open_orders)}\n"
                f"Filled Orders: {len(self._filled_orders)}\n"
                f"Cancelled Orders: {len(self._cancelled_orders)}\n")

    # CASH functions ############################################################
    def deposit_cash(self, amount: float) -> float:
        '''Returns the current cash amount after deposit'''
        self._total_cash += amount
        return self._total_cash

    def calc_needed_cash(self) -> None:
        '''Calcs the cash needed to execute the open orders'''
        pass

    @property
    def needed_cash(self) -> float:
        self.calc_needed_cash()
        return self._needed_cash
    @property
    def total_cash(self) -> float:
        return self._total_cash
    @property
    def settled_cash(self) -> float:
        return self._settled_cash
    @property
    def unsettled_cash(self) -> float:
        return self._unsettled_cash

    # ORDER functions #########################################################
    @property
    def open_orders(self) -> list['GenericLimitOrder']:
        return self._open_orders
    
    def add_open_order(self, order: 'GenericLimitOrder') -> None:
        '''Adds the order to the list of open orders'''
        self._open_orders.append(order)

    def fill_open_order(self, order_id: int) -> int:
        '''Moves the order specified by order_id from open_orders to filled_orders
        Returns 0 on success. 1 if not'''
        for order in self._open_orders:
            if order.id == order_id:
                self._filled_orders.append(order)
                self._open_orders.remove(order)

                return 0  # Success
        return 1  # Order not found
    
    def cancel_open_order(self, order_id: int) -> int:
        '''Moves the order specified by order_id from open_orders to cancelled_orders
        Returns 0 on success, 1 if not'''
        for order in self._open_orders:
            if order.id == order_id:
                self._cancelled_orders.append(order)
                self._open_orders.remove(order)
                return 0  # Success
        return 1  # Order not found

    # STOCK POSITION functions ###################################################
    def add_position_from_filled_response(self, response: 'OrderFilledResponse', method: str = "fifo") -> int:
        '''Given a response, determine how it plays into the CumulativeStockPosition.
        First calculate the current position (long or short). If the response is adding to the position,
        just append a position. If it is taking away from the position (making it more neutral), determine which
        previous filled response to take away from.
        
        returns the number of shares closed'''
        symbol = response.symbol
        if symbol not in self._cumulative_stock_positions:
            self._cumulative_stock_positions[symbol] = CumulativeStockPosition(symbol)
        quantity_closed = self._cumulative_stock_positions[symbol].add_position_from_filled_response(response, method)
        return quantity_closed




class CumulativeStockPosition():
    '''Holds the current cumulative position regarding a single stock.
    This will handle which shares are sold (earliest acquired, cheapest, etc)'''
    def __init__(self, symbol: str):
        self._symbol = symbol

        # Keep track of the invidual positions for this stock
            # Only the responses that actively contribute to the cumulative position
        self._single_stock_positions: list[SingleStockPosition] = []

        self._quantity: float = 0.0 # Total number of shares we are long or short
        self._cost_basis: float = 0.0
        self._cost_per_share: float = 0.0

    def calculate_cumulative_position(self) -> None:
        '''Calculate the total quantity, cost basis, and cost per 
        share of the current single positions'''

        self._quantity = sum(position.quantity for position in self._single_stock_positions)
        self._cost_basis = sum(position.quantity * position.cost_per_share for position in self._single_stock_positions)
        self._cost_per_share = self._cost_basis / self._quantity if self._quantity != 0 else 0.0

    def close_position(self, new_position: 'SingleStockPosition', method: str = "fifo") -> float:
        '''fifo = close the first placed position
        lifo = close the more position
        cheapest = close cheapest first,
        priciest = close most expensive first
        max_loss = close the stock to maximize losses
        max_gain = close the stock to maximize gains
        
        returns the number of shares closed'''
        
        quantity_to_close = abs(new_position.quantity)
        quantity_closed = 0

        while quantity_closed < quantity_to_close and self._single_stock_positions:
            if method == "fifo":
                position_to_use = self._single_stock_positions[0]
            elif method == "lifo":
                position_to_use = self._single_stock_positions[-1]
            elif method == "cheapest":
                position_to_use = min(self._single_stock_positions, key=lambda x: x.cost_per_share)
            elif method == "priciest":
                position_to_use = max(self._single_stock_positions, key=lambda x: x.cost_per_share)
            elif method == "max_loss":
                position_to_use = min(self._single_stock_positions, key=lambda x: x.cost_per_share - new_position.cost_per_share)
            elif method == "max_gain":
                position_to_use = max(self._single_stock_positions, key=lambda x: x.cost_per_share - new_position.cost_per_share)
            else:
                raise ValueError("Incorrectly specified method for closing position")

            if quantity_to_close < position_to_use.quantity:
                position_to_use.quantity -= quantity_to_close
                quantity_closed += quantity_to_close
                break
            else:
                quantity_to_close -= position_to_use.quantity
                quantity_closed += position_to_use.quantity
                self._single_stock_positions.remove(position_to_use)

        if quantity_closed < abs(new_position.quantity):
            new_position.quantity = -quantity_to_close
            self._single_stock_positions.append(new_position)

        return quantity_closed
            
    def make_single_stock_position(self, response: 'OrderFilledResponse') -> 'SingleStockPosition':
        single_stock_position = SingleStockPosition(
            symbol = response.symbol,
            matching_order_id = response.matching_order_id,
            datetime_filled = response.datetime_filled,
            quantity = response.quantity,
            cost_per_share = response.cost_per_share
        )
        return single_stock_position
    
    def add_position_from_filled_response(self, response: 'OrderFilledResponse', method: str = "fifo") -> int:
        '''Given a response, determine how it plays into the CumulativeStockPosition.
        First calculate the current position (long or short). If the response is adding to the position,
        just append a position. If it is taking away from the position (making it more neutral), determine which
        previous filled response to take away from.
        
        returns the number of shares closed'''

        # Make a SingleStockPosition from the filled response
        new_position = self.make_single_stock_position(response)

        quantity_closed = 0
        # Determine if it adds to the current position or takes away
        if self._quantity * new_position.quantity >= 0:    # Same sign means the response is adding
            self._single_stock_positions.append(new_position)
        else:                                                       # The response is taking away
            quantity_closed = self.close_position(new_position, method)

        # Recalculate current position
        self.calculate_cumulative_position()

        return quantity_closed

class SingleStockPosition():
    def __init__(self, symbol: str, matching_order_id: int, datetime_filled: 'datetime', quantity: float, cost_per_share: float):
        self.symbol = symbol
        self.matching_order_id = matching_order_id
        self.datetime_filled = datetime_filled
        self.quantity = quantity
        self.cost_per_share = cost_per_share