{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "source": [
    "## A Simple Pair Trading Strategy\n",
    "\n",
    "**_Please go through the \"building strategies\" notebook before looking at this notebook._**\n",
    "\n",
    "Let's build a aimple pair trading strategy to show how you can trade multiple symbols in a strategy.  We will trade 2 stocks, Coca-Cola (KO) and Pepsi (PEP)\n",
    "\n",
    "1.  We will buy KO and sell PEP when the price ratio KO / PEP is more than 1 standard deviation lower than its 5 day simple moving average.  \n",
    "2.  We will buy PEP and sell KO when the price ratio KO / PEP is more than 1 standard deviation higher than its 5 day simple moving average.\n",
    "3.  We will exit when the price ratio is less than +/- 0.5 standard deviations away from its simple moving average\n",
    "4.  We will size the trades in 1 and 2 by allocating 10% of our capital to each trade.\n",
    "\n",
    "First lets load some price data in fifteen minute bars."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-16T02:39:13.224515Z",
     "start_time": "2022-08-16T02:39:12.978549Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import datetime\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.stats\n",
    "import os\n",
    "from types import SimpleNamespace\n",
    "\n",
    "import pyqstrat as pq\n",
    "\n",
    "pq.set_defaults() # Set some display defaults to make dataframes and plots easier to look at\n",
    "\n",
    "try:\n",
    "    ko_file = os.path.dirname(os.path.realpath(__file__)) + './support/coke_15_min_prices.csv.gz'\n",
    "    pep_file = os.path.dirname(os.path.realpath(__file__)) + './support/pepsi_15_min_prices.csv.gz' # If we are running from unit tests\n",
    "    \n",
    "except:\n",
    "    ko_file_path = '../notebooks/support/coke_15_min_prices.csv.gz'\n",
    "    pep_file_path = '../notebooks/support/pepsi_15_min_prices.csv.gz'\n",
    "\n",
    "ko_prices = pd.read_csv(ko_file_path)\n",
    "pep_prices = pd.read_csv(pep_file_path)\n",
    "\n",
    "ko_prices['timestamp'] = pd.to_datetime(ko_prices.date)\n",
    "pep_prices['timestamp'] = pd.to_datetime(pep_prices.date)\n",
    "\n",
    "timestamps = ko_prices.timestamp.values\n",
    "\n",
    "ko_contract_group = pq.ContractGroup.create('KO')\n",
    "pep_contract_group = pq.ContractGroup.create('PEP')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-02-27T02:39:20.994825Z",
     "start_time": "2019-02-27T02:39:20.991868Z"
    }
   },
   "source": [
    "Lets compute the ratio of the two prices and add it to the market data.  Since the two price series have the exact same timestamps, we can simply divide the two close price series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-16T02:39:13.227040Z",
     "start_time": "2022-08-16T02:39:13.225268Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "ratio = ko_prices.c / pep_prices.c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-02-27T02:41:42.162910Z",
     "start_time": "2019-02-27T02:41:42.158983Z"
    }
   },
   "source": [
    "Next, lets create an indicator for the zscore, and plot it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-16T02:39:13.349445Z",
     "start_time": "2022-08-16T02:39:13.228604Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 102504 characters]\n",
      "text/plain": [
       "<Figure size 1080x576 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def zscore_indicator(symbol, timestamps, indicators, strategy_context): # simple moving average\n",
    "    ratio = indicators.ratio\n",
    "    r = pd.Series(ratio).rolling(window = 130)\n",
    "    mean = r.mean()\n",
    "    std = r.std(ddof = 0)\n",
    "    zscore = (ratio - mean) / std\n",
    "    zscore = np.nan_to_num(zscore)\n",
    "    return zscore\n",
    "\n",
    "ko_zscore = zscore_indicator(None, None, SimpleNamespace(ratio = ratio), None)\n",
    "\n",
    "ratio_subplot = pq.Subplot([pq.TimeSeries('ratio', timestamps, ratio)], ylabel = 'Ratio')\n",
    "zscore_subplot = pq.Subplot([pq.TimeSeries('zscore', timestamps, ko_zscore)], ylabel = 'ZScore')\n",
    "plot = pq.Plot([ratio_subplot, zscore_subplot], title = 'KO')\n",
    "plot.draw();\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now lets create the signal that will tell us to get in when the zscore is +/-1 and get out when its less than +/- 0.5.  We use a signal value of 2 to figure out when to go long, and -2 to figure out when to go short.  A value of 1 means get out of a long position, and -1 means get out of a short position.  We also plot the signal to check it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-16T02:39:13.428009Z",
     "start_time": "2022-08-16T02:39:13.350320Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 99820 characters]\n",
      "text/plain": [
       "<Figure size 1080x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def pair_strategy_signal(contract_group, timestamps, indicators, parent_signals, strategy_context):\n",
    "    zscore = indicators.zscore\n",
    "    signal = np.where(zscore > 1, 2, 0)\n",
    "    signal = np.where(zscore < -1, -2, signal)\n",
    "    signal = np.where((zscore > 0.5) & (zscore < 1), 1, signal)\n",
    "    signal = np.where((zscore < -0.5) & (zscore > -1), -1, signal)\n",
    "    if contract_group.name == 'PEP': signal = -1. * signal\n",
    "    return signal\n",
    "\n",
    "signal = pair_strategy_signal(ko_contract_group, timestamps, SimpleNamespace(zscore = ko_zscore), None, None)\n",
    "signal_subplot = pq.Subplot([pq.TimeSeries('signal', timestamps, signal)], ylabel = 'Signal')\n",
    "plot = pq.Plot([signal_subplot], title = 'KO', show_date_gaps = False)\n",
    "plot.draw();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we create the trading rule and market simulator functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-16T02:39:13.433703Z",
     "start_time": "2022-08-16T02:39:13.428863Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "def pair_trading_rule(contract_group, i, timestamps, indicators, signal, account, strategy_context):\n",
    "    timestamp = timestamps[i]\n",
    "    curr_pos = account.position(contract_group, timestamp)\n",
    "    signal_value = signal[i]\n",
    "    risk_percent = 0.1\n",
    "    \n",
    "    orders = []\n",
    "    \n",
    "    symbol = contract_group.name\n",
    "    \n",
    "    contract = contract_group.get_contract(symbol)\n",
    "    if contract is None:\n",
    "        contract = pq.Contract.create(symbol, contract_group = contract_group)\n",
    "    \n",
    "    # if we don't already have a position, check if we should enter a trade\n",
    "    if math.isclose(curr_pos, 0):\n",
    "        if signal_value == 2 or signal_value == -2:\n",
    "            curr_equity = account.equity(timestamp)\n",
    "            order_qty = np.round(curr_equity * risk_percent / indicators.c[i] * np.sign(signal_value))\n",
    "            trigger_price = indicators.c[i]\n",
    "            reason_code = pq.ReasonCode.ENTER_LONG if order_qty > 0 else pq.ReasonCode.ENTER_SHORT\n",
    "            orders.append(pq.MarketOrder(contract, timestamp, order_qty, reason_code = reason_code))\n",
    "        \n",
    "    else: # We have a current position, so check if we should exit\n",
    "        if (curr_pos > 0 and signal_value == -1) or (curr_pos < 0 and signal_value == 1):\n",
    "            order_qty = -curr_pos\n",
    "            reason_code = pq.ReasonCode.EXIT_LONG if order_qty < 0 else pq.ReasonCode.EXIT_SHORT\n",
    "            orders.append(pq.MarketOrder(contract, timestamp, order_qty, reason_code = reason_code))\n",
    "    return orders\n",
    "\n",
    "def market_simulator(orders, i, timestamps, indicators, signals, strategy_context):\n",
    "    trades = []\n",
    "    \n",
    "    timestamp = timestamps[i]\n",
    "    \n",
    "    for order in orders:\n",
    "        trade_price = np.nan\n",
    "        \n",
    "        contract_group = order.contract.contract_group\n",
    "        ind = indicators[contract_group]\n",
    "        \n",
    "        o, h, l, c = ind.o[i], ind.h[i], ind.l[i], ind.c[i]\n",
    "        \n",
    "        if isinstance(order, pq.MarketOrder):\n",
    "            trade_price = 0.5 * (o + h) if order.qty > 0 else 0.5 * (o + l)\n",
    "        else:\n",
    "            raise Exception(f'unexpected order type: {order}')\n",
    "            \n",
    "        if np.isnan(trade_price): continue\n",
    "            \n",
    "        trade = pq.Trade(order.contract, order, timestamp, order.qty, trade_price, commission = 0, fee = 0)\n",
    "        order.status = 'filled'\n",
    "                           \n",
    "        trades.append(trade)\n",
    "                           \n",
    "    return trades\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets run the strategy, plot the results and look at the returns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-16T02:39:14.004452Z",
     "start_time": "2022-08-16T02:39:13.434831Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    },
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 212732 characters]\n",
      "text/plain": [
       "<Figure size 1440x1080 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 211404 characters]\n",
      "text/plain": [
       "<Figure size 1440x1080 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def get_price(contract, timestamps, i, strategy_context):\n",
    "    if contract.symbol == 'KO':\n",
    "        return strategy_context.ko_price[i]\n",
    "    elif contract.symbol == 'PEP':\n",
    "        return strategy_context.pep_price[i]\n",
    "    raise Exception(f'Unknown symbol: {symbol}')\n",
    "\n",
    "strategy_context = SimpleNamespace(ko_price = ko_prices.c.values, pep_price = pep_prices.c.values)\n",
    "\n",
    "strategy = pq.Strategy(timestamps, [ko_contract_group, pep_contract_group], get_price, trade_lag = 1, strategy_context = strategy_context)\n",
    "for tup in [(ko_contract_group, ko_prices), (pep_contract_group, pep_prices)]:\n",
    "    for column in ['o', 'h', 'l', 'c']:\n",
    "        strategy.add_indicator(column, tup[1][column].values, contract_groups = [tup[0]])\n",
    "        \n",
    "strategy.add_indicator('ratio', ratio)\n",
    "strategy.add_indicator('zscore', zscore_indicator, depends_on = ['ratio'])\n",
    "\n",
    "strategy.add_signal('pair_strategy_signal', pair_strategy_signal, depends_on_indicators = ['zscore'])\n",
    "\n",
    "# ask pqstrat to call our trading rule when the signal has one of the values [-2, -1, 1, 2]\n",
    "strategy.add_rule('pair_trading_rule', pair_trading_rule, \n",
    "                  signal_name = 'pair_strategy_signal', sig_true_values = [-2, -1, 1, 2])\n",
    "\n",
    "strategy.add_market_sim(market_simulator)\n",
    "\n",
    "portfolio = pq.Portfolio()\n",
    "portfolio.add_strategy('pair_strategy', strategy)\n",
    "portfolio.run()\n",
    "\n",
    "strategy.plot(primary_indicators = ['c'], secondary_indicators = ['zscore'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-16T02:39:14.169791Z",
     "start_time": "2022-08-16T02:39:14.005334Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>gmean</th>\n",
       "      <th>amean</th>\n",
       "      <th>std</th>\n",
       "      <th>shrp</th>\n",
       "      <th>srt</th>\n",
       "      <th>k</th>\n",
       "      <th>calmar</th>\n",
       "      <th>mar</th>\n",
       "      <th>mdd_pct</th>\n",
       "      <th>mdd_dates</th>\n",
       "      <th>dd_3y_pct</th>\n",
       "      <th>dd_3y_timestamps</th>\n",
       "      <th>up_dwn</th>\n",
       "      <th>2019</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>0.07401</td>\n",
       "      <td>0.07166</td>\n",
       "      <td>0.001428</td>\n",
       "      <td>3.161</td>\n",
       "      <td>14.04</td>\n",
       "      <td>4.849</td>\n",
       "      <td>21.44</td>\n",
       "      <td>21.44</td>\n",
       "      <td>0.003343</td>\n",
       "      <td>2019-01-22/2019-02-13</td>\n",
       "      <td>0.003343</td>\n",
       "      <td>2019-01-22/2019-02-13</td>\n",
       "      <td>11/12/0.478</td>\n",
       "      <td>0.07401</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    gmean    amean       std   shrp    srt      k calmar    mar   mdd_pct  \\\n",
       "  0.07401  0.07166  0.001428  3.161  14.04  4.849  21.44  21.44  0.003343   \n",
       "\n",
       "              mdd_dates dd_3y_pct       dd_3y_timestamps       up_dwn     2019  \n",
       "  2019-01-22/2019-02-13  0.003343  2019-01-22/2019-02-13  11/12/0.478  0.07401  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 102152 characters]\n",
      "text/plain": [
       "<Figure size 1080x576 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "strategy.evaluate_returns();"
   ]
  }
 ],
 "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
}
