{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-15T22:15:07.547650Z",
     "start_time": "2022-08-15T22:15:07.029647Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running typecheck\n",
      "Success: no issues found in 1 source file\n",
      "\n",
      "running flake8\n",
      "stdin:517:17: W503 line break before binary operator\n",
      "stdin:518:17: W503 line break before binary operator\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%%checkall\n",
    "import copy\n",
    "from collections import defaultdict, deque\n",
    "from sortedcontainers import SortedDict\n",
    "import math\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pyqstrat.pq_types import ContractGroup, Trade, Contract, RoundTripTrade\n",
    "from types import SimpleNamespace\n",
    "from typing import Sequence, Any, Tuple, Callable, Union, MutableSet, MutableSequence, MutableMapping, List, Optional, Dict\n",
    "\n",
    "\n",
    "def calc_trade_pnl(open_qtys: np.ndarray, \n",
    "                   open_prices: np.ndarray, \n",
    "                   new_qtys: np.ndarray, \n",
    "                   new_prices: np.ndarray, \n",
    "                   multiplier: float) -> Tuple[np.ndarray, np.ndarray, float, float, float]:\n",
    "    '''\n",
    "    >>> print(calc_trade_pnl(\n",
    "    ...          open_qtys = np.array([], dtype = float), open_prices = np.array([], dtype = float), \n",
    "    ...            new_qtys = np.array([-8, 9, -4]), new_prices = np.array([10, 11, 6]), multiplier = 100))\n",
    "    (array([-3.]), array([6.]), -3.0, 6.0, -1300.0)\n",
    "    >>> print(calc_trade_pnl(open_qtys = np.array([], dtype = float), open_prices = np.array([], dtype = float), new_qtys = np.array([3, 10, -5]), \n",
    "    ...          new_prices = np.array([51, 50, 45]), multiplier = 100))\n",
    "    (array([8.]), array([50.]), 8.0, 50.0, -2800.0)\n",
    "    >>> print(calc_trade_pnl(open_qtys = np.array([]), open_prices = np.array([]), \n",
    "    ...                new_qtys = np.array([-58, -5, -5, 6, -8, 5, 5, -5, 19, 7, 5, -5, 39]),\n",
    "    ...                new_prices = np.array([2080, 2075.25, 2070.75, 2076, 2066.75, 2069.25, 2074.75, 2069.75, 2087.25, 2097.25, 2106, 2088.25, 2085.25]),\n",
    "    ...                multiplier = 50))\n",
    "    (array([], dtype=float64), array([], dtype=float64), 0.0, 0, -33762.5)    '''\n",
    "    # TODO: Cythonize this\n",
    "    \n",
    "    realized = 0.\n",
    "    \n",
    "    new_qtys = new_qtys.copy()\n",
    "    new_prices = new_prices.copy()\n",
    "\n",
    "    _open_prices = np.zeros(len(open_prices) + len(new_prices), dtype=float)\n",
    "    _open_prices[:len(open_prices)] = open_prices\n",
    "    \n",
    "    _open_qtys = np.zeros(len(open_qtys) + len(new_qtys), dtype=float)\n",
    "    _open_qtys[:len(open_qtys)] = open_qtys\n",
    "    \n",
    "    new_qty_indices = np.nonzero(new_qtys)[0]\n",
    "    open_qty_indices = np.zeros(len(_open_qtys), dtype=int)\n",
    "    nonzero_indices = np.nonzero(_open_qtys)[0]\n",
    "    open_qty_indices[:len(nonzero_indices)] = nonzero_indices \n",
    "\n",
    "    i = 0                      # index into new_qty_indices to get idx of the new qty we are currently netting\n",
    "    o = len(nonzero_indices)  # virtual length of open_qty_indices\n",
    "    j = 0                      # index into open_qty_indices to get idx of the open qty we are currently netting\n",
    "    k = len(open_qtys)         # virtual length of _open_qtys\n",
    "    \n",
    "    # Try to net all new trades against existing non-netted trades.\n",
    "    # Append any remaining non-netted new trades to end of existing trades\n",
    "    while i < len(new_qty_indices):\n",
    "        # Always try to net first non-zero new trade against first non-zero existing trade\n",
    "        # FIFO acccounting\n",
    "        new_idx = new_qty_indices[i]\n",
    "        new_qty, new_price = new_qtys[new_idx], new_prices[new_idx]\n",
    "        \n",
    "        # print(f'i: {i} j: {j} k: {k} o: {o} oq: {_open_qtys} oqi: {open_qty_indices} op: {_open_prices} nq: {new_qtys} np: {new_prices}')\n",
    "        \n",
    "        if j < o:  # while we still have open positions to net against\n",
    "            open_idx = open_qty_indices[j]\n",
    "            open_qty, open_price = _open_qtys[open_idx], _open_prices[open_idx]\n",
    "            \n",
    "            if math.copysign(1, open_qty) == math.copysign(1, new_qty):\n",
    "                # Nothing to net against so add this trade to the array and wait for the next offsetting trade\n",
    "                \n",
    "                _open_qtys[k] = new_qty\n",
    "                _open_prices[k] = new_price\n",
    "                open_qty_indices[o] = k\n",
    "                k += 1\n",
    "                o += 1\n",
    "\n",
    "                new_qtys[new_idx] = 0\n",
    "                i += 1\n",
    "\n",
    "            elif abs(new_qty) > abs(open_qty):\n",
    "                # New trade has more qty than offsetting trade so:\n",
    "                # a. net against offsetting trade\n",
    "                # b. remove the offsetting trade\n",
    "                # c. reduce qty of new trade\n",
    "                open_qty, open_price = _open_qtys[open_idx], _open_prices[open_idx]\n",
    "                realized += open_qty * (new_price - open_price)\n",
    "                # print(f'open_qty: {open_qty} open_price: {open_price} open_idx: {open_idx} i: {i}\n",
    "                # j: {j} k: {k} l: {l} oq: {_open_qtys} oqi: {open_qty_indices} op: {_open_prices} nq: {new_qtys} np: {new_prices}')\n",
    "                _open_qtys[open_idx] = 0\n",
    "                j += 1\n",
    "\n",
    "                new_qtys[new_idx] += open_qty\n",
    "            else:\n",
    "                # New trade has less qty than offsetting trade so:\n",
    "                # a. net against offsetting trade\n",
    "                # b. remove new trade\n",
    "                # c. reduce qty of offsetting trade\n",
    "                realized += new_qty * (open_price - new_price)\n",
    "                new_qtys[new_idx] = 0\n",
    "                i += 1\n",
    "                _open_qtys[open_idx] += new_qty\n",
    "        else:\n",
    "            # Nothing to net against so add this trade to the open trades array and wait for the next offsetting trade\n",
    "            _open_qtys[k] = new_qty\n",
    "            _open_prices[k] = new_price\n",
    "            open_qty_indices[o] = k\n",
    "            k += 1\n",
    "            o += 1\n",
    "\n",
    "            new_qtys[new_idx] = 0\n",
    "            i += 1\n",
    "\n",
    "    mask = _open_qtys != 0\n",
    "    _open_qtys = _open_qtys[mask]\n",
    "    _open_prices = _open_prices[mask]\n",
    "    open_qty = np.sum(_open_qtys)\n",
    "    if math.isclose(open_qty, 0):\n",
    "        weighted_avg_price = 0\n",
    "    else:\n",
    "        weighted_avg_price = np.sum(_open_qtys * _open_prices) / open_qty\n",
    "        \n",
    "    return _open_qtys, _open_prices, open_qty, weighted_avg_price, realized * multiplier\n",
    "\n",
    "\n",
    "def leading_nan_to_zero(df: pd.DataFrame, columns: Sequence[str]) -> pd.DataFrame:\n",
    "    for column in columns:\n",
    "        vals = df[column].values\n",
    "        first_non_nan_index_ = np.ravel(np.nonzero(~np.isnan(vals)))  # type: ignore\n",
    "        if len(first_non_nan_index_):\n",
    "            first_non_nan_index = first_non_nan_index_[0]\n",
    "        else:\n",
    "            first_non_nan_index = -1\n",
    "\n",
    "        if first_non_nan_index > 0 and first_non_nan_index < len(vals):\n",
    "            vals[:first_non_nan_index] = np.nan_to_num(vals[:first_non_nan_index])\n",
    "            df[column] = vals\n",
    "    return df\n",
    "\n",
    "\n",
    "def find_last_non_nan_index(array: np.ndarray) -> int:\n",
    "    i = np.nonzero(np.isfinite(array))[0]\n",
    "    if len(i): return i[-1]\n",
    "    return 0\n",
    "\n",
    "\n",
    "def find_index_before(sorted_dict: SortedDict, key: Any) -> int:\n",
    "    '''\n",
    "    Find index of the first key in a sorted dict that is less than or equal to the key passed in.\n",
    "    If the key is less than the first key in the dict, return -1\n",
    "    '''\n",
    "    size = len(sorted_dict)\n",
    "    if not size: return -1\n",
    "    i = sorted_dict.bisect_left(key)\n",
    "    if i == size: return size - 1\n",
    "    if sorted_dict.keys()[i] != key:\n",
    "        return i - 1\n",
    "    return i\n",
    "\n",
    "\n",
    "class ContractPNL:\n",
    "    '''Computes pnl for a single contract over time given trades and market data\n",
    "    >>> from pyqstrat.pq_types import MarketOrder\n",
    "    >>> ContractGroup.clear()\n",
    "    >>> Contract.clear()\n",
    "    >>> aapl_cg = ContractGroup.create('AAPL')\n",
    "    >>> aapl_contract = Contract.create('AAPL', contract_group=aapl_cg)\n",
    "    >>> timestamps = np.arange(np.datetime64('2018-01-01'), np.datetime64('2018-01-04'))\n",
    "    >>> def get_price(contract, timestamps, idx, strategy_context):\n",
    "    ...    assert contract.symbol == 'AAPL', f'unknown contract: {contract}'\n",
    "    ...    return idx + 10.1\n",
    "\n",
    "    >>> contract_pnl = ContractPNL(aapl_contract, timestamps, get_price, SimpleNamespace()) \n",
    "    >>> trade_5 = Trade(aapl_contract, MarketOrder(aapl_contract, timestamps[1], 20), timestamps[2], 10, 16.2)\n",
    "    >>> trade_6 = Trade(aapl_contract, MarketOrder(aapl_contract, timestamps[1], -20), timestamps[2], -10, 16.5)\n",
    "    >>> trade_7 = Trade(aapl_contract, MarketOrder(aapl_contract, timestamps[1], -20), timestamps[2], -10, 16.5)\n",
    "    >>> contract_pnl._add_trades([trade_5, trade_6])\n",
    "    >>> contract_pnl._add_trades([trade_7])\n",
    "    >>> df = contract_pnl.df()\n",
    "    >>> assert (len(df == 1))\n",
    "    >>> row = df.iloc[0]\n",
    "    >>> assert row.to_dict() == {'symbol': 'AAPL',\n",
    "    ... 'timestamp': pd.Timestamp('2018-01-03 00:00:00'),\n",
    "    ... 'position': -10,\n",
    "    ... 'price': 12.1,\n",
    "    ... 'unrealized': 44.0,\n",
    "    ... 'realized': 3.000000000000007,\n",
    "    ... 'commission': 0.0,\n",
    "    ... 'fee': 0.0,\n",
    "    ... 'net_pnl': 47.00000000000001}\n",
    "    '''\n",
    "    def __init__(self, \n",
    "                 contract: Contract, \n",
    "                 account_timestamps: np.ndarray, \n",
    "                 price_function: Callable[[Contract, np.ndarray, int, SimpleNamespace], float],\n",
    "                 strategy_context: SimpleNamespace) -> None:\n",
    "        self.contract = contract\n",
    "        self._price_function = price_function\n",
    "        self.strategy_context = strategy_context\n",
    "        self._account_timestamps = account_timestamps\n",
    "        self._trade_pnl = SortedDict()\n",
    "        self._net_pnl = SortedDict()\n",
    "        # Store trades that are not offset so when new trades come in we can offset against these to calc pnl\n",
    "        self.open_qtys = np.empty(0, dtype=int)\n",
    "        self.open_prices = np.empty(0, dtype=float)\n",
    "        self.first_trade_timestamp: Optional[np.datetime64] = None\n",
    "        self.final_pnl = np.nan\n",
    "        self.new_trades_added = False\n",
    "        \n",
    "    def _add_trades(self, trades: Sequence[Trade]) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            trades: Must be sorted by timestamp\n",
    "        '''\n",
    "        if not len(trades): return\n",
    "        timestamps = np.unique([trade.timestamp for trade in trades])\n",
    "        if len(self._trade_pnl):\n",
    "            prev_max_timestamp, _ = self._trade_pnl.peekitem(-1)\n",
    "            assert timestamps[0] >= prev_max_timestamp, \\\n",
    "                f'Trades can only be added with non-decreasing timestamps current: {timestamps[0]} prev max: {prev_max_timestamp}'\n",
    "            \n",
    "        if self.first_trade_timestamp is None: self.first_trade_timestamp = timestamps[0]\n",
    "            \n",
    "        self.new_trades_added = True\n",
    "                \n",
    "        for i, timestamp in enumerate(timestamps):\n",
    "            t_trades = [trade for trade in trades if trade.timestamp == timestamp]\n",
    "            open_qtys, open_prices, open_qty, weighted_avg_price, realized_chg = calc_trade_pnl(\n",
    "                self.open_qtys, self.open_prices, \n",
    "                np.array([trade.qty for trade in t_trades]), \n",
    "                np.array([trade.price for trade in t_trades]),\n",
    "                self.contract.multiplier)\n",
    "            self.open_qtys = open_qtys\n",
    "            self.open_prices = open_prices\n",
    "            position_chg = sum([trade.qty for trade in t_trades])\n",
    "            commission_chg = sum([trade.commission for trade in t_trades])\n",
    "            fee_chg = sum([trade.fee for trade in t_trades])\n",
    "            index = find_index_before(self._trade_pnl, timestamp)\n",
    "            if index == -1:\n",
    "                self._trade_pnl[timestamp] = (position_chg, realized_chg, fee_chg, commission_chg, open_qty, weighted_avg_price)\n",
    "            else:\n",
    "                prev_timestamp, (prev_position, prev_realized, prev_fee, prev_commission, _, _) = self._trade_pnl.peekitem(index)\n",
    "                self._trade_pnl[timestamp] = (prev_position + position_chg, prev_realized + realized_chg,\n",
    "                                              prev_fee + fee_chg, prev_commission + commission_chg, open_qty, weighted_avg_price)\n",
    "            self.calc_net_pnl(timestamp)\n",
    "            \n",
    "    def calc_net_pnl(self, timestamp: np.datetime64) -> None:\n",
    "        # If we already calculated unrealized pnl for this timestamp and no new trades were added no need to do anything\n",
    "        if timestamp in self._net_pnl and not self.new_trades_added: return\n",
    "        if self.first_trade_timestamp is None or timestamp < self.first_trade_timestamp: return\n",
    "        # TODO: Option expiry should be a special case.  If option expires at 3:00 pm, we put in an expiry order at 3 pm and the\n",
    "        # trade comes in at 3:01 pm.  In this case, the final pnl is recorded at 3:01 but should be at 3 pm.\n",
    "        if self.contract.expiry is not None and timestamp > self.contract.expiry and not math.isnan(self.final_pnl): return\n",
    "        \n",
    "        # make sure timestamp is in the sequence of timestamps we were given \n",
    "        i = np.searchsorted(self._account_timestamps, timestamp)\n",
    "        assert self._account_timestamps[i] == timestamp, f'timestamp {timestamp} not found'\n",
    "\n",
    "        # Find most current trade PNL, i.e. with the index before or equal to current timestamp.  If not found, set to 0's\n",
    "        trade_pnl_index = find_index_before(self._trade_pnl, timestamp)\n",
    "        if trade_pnl_index == -1:\n",
    "            realized, fee, commission, open_qty, open_qty, weighted_avg_price = 0, 0, 0, 0, 0, 0\n",
    "        else:\n",
    "            _, (_, realized, fee, commission, open_qty, weighted_avg_price) = self._trade_pnl.peekitem(trade_pnl_index)\n",
    "\n",
    "        price = np.nan\n",
    "\n",
    "        if math.isclose(open_qty, 0):\n",
    "            unrealized = 0.0\n",
    "        else:\n",
    "            price = self._price_function(self.contract, self._account_timestamps, i, self.strategy_context)  # type: ignore\n",
    "            assert np.isreal(price), \\\n",
    "                f'Unexpected price type: {price} {type(price)} for contract: {self.contract} timestamp: {self._account_timestamps[i]}'\n",
    "\n",
    "            if math.isnan(price):\n",
    "                index = find_index_before(self._net_pnl, timestamp)  # Most recent unrealized pnl\n",
    "                if index == -1:\n",
    "                    prev_unrealized, prev_open_qty = 0, 0\n",
    "                else:\n",
    "                    _, (_, prev_open_qty, prev_unrealized, _) = self._net_pnl.peekitem(index)\n",
    "                unrealized = prev_unrealized + (open_qty - prev_open_qty) * (price - weighted_avg_price) * self.contract.multiplier\n",
    "            else:\n",
    "                unrealized = open_qty * (price - weighted_avg_price) * self.contract.multiplier\n",
    "                \n",
    "        net_pnl = realized + unrealized - commission - fee\n",
    "\n",
    "        self._net_pnl[timestamp] = (price, open_qty, unrealized, net_pnl)\n",
    "        if self.contract.expiry is not None and timestamp > self.contract.expiry:\n",
    "            self.final_pnl = net_pnl\n",
    "        self.new_trades_added = False\n",
    "        \n",
    "    def position(self, timestamp: np.datetime64) -> float:\n",
    "        index = find_index_before(self._trade_pnl, timestamp)\n",
    "        if index == -1: return 0.\n",
    "        _, (position, _, _, _, _, _) = self._trade_pnl.peekitem(index)  # Less than or equal to timestamp\n",
    "        return position\n",
    "    \n",
    "    def net_pnl(self, timestamp: np.datetime64) -> float:\n",
    "        if self.contract.expiry is not None and timestamp > self.contract.expiry and not math.isnan(self.final_pnl):\n",
    "            return self.final_pnl\n",
    "        index = find_index_before(self._net_pnl, timestamp)\n",
    "        if index == -1: return 0.\n",
    "        _, (_, _, _, net_pnl) = self._net_pnl.peekitem(index)  # Less than or equal to timestamp\n",
    "        return net_pnl\n",
    "    \n",
    "    def pnl(self, timestamp: np.datetime64) -> Tuple[float, float, float, float, float, float, float]:\n",
    "        index = find_index_before(self._trade_pnl, timestamp)\n",
    "        position, realized, fee, commission, price, unrealized, net_pnl = 0, 0, 0, 0, 0, 0, 0\n",
    "        if index != -1:\n",
    "            _, (position, realized, fee, commission, _, _) = self._trade_pnl.peekitem(index)  # Less than or equal to timestamp\n",
    "        \n",
    "        index = find_index_before(self._net_pnl, timestamp)\n",
    "        if index != -1:\n",
    "            _, (price, open_position, unrealized, net_pnl) = self._net_pnl.peekitem(index)  # Less than or equal to timestamp\n",
    "        return position, price, realized, unrealized, fee, commission, net_pnl\n",
    "    \n",
    "    def df(self) -> pd.DataFrame:\n",
    "        '''Returns a pandas dataframe with pnl data'''\n",
    "        df_trade_pnl = pd.DataFrame.from_records([\n",
    "            (k, v[0], v[1], v[2], v[3]) for k, v in self._trade_pnl.items()],\n",
    "            columns=['timestamp', 'position', 'realized', 'fee', 'commission'])\n",
    "        df_net_pnl = pd.DataFrame.from_records([\n",
    "            (k, v[0], v[2], v[3]) for k, v in self._net_pnl.items()],\n",
    "            columns=['timestamp', 'price', 'unrealized', 'net_pnl'])\n",
    "        all_timestamps = np.unique(np.concatenate((df_trade_pnl.timestamp.values, df_net_pnl.timestamp.values)))\n",
    "        df_trade_pnl = df_trade_pnl.set_index('timestamp').reindex(all_timestamps, method='ffill').reset_index()\n",
    "        df_trade_pnl = leading_nan_to_zero(df_trade_pnl, ['position', 'realized', 'fee', 'commission'])\n",
    "        df_net_pnl = df_net_pnl.set_index('timestamp').reindex(all_timestamps, method='ffill').reset_index()\n",
    "        del df_net_pnl['timestamp']\n",
    "        df = pd.concat([df_trade_pnl, df_net_pnl], axis=1)\n",
    "        df['symbol'] = self.contract.symbol\n",
    "        df = df[['symbol', 'timestamp', 'position', 'price', 'unrealized', 'realized', 'commission', 'fee', 'net_pnl']]\n",
    "        return df\n",
    "         \n",
    "\n",
    "def _get_calc_timestamps(timestamps: np.ndarray, pnl_calc_time: int) -> np.ndarray:\n",
    "    time_delta = np.timedelta64(pnl_calc_time, 'm')\n",
    "    calc_timestamps = np.unique(timestamps.astype('M8[D]')) + time_delta\n",
    "    calc_indices = np.searchsorted(timestamps, calc_timestamps, side='left') - 1\n",
    "    if calc_indices[0] == -1: calc_indices[0] = 0\n",
    "    return np.unique(timestamps[calc_indices])\n",
    "\n",
    "\n",
    "def _net_trade(stacks: Dict[str, deque], trade: Trade) -> Optional[RoundTripTrade]:\n",
    "    stack = stacks[trade.contract.symbol]\n",
    "    if not len(stack) or np.sign(trade.qty) == np.sign(stack[0].qty):\n",
    "        stack.append(trade)\n",
    "        return None\n",
    "    \n",
    "    entry = stack[0]\n",
    "    qty = min(abs(entry.qty), abs(trade.qty)) * np.sign(entry.qty)\n",
    "    pnl = qty * (trade.price - entry.price) * entry.contract.multiplier - trade.commission - entry.commission\n",
    "    rt = RoundTripTrade(entry.contract, \n",
    "                        entry.order, trade.order, \n",
    "                        entry.timestamp, trade.timestamp,\n",
    "                        qty, \n",
    "                        entry.price, trade.price, \n",
    "                        entry.order.reason_code, trade.order.reason_code,\n",
    "                        entry.commission, trade.commission,\n",
    "                        entry.properties, trade.properties,\n",
    "                        pnl)\n",
    "    resid = entry.qty - qty\n",
    "    trade.qty += qty\n",
    "    if resid == 0:\n",
    "        stack.popleft()\n",
    "    return rt\n",
    "\n",
    "\n",
    "class Account:\n",
    "    '''An Account calculates pnl for a set of contracts'''\n",
    "    def __init__(self, \n",
    "                 contract_groups: Sequence[ContractGroup], \n",
    "                 timestamps: np.ndarray,\n",
    "                 price_function: Callable[[Contract, np.ndarray, int, SimpleNamespace], float],\n",
    "                 strategy_context: SimpleNamespace,\n",
    "                 starting_equity: float = 1.0e6, \n",
    "                 pnl_calc_time: int = 15 * 60) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            contract_groups: Contract groups that we want to compute PNL for\n",
    "            timestamps: Timestamps that we might compute PNL at\n",
    "            price_function: Function that returns contract prices used to compute pnl\n",
    "            strategy_context: This is passed into the price function so we can use current state of strategy to compute prices\n",
    "            starting_equity: Starting equity in account currency.  Default 1.e6\n",
    "            pnl_calc_time: Number of minutes past midnight that we should calculate PNL at.  Default 15 * 60, i.e. 3 pm\n",
    "        '''\n",
    "        self.starting_equity = starting_equity\n",
    "        self._price_function = price_function\n",
    "        self.strategy_context = strategy_context\n",
    "        \n",
    "        self.timestamps = timestamps\n",
    "        self.calc_timestamps = _get_calc_timestamps(timestamps, pnl_calc_time)\n",
    "        \n",
    "        self.contracts: MutableSet[Contract] = set()\n",
    "        self._trades: MutableSequence[Trade] = []\n",
    "        self._pnl = SortedDict()\n",
    "        self.symbol_pnls_by_contract_group: MutableMapping[str, MutableSequence[ContractPNL]] = defaultdict(list)\n",
    "        \n",
    "        self.symbol_pnls: MutableMapping[str, ContractPNL] = {}\n",
    "        \n",
    "    def symbols(self) -> MutableSequence[str]:\n",
    "        return [contract.symbol for contract in self.contracts]\n",
    "        \n",
    "    def _add_contract(self, contract: Contract, timestamp: np.datetime64) -> None:\n",
    "        if contract.symbol in self.symbol_pnls: \n",
    "            raise Exception(f'Already have contract with symbol: {contract.symbol} {contract}')\n",
    "        contract_pnl = ContractPNL(contract, self.timestamps, self._price_function, self.strategy_context)\n",
    "        self.symbol_pnls[contract.symbol] = contract_pnl\n",
    "        # For fast lookup in position function\n",
    "        self.symbol_pnls_by_contract_group[contract.contract_group.name].append(contract_pnl)\n",
    "        self.contracts.add(contract)\n",
    "        \n",
    "    def add_trades(self, trades: Sequence[Trade]) -> None:\n",
    "        trades = sorted(trades, key=lambda x: getattr(x, 'timestamp'))\n",
    "        # Break up trades by contract so we can add them in a batch\n",
    "        trades_by_contract: MutableMapping[Contract, List[Trade]] = defaultdict(list)\n",
    "        for trade in trades:\n",
    "            contract = trade.contract\n",
    "            if contract not in self.contracts: self._add_contract(contract, trade.timestamp)\n",
    "            trades_by_contract[contract].append(trade)\n",
    "            \n",
    "        for contract, contract_trades in trades_by_contract.items():\n",
    "            contract_trades.sort(key=lambda x: x.timestamp)  # type: ignore\n",
    "            self.symbol_pnls[contract.symbol]._add_trades(contract_trades)\n",
    "            \n",
    "        self._trades += trades\n",
    "        \n",
    "    def calc(self, timestamp: np.datetime64) -> None:\n",
    "        '''\n",
    "        Computes P&L and stores it internally for all contracts.\n",
    "        \n",
    "        Args:\n",
    "            timestamp: timestamp to compute P&L at.  Account remembers the last timestamp it computed P&L up to and will compute P&L\n",
    "                between these and including timestamp. If there is more than one day between the last index and current index, we will \n",
    "                include pnl for at the defined pnl_calc_time for those dates as well.\n",
    "        '''\n",
    "        if timestamp in self._pnl: return\n",
    "            \n",
    "        prev_idx = find_index_before(self._pnl, timestamp)\n",
    "        prev_timestamp = None if prev_idx == -1 else self.timestamps[prev_idx]\n",
    "            \n",
    "        # Find the last timestamp per day that is between the previous index we computed and the current index,\n",
    "        # so we can compute daily pnl in addition to the current index pnl\n",
    "        calc_timestamps = self.calc_timestamps\n",
    "        intermediate_calc_timestamps = calc_timestamps[calc_timestamps <= timestamp]\n",
    "        if prev_timestamp is not None:\n",
    "            intermediate_calc_timestamps = intermediate_calc_timestamps[intermediate_calc_timestamps > prev_timestamp]\n",
    "\n",
    "        if not len(intermediate_calc_timestamps) or intermediate_calc_timestamps[-1] != timestamp: \n",
    "            intermediate_calc_timestamps = np.append(intermediate_calc_timestamps, timestamp)\n",
    "            \n",
    "        for ts in intermediate_calc_timestamps:\n",
    "            net_pnl = 0.\n",
    "            for symbol_pnl in self.symbol_pnls.values():\n",
    "                symbol_pnl.calc_net_pnl(ts)\n",
    "                net_pnl += symbol_pnl.net_pnl(ts)\n",
    "            self._pnl[ts] = net_pnl\n",
    "        \n",
    "    def position(self, contract_group: ContractGroup, timestamp: np.datetime64) -> float:\n",
    "        '''Returns netted position for a contract_group at a given date in number of contracts or shares.'''\n",
    "        position = 0.\n",
    "        for symbol_pnl in self.symbol_pnls_by_contract_group[contract_group.name]:\n",
    "            position += symbol_pnl.position(timestamp)\n",
    "        return position\n",
    "                \n",
    "    def positions(self, contract_group: ContractGroup, timestamp: np.datetime64) -> MutableSequence[Tuple[Contract, float]]:\n",
    "        '''\n",
    "        Returns all non-zero positions in a contract group\n",
    "        '''\n",
    "        positions = []\n",
    "        for contract in contract_group.contracts:\n",
    "            symbol = contract.symbol\n",
    "            if symbol not in self.symbol_pnls: continue\n",
    "            position = self.symbol_pnls[symbol].position(timestamp)\n",
    "            if not math.isclose(position, 0): positions.append((contract, position))\n",
    "        return positions\n",
    "    \n",
    "    def equity(self, timestamp: np.datetime64) -> float:\n",
    "        '''Returns equity in this account in Account currency.  Will cause calculation if Account has not previously \n",
    "            calculated up to this date'''\n",
    "        pnl = self._pnl.get(timestamp)\n",
    "        if pnl is None:\n",
    "            self.calc(timestamp)\n",
    "            pnl = self._pnl[timestamp]\n",
    "        return self.starting_equity + pnl\n",
    "    \n",
    "    def trades(self,\n",
    "               contract_group: ContractGroup = None, \n",
    "               start_date: np.datetime64 = None, \n",
    "               end_date: np.datetime64 = None) -> MutableSequence[Trade]:\n",
    "        '''Returns a list of trades with the given symbol and with trade date between (and including) start date \n",
    "            and end date if they are specified. If symbol is None trades for all symbols are returned'''\n",
    "        return [trade for trade in self._trades if (start_date is None or trade.timestamp >= start_date) and (\n",
    "            end_date is None or trade.timestamp <= end_date) and (\n",
    "            contract_group is None or trade.contract.contract_group == contract_group)]\n",
    "    \n",
    "    def roundtrip_trades(self,\n",
    "                         contract_group: ContractGroup = None, \n",
    "                         start_date: np.datetime64 = None, \n",
    "                         end_date: np.datetime64 = None) -> MutableSequence[RoundTripTrade]:\n",
    "        '''Returns a list of round trip trades with the given symbol and with trade date \n",
    "            between (and including) start date and end date if they are specified. \n",
    "            If symbol is None trades for all symbols are returned'''\n",
    "        rts: List[RoundTripTrade] = []\n",
    "        stacks: Dict[str, deque] = defaultdict(deque)\n",
    "\n",
    "        trades = copy.deepcopy(self._trades)    \n",
    "\n",
    "        for _trade in trades:\n",
    "            trade = copy.deepcopy(_trade)\n",
    "            while True:\n",
    "                rt = _net_trade(stacks, trade)\n",
    "                if rt is None: break\n",
    "                rts.append(rt)\n",
    "                if trade.qty == 0: break\n",
    "                    \n",
    "        return [rt for rt in rts if (start_date is None or rt.entry_timestamp >= start_date)\n",
    "                and (end_date is None or rt.exit_timestamp <= end_date)\n",
    "                and (contract_group is None or rt.contract.contract_group == contract_group)]\n",
    "    \n",
    "    def df_pnl(self, contract_groups: Union[ContractGroup, Sequence[ContractGroup]] = None) -> pd.DataFrame:\n",
    "        '''\n",
    "        Returns a dataframe with P&L columns broken down by contract group and symbol\n",
    "        \n",
    "        Args:\n",
    "            contract_group: Return PNL for this contract group.  If None (default), include all contract groups\n",
    "        '''\n",
    "        if contract_groups is None: \n",
    "            contract_groups = list(set([contract.contract_group for contract in self.contracts]))\n",
    "\n",
    "        if isinstance(contract_groups, ContractGroup): contract_groups = [contract_groups]\n",
    "\n",
    "        dfs = []\n",
    "        for contract_group in contract_groups:\n",
    "            for contract in contract_group.contracts:\n",
    "                symbol = contract.symbol\n",
    "                if symbol not in self.symbol_pnls: continue\n",
    "                df = self.symbol_pnls[symbol].df()\n",
    "                if len(df) > 1:\n",
    "                    net_pnl_diff = np.diff(df.net_pnl.values)  # np.diff returns a vector one shorter than the original\n",
    "                    last_index = np.nonzero(net_pnl_diff)\n",
    "                    if len(last_index[0]): \n",
    "                        last_index_ = last_index[0][-1] + 1\n",
    "                        df = df.iloc[:last_index_ + 1]\n",
    "                df['contract_group'] = contract_group.name\n",
    "                dfs.append(df)\n",
    "        ret_df = pd.concat(dfs)\n",
    "        ret_df = ret_df.sort_values(by=['timestamp', 'contract_group', 'symbol'])\n",
    "        ret_df = ret_df[['timestamp', 'contract_group', 'symbol', 'position', 'price', 'unrealized', 'realized', \n",
    "                         'commission', 'fee', 'net_pnl']]\n",
    "        return ret_df\n",
    "    \n",
    "    def df_account_pnl(self, contract_group: ContractGroup = None) -> pd.DataFrame:\n",
    "        '''\n",
    "        Returns PNL at the account level.\n",
    "        \n",
    "        Args:\n",
    "            contract_group: If set, we only return pnl for this contract_group.  Otherwise we return pnl for all contract groups\n",
    "        '''\n",
    "\n",
    "        if contract_group is not None:\n",
    "            symbols = [contract.symbol for contract in contract_group.contracts if contract.symbol in self.symbol_pnls]\n",
    "            symbol_pnls = [self.symbol_pnls[symbol] for symbol in symbols]\n",
    "        else:\n",
    "            symbol_pnls = list(self.symbol_pnls.values())\n",
    "\n",
    "        timestamps = self.calc_timestamps\n",
    "        position = np.full(len(timestamps), 0., dtype=float)\n",
    "        realized = np.full(len(timestamps), 0., dtype=float)\n",
    "        unrealized = np.full(len(timestamps), 0., dtype=float)\n",
    "        fee = np.full(len(timestamps), 0., dtype=float)\n",
    "        commission = np.full(len(timestamps), 0., dtype=float)\n",
    "        net_pnl = np.full(len(timestamps), 0., dtype=float)\n",
    "\n",
    "        for i, timestamp in enumerate(timestamps):\n",
    "            for symbol_pnl in symbol_pnls:\n",
    "                _position, _price, _realized, _unrealized, _fee, _commission, _net_pnl = symbol_pnl.pnl(timestamp)\n",
    "                if math.isfinite(_position): position[i] += _position\n",
    "                if math.isfinite(_realized): realized[i] += _realized\n",
    "                if math.isfinite(_unrealized): unrealized[i] += _unrealized\n",
    "                if math.isfinite(_fee): fee[i] += _fee\n",
    "                if math.isfinite(_commission): commission[i] += _commission\n",
    "                if math.isfinite(_net_pnl): net_pnl[i] += _net_pnl\n",
    "\n",
    "        df = pd.DataFrame.from_records(zip(timestamps, position, unrealized, realized, commission, fee, net_pnl), \n",
    "                                       columns=['timestamp', 'position', 'unrealized', 'realized', 'commission', 'fee', 'net_pnl'])\n",
    "        df['equity'] = self.starting_equity + df.net_pnl\n",
    "        return df[['timestamp', 'position', 'unrealized', 'realized', 'commission', 'fee', 'net_pnl', 'equity']]\n",
    "    \n",
    "    def df_trades(self, \n",
    "                  contract_group: ContractGroup = None, \n",
    "                  start_date: np.datetime64 = None, \n",
    "                  end_date: np.datetime64 = None) -> pd.DataFrame:\n",
    "        '''\n",
    "        Returns a dataframe of trades\n",
    "        \n",
    "        Args:\n",
    "            contract_group: Return trades for this contract group. If None (default), include all contract groups\n",
    "            start_date: Include trades with date greater than or equal to this timestamp.\n",
    "            end_date: Include trades with date less than or equal to this timestamp.\n",
    "        '''\n",
    "        # start_date, end_date = str2date(start_date), str2date(end_date)\n",
    "        trades = self.trades(contract_group, start_date, end_date)\n",
    "        df = pd.DataFrame.from_records([(\n",
    "            trade.contract.symbol, \n",
    "            trade.timestamp, \n",
    "            trade.qty, \n",
    "            trade.price, \n",
    "            trade.fee, \n",
    "            trade.commission, \n",
    "            trade.order.timestamp, \n",
    "            trade.order.qty, \n",
    "            trade.order.reason_code, \n",
    "            (str(trade.order.properties.__dict__) if trade.order.properties.__dict__ else ''), \n",
    "            (str(trade.contract.properties.__dict__) if trade.contract.properties.__dict__ else '')) for trade in trades],\n",
    "            columns=['symbol', 'timestamp', 'qty', 'price', 'fee', 'commission', 'order_date', 'order_qty',\n",
    "                     'reason_code', 'order_props', 'contract_props'])\n",
    "        df = df.sort_values(by=['timestamp', 'symbol'])\n",
    "        return df\n",
    "    \n",
    "    def df_roundtrip_trades(self, \n",
    "                            contract_group: ContractGroup = None, \n",
    "                            start_date: np.datetime64 = None, \n",
    "                            end_date: np.datetime64 = None) -> pd.DataFrame:\n",
    "        '''\n",
    "        Returns a dataframe of round trip trades\n",
    "        \n",
    "        Args:\n",
    "            contract_group: Return trades for this contract group. If None (default), include all contract groups\n",
    "            start_date: Include trades with date greater than or equal to this timestamp.\n",
    "            end_date: Include trades with date less than or equal to this timestamp.\n",
    "        '''\n",
    "        rt_trades = self.roundtrip_trades(contract_group, start_date, end_date)\n",
    "        df_rts = pd.DataFrame.from_records([dict(\n",
    "            symbol=s.contract.symbol, \n",
    "            multiplier=s.contract.multiplier, \n",
    "            entry_timestamp=s.entry_timestamp,\n",
    "            exit_timestamp=s.exit_timestamp,\n",
    "            qty=s.qty,\n",
    "            entry_price=s.entry_price,\n",
    "            exit_price=s.exit_price,\n",
    "            entry_reason=s.entry_reason,\n",
    "            exit_reason=s.exit_reason,\n",
    "            entry_commission=s.entry_commission,\n",
    "            exit_commission=s.exit_commission,\n",
    "            net_pnl=s.net_pnl) for s in rt_trades])\n",
    "        df_rts = df_rts.sort_values(by=['entry_timestamp', 'symbol'])\n",
    "        return df_rts\n",
    " \n",
    "\n",
    "def test_account():\n",
    "    from pyqstrat.pq_types import MarketOrder\n",
    "\n",
    "    def get_close_price(contract, timestamps, idx, strategy_context):\n",
    "        if contract.symbol == \"IBM\":\n",
    "            price = idx + 10.1\n",
    "        elif contract.symbol == \"MSFT\":\n",
    "            price = idx + 15.3\n",
    "        elif contract.symbol == 'AAPL':\n",
    "            price = idx + 10\n",
    "        else:\n",
    "            raise Exception(f'unknown contract: {contract}')\n",
    "        return price\n",
    "\n",
    "    ContractGroup.clear()\n",
    "    Contract.clear()\n",
    "    ibm_cg = ContractGroup.create('IBM')\n",
    "    msft_cg = ContractGroup.create('MSFT')\n",
    "\n",
    "    ibm_contract = Contract.create('IBM', contract_group=ibm_cg)\n",
    "    msft_contract = Contract.create('MSFT', contract_group=msft_cg)\n",
    "\n",
    "    timestamps = np.array(['2018-01-01 09:00',\n",
    "                           '2018-01-02 08:00', \n",
    "                           '2018-01-02 09:00', \n",
    "                           '2018-01-05 13:35',\n",
    "                           '2018-01-05 13:36'], dtype='M8[m]')\n",
    "    trade_1 = Trade(ibm_contract, MarketOrder(ibm_contract, np.datetime64('2018-01-01 09:00'), 10), \n",
    "                    np.datetime64('2018-01-02 08:00'), 10, 10.1, commission=0.01)\n",
    "    trade_2 = Trade(ibm_contract, MarketOrder(ibm_contract, np.datetime64('2018-01-01 09:00'), -20),\n",
    "                    np.datetime64('2018-01-02 09:00'), -20, 15.1, commission=0.02)\n",
    "    trade_3 = Trade(msft_contract, MarketOrder(msft_contract, timestamps[1], 15), timestamps[1], 20, 13.2, commission=0.04)\n",
    "    trade_4 = Trade(msft_contract, MarketOrder(msft_contract, timestamps[2], 20), timestamps[2], 20, 16.2, commission=0.05)\n",
    "    trade_5 = Trade(msft_contract, MarketOrder(msft_contract, timestamps[2], 5), timestamps[2], 5, 16.21, commission=0.03)\n",
    "\n",
    "    account = Account([ibm_cg, msft_cg], timestamps, get_close_price, None)\n",
    "    account.add_trades([trade_1, trade_2, trade_3, trade_4, trade_5])\n",
    "    account.calc(np.datetime64('2018-01-05 13:35'))\n",
    "\n",
    "    assert(len(account.df_trades()) == 5)\n",
    "    assert(len(account.df_pnl()) == 6)\n",
    "    assert(np.allclose(np.array([9.99, 61.96, 79.97, 109.33, 69.97, 154.33]), account.df_pnl().net_pnl.values, rtol=0))\n",
    "    assert(np.allclose(np.array([10, 20, -10, 45, -10, 45]), account.df_pnl().position.values, rtol=0))\n",
    "    assert(np.allclose(np.array([1000000., 1000189.3, 1000224.3]), account.df_account_pnl().equity.values, rtol=0))\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_account()\n",
    "    import doctest\n",
    "    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
