{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-11-06T21:13:07.984296Z",
     "start_time": "2018-11-06T21:13:07.146799Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "from collections import defaultdict, deque\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from copy import copy\n",
    "from pyqstrat.pq_utils import str2date\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-11-06T21:13:08.022237Z",
     "start_time": "2018-11-06T21:13:07.987579Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "def _calc_pnl(open_trades, new_trades, ending_close, multiplier):\n",
    "    '''\n",
    "    >>> from collections import deque\n",
    "    >>> from pyqstrat.pq_types import Trade\n",
    "    >>> trades = deque([Trade('IBM', np.datetime64('2018-01-01 10:15:00'), 3, 51.),\n",
    "    ...              Trade('IBM', np.datetime64('2018-01-01 10:20:00'), 10, 50.),\n",
    "    ...              Trade('IBM', np.datetime64('2018-01-02 11:20:00'), -5, 45.)])\n",
    "    >>> print(_calc_pnl(open_trades = deque(), new_trades = trades, ending_close = 54, multiplier = 100))\n",
    "    (deque([IBM 2018-01-01 10:20 qty: 8 prc: 50.0 order: None]), 3200.0, -2800.0)\n",
    "    >>> trades = deque([Trade('IBM', np.datetime64('2018-01-01 10:15:00'), -8, 10.),\n",
    "    ...          Trade('IBM', np.datetime64('2018-01-01 10:20:00'), 9, 11.),\n",
    "    ...          Trade('IBM', np.datetime64('2018-01-02 11:20:00'), -4, 6.)])\n",
    "    >>> print(_calc_pnl(open_trades = deque(), new_trades = trades, ending_close = 5.8, multiplier = 100))\n",
    "    (deque([IBM 2018-01-02 11:20 qty: -3 prc: 6.0 order: None]), 60.00000000000006, -1300.0)\n",
    "    '''\n",
    "    \n",
    "    realized = 0.\n",
    "    unrealized = 0.\n",
    "    \n",
    "    trades = copy(new_trades)\n",
    "    \n",
    "    while (len(trades)):\n",
    "        trade = trades[0]\n",
    "        if not len(open_trades) or (np.sign(open_trades[-1].qty) == np.sign(trade.qty)):\n",
    "            open_trades.append(copy(trade))\n",
    "            trades.popleft()\n",
    "            continue\n",
    "            \n",
    "        if abs(trade.qty) > abs(open_trades[0].qty):\n",
    "            open_trade = open_trades.popleft()\n",
    "            realized += open_trade.qty * multiplier * (trade.price - open_trade.price)\n",
    "            trade.qty += open_trade.qty\n",
    "        else:\n",
    "            open_trade = open_trades[0]\n",
    "            realized += trade.qty * multiplier * (open_trades[-1].price - trade.price)\n",
    "            trades.popleft()\n",
    "            open_trade.qty += trade.qty\n",
    " \n",
    "    unrealized = sum([open_trade.qty * (ending_close - open_trade.price) for open_trade in open_trades]) * multiplier\n",
    "\n",
    "    return open_trades, unrealized, realized\n",
    "        \n",
    "class ContractPNL:\n",
    "    '''Computes pnl for a single contract over time given trades and market data'''\n",
    "    def __init__(self, contract, marketdata):\n",
    "        self.symbol = contract.symbol\n",
    "        self.multiplier = contract.multiplier\n",
    "        self.marketdata = marketdata\n",
    "        self.dates = self.marketdata.dates\n",
    "        self.unrealized = np.full(len(self.dates), np.nan, dtype = np.float); self.unrealized[0] = 0\n",
    "        self.realized = np.full(len(self.dates), np.nan, dtype = np.float); self.realized[0] = 0\n",
    "        \n",
    "        self.commission = np.zeros(len(self.dates), dtype = np.float);\n",
    "        self.fee = np.zeros(len(self.dates), dtype = np.float);\n",
    "        \n",
    "        self.net_pnl = np.full(len(self.dates), np.nan, dtype = np.float); self.net_pnl[0] = 0\n",
    "        self.position = np.full(len(self.dates), np.nan, dtype = np.float); self.position[0] = 0\n",
    "        self.close = self.marketdata.c\n",
    "        self._trades = []\n",
    "        self.open_trades = deque()\n",
    "        \n",
    "    def add_trades(self, trades):\n",
    "        '''Args:\n",
    "            trades: A list of Trade objects\n",
    "        '''\n",
    "        self._trades += trades\n",
    "        \n",
    "    def calc(self, prev_i, i):\n",
    "        '''Compute pnl and store it internally\n",
    "        \n",
    "        Args:\n",
    "            prev_i: Start index to compute pnl from\n",
    "            i: End index to compute pnl to\n",
    "        '''\n",
    "        calc_trades = deque([trade for trade in self._trades if trade.date > self.dates[prev_i] and trade.date <= self.dates[i]])\n",
    "        \n",
    "        if not np.isfinite(self.close[i]):\n",
    "            unrealized = self.unrealized[prev_i]\n",
    "            realized = 0. \n",
    "        else:\n",
    "            open_trades, unrealized, realized = _calc_pnl(self.open_trades, calc_trades, self.close[i], self.multiplier)\n",
    "            self.open_trades = open_trades\n",
    "            \n",
    "        self.unrealized[i] = unrealized\n",
    "        self.realized[i] = self.realized[prev_i] + realized\n",
    "        trade_qty = sum([trade.qty for trade in calc_trades])\n",
    "        self.position[i] = self.position[prev_i] + trade_qty\n",
    "        self.commission[i] = self.commission[prev_i] + sum([trade.commission for trade in calc_trades])\n",
    "        self.fee[i] = self.fee[prev_i] + sum([trade.fee for trade in calc_trades])\n",
    "        self.net_pnl[i] = self.realized[i] + self.unrealized[i] - self.commission[i] - self.fee[i]\n",
    "        if np.isnan(self.net_pnl[i]):\n",
    "            raise Exception(f'net_pnl: nan i: {i} realized: {self.realized[i]} unrealized: {self.unrealized[i]} commission: {self.commission[i]} fee: {self.fee[i]}')\n",
    "        \n",
    "    def trades(self, start_date = None, end_date = None):\n",
    "        '''Get a list of trades\n",
    "        \n",
    "        Args:\n",
    "            start_date: A string or numpy datetime64.  Trades with trade dates >= start_date will be returned.  Default None\n",
    "            end_date: A string or numpy datetime64.  Trades with trade dates <= end_date will be returned.  Default None\n",
    "        '''\n",
    "        start_date, end_date = str2date(start_date), str2date(end_date)\n",
    "        trades = [trade for trade in self._trades if (start_date is None or trade.date >= start_date) and (end_date is None or trade.date <= end_date)]\n",
    "        return trades\n",
    "         \n",
    "    def df(self):\n",
    "        '''Returns a pandas dataframe with pnl data, indexed by date'''\n",
    "        df = pd.DataFrame({'date' : self.dates, 'unrealized' : self.unrealized, 'realized' : self.realized, \n",
    "                           'fee' : self.fee, 'net_pnl' : self.net_pnl, 'position' : self.position})\n",
    "        df.dropna(subset = ['unrealized', 'realized'], inplace = True)\n",
    "        df['symbol'] = self.symbol\n",
    "        return df[['symbol', 'date', 'unrealized', 'realized', 'fee', 'net_pnl', 'position']].set_index('date')\n",
    "\n",
    "class Account:\n",
    "    '''An Account calculates pnl for a set of contracts'''\n",
    "    def __init__(self, contracts, marketdata_collection, starting_equity = 1.0e6, calc_frequency = 'D'):\n",
    "        '''\n",
    "        Args:\n",
    "            contracts (list of Contract): Contracts that we want to compute PNL for\n",
    "            marketdata_collection (MarketDataCollection): MarketData corresponding to contracts \n",
    "            starting_equity (float, optional): Starting equity in account currency.  Default 1.e6\n",
    "            calc_frequency (str, optional): Account will calculate pnl at this frequency.  Default 'D' for daily\n",
    "       \n",
    "        >>> from pyqstrat.marketdata import MarketData, MarketDataCollection\n",
    "        >>> from pyqstrat.pq_types import Contract\n",
    "        >>> dates = np.array(['2018-01-01', '2018-01-02'], dtype = 'M8[D]')\n",
    "        >>> account = Account([Contract(\"IBM\")], MarketDataCollection([\"IBM\"], [MarketData(dates, [8.1, 8.2])]))\n",
    "        >>> np.set_printoptions(formatter = {'float' : lambda x : f'{x:.4f}'})  # After numpy 1.13 positive floats don't have a leading space for sign\n",
    "        >>> print(account.marketdata['IBM'].c)\n",
    "        [8.1000 8.2000]\n",
    "        '''\n",
    "        if calc_frequency != 'D': raise Exception('unknown calc frequency: {}'.format(calc_frequency))\n",
    "        self.calc_freq = calc_frequency\n",
    "        self.contract_pnls = defaultdict()\n",
    "        self.current_calc_index = 0\n",
    "        self.all_dates = None\n",
    "        self.starting_equity = starting_equity\n",
    "        \n",
    "        if contracts is not None or marketdata_collection is not None:\n",
    "            if contracts is None or marketdata_collection is None:\n",
    "                raise Exception(\"either contracts and marketdata_collection must both be None or they both must be non-None\")\n",
    "            contract_symbols = sorted([contract.symbol for contract in contracts])\n",
    "            md_symbols = sorted([tup[0] for tup in marketdata_collection.items()])\n",
    "            if contract_symbols != md_symbols:\n",
    "                raise Exception(f\"contracts and marketdata must have the same symbols: {contract_symbols} {md_symbols}\")\n",
    "                \n",
    "        self.marketdata = {}\n",
    "                \n",
    "        if marketdata_collection is not None:\n",
    "            for symbol, marketdata in marketdata_collection.items():\n",
    "                self.add_marketdata(symbol, marketdata)\n",
    "\n",
    "        if contracts is not None:\n",
    "            for contract in contracts: \n",
    "                self.add_contract(contract)\n",
    "                \n",
    "    \n",
    "    def _set_dates(self, dates):\n",
    "        if self.all_dates is not None and not np.array_equal(dates, all_dates):\n",
    "            raise Exception('all symbols in a strategy must have the same dates')\n",
    "        self.all_dates = dates\n",
    "        calc_dates = dates.astype('M8[D]')\n",
    "        self.calc_dates = np.unique(calc_dates)\n",
    "        self.calc_indices = np.searchsorted(dates, self.calc_dates, side='left') - 1\n",
    "        if self.calc_indices[0] == -1: self.calc_indices[0] = 0\n",
    "        self._equity = np.full(len(dates), np.nan, dtype = np.float); \n",
    "        self._equity[0] = self.starting_equity\n",
    "        \n",
    "    def symbols(self):\n",
    "        return list(self.contract_pnls.keys())\n",
    "        \n",
    "    def add_contract(self, contract):\n",
    "         self.contract_pnls[contract.symbol] = ContractPNL(contract, self.marketdata[contract.symbol])\n",
    "        \n",
    "    def add_marketdata(self, symbol, marketdata):\n",
    "        if self.all_dates is None: self._set_dates(marketdata.dates)\n",
    "        self.marketdata[symbol] = marketdata\n",
    "        \n",
    "    def _add_trades(self, symbol, trades):\n",
    "        self.contract_pnls[symbol].add_trades(trades)\n",
    "        \n",
    "    def calc(self, i):\n",
    "        '''\n",
    "        Computes P&L and stores it internally for all contracts.\n",
    "        \n",
    "        Args:\n",
    "            i: Index to compute P&L at.  Account remembers the last index it computed P&L up to and will compute P&L between these two indices\n",
    "        '''\n",
    "        calc_indices = self.calc_indices[:]\n",
    "        if self.current_calc_index == i: return\n",
    "        intermediate_calc_indices = np.ravel(np.where(np.logical_and(calc_indices > self.current_calc_index, calc_indices <= i)))\n",
    "\n",
    "        if not len(intermediate_calc_indices) or calc_indices[intermediate_calc_indices[-1]] != i: \n",
    "            calc_indices = np.append(calc_indices, i)\n",
    "            intermediate_calc_indices = np.append(intermediate_calc_indices, len(calc_indices) - 1)\n",
    "            \n",
    "        for symbol, symbol_pnl in self.contract_pnls.items():\n",
    "            prev_calc_index = self.current_calc_index\n",
    "            for idx in intermediate_calc_indices:\n",
    "                calc_index = calc_indices[idx]\n",
    "                symbol_pnl.calc(prev_calc_index, calc_index)\n",
    "                self._equity[calc_index] = self._equity[prev_calc_index] + symbol_pnl.net_pnl[calc_index] - symbol_pnl.net_pnl[prev_calc_index]\n",
    "                # print(f'prev_calc_index: {prev_calc_index} calc_index: {calc_index} prev_equity: {self._equity[prev_calc_index]} net_pnl: {symbol_pnl.net_pnl[calc_index]} prev_net_pnl: {symbol_pnl.net_pnl[prev_calc_index]}')\n",
    "                prev_calc_index = calc_index\n",
    "                \n",
    "        self.current_calc_index = i\n",
    "        \n",
    "    def position(self, symbol, date):\n",
    "        '''Returns position for a symbol at a given date in number of contracts or shares.  Will cause calculation if Account has not previously calculated\n",
    "          up to this date'''\n",
    "        i = self.find_index_before(date)\n",
    "        self.calc(i)\n",
    "        return self.contract_pnls[symbol].position[i]\n",
    "    \n",
    "    def equity(self, date):\n",
    "        '''Returns equity in this account in Account currency.  Will cause calculation if Account has not previously calculated up to this date'''\n",
    "        i = self.find_index_before(date)\n",
    "        self.calc(i)\n",
    "        return self._equity[i]\n",
    "    \n",
    "    def trades(self, symbol = None, start_date = None, end_date = None):\n",
    "        '''Returns a list of trades with the given symbol and with trade date between (and including) start date and end date if they are specified.\n",
    "          If symbol is None trades for all symbols are returned'''\n",
    "        start_date, end_date = str2date(start_date), str2date(end_date)\n",
    "        if symbol is None:\n",
    "            trades = []\n",
    "            for symbol, sym_pnl in self.contract_pnls.items():\n",
    "                trades += sym_pnl.trades(start_date, end_date)\n",
    "            return trades\n",
    "        else:\n",
    "            return self.contract_pnls[symbol].trades(start_date, end_date)\n",
    "        \n",
    "    def find_index_before(self, date):\n",
    "        '''Returns the market data index before or at date'''\n",
    "        return np.searchsorted(self.all_dates, date)\n",
    "        \n",
    "    def transfer_cash(self, date, amount):\n",
    "        '''Move cash from one portfolio to another'''\n",
    "        i = self.find_index_before(date)\n",
    "        curr_equity = self.equity(date)\n",
    "        if (amount > curr_equity): amount = curr_equity # Cannot make equity negative\n",
    "        self._equity[i] -= amount\n",
    "        return amount\n",
    "\n",
    "    def df_pnl(self, symbol = None):\n",
    "        '''Returns a dataframe with P&L columns.  If symbol is set to None (default), sums up P&L across symbols'''\n",
    "        if symbol:\n",
    "            ret = self.contract_pnls[symbol].df()\n",
    "        else:\n",
    "            dfs = []\n",
    "            for symbol, symbol_pnl in self.contract_pnls.items():\n",
    "                df = symbol_pnl.df()\n",
    "                dfs.append(df)\n",
    "            ret = pd.concat(dfs)\n",
    "            ret = ret.reset_index().groupby('date').sum()\n",
    "        df_equity = pd.DataFrame({'equity' : self._equity}, index = self.all_dates).dropna()\n",
    "        ret = pd.merge(ret, df_equity, left_index = True, right_index = True, how = 'outer')\n",
    "        ret.index.name = 'date'\n",
    "        return ret\n",
    "    \n",
    "    def df_trades(self, symbol = None, start_date = None, end_date = None):\n",
    "        '''Returns a dataframe with data from trades with the given symbol and with trade date between (and including) start date and end date\n",
    "          if they are specified.  If symbol is None, trades for all symbols are returned'''\n",
    "        start_date, end_date = str2date(start_date), str2date(end_date)\n",
    "        if symbol:\n",
    "            trades = self.contract_pnls[symbol].trades(start_date, end_date)\n",
    "        else:\n",
    "            trades = [v.trades(start_date, end_date) for v in self.contract_pnls.values()]\n",
    "            trades = [trade for sublist in trades for trade in sublist] # flatten list\n",
    "        df = pd.DataFrame.from_records([(trade.symbol, trade.date, trade.qty, trade.price, trade.fee, trade.commission, trade.order.date, trade.order.qty, trade.order.params()) for trade in trades],\n",
    "                    columns = ['symbol', 'date', 'qty', 'price', 'fee', 'commission', 'order_date', 'order_qty', 'order_params'])\n",
    "        return df"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
