{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-09T23:44:56.522521Z",
     "start_time": "2022-11-09T23:44:55.824754Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running typecheck\n",
      "Success: no issues found in 1 source file\n",
      "\n",
      "running flake8\n",
      "stdin:738:13: F841 local variable '_date_range' is assigned to but never used\n",
      "\n",
      "order_qty: 2090.0 curr_equity: 1000000.0 timestamp: 2019-01-23T11:15:00.000000000 risk_percent: 0.1 indicator: 47.84 signal_value: 2\n",
      "order_qty: -912.0 curr_equity: 1000000.0 timestamp: 2019-01-23T11:15:00.000000000 risk_percent: 0.1 indicator: 109.63 signal_value: -2.0\n",
      "trade: KO 2019-01-23 11:30:00 qty: 2090.0 prc: 47.87 order: KO 2019-01-23 11:15:00 qty: 2090.0 enter long filled\n",
      "trade: PEP 2019-01-23 11:30:00 qty: -912.0 prc: 109.585 order: PEP 2019-01-23 11:15:00 qty: -912.0 enter short filled\n",
      "trade: KO 2019-01-25 11:15:00 qty: -2090.0 prc: 47.6225 order: KO 2019-01-25 11:00:00 qty: -2090.0 exit long filled\n",
      "trade: PEP 2019-01-25 11:15:00 qty: 912.0 prc: 110.337 order: PEP 2019-01-25 11:00:00 qty: 912.0 exit short filled\n",
      "order_qty: -2118.0 curr_equity: 998796.901 timestamp: 2019-01-28T09:45:00.000000000 risk_percent: 0.1 indicator: 47.155 signal_value: -2\n",
      "order_qty: 915.0 curr_equity: 998796.901 timestamp: 2019-01-28T09:45:00.000000000 risk_percent: 0.1 indicator: 109.2107 signal_value: 2.0\n",
      "trade: KO 2019-01-28 10:00:00 qty: -2118.0 prc: 47.0575 order: KO 2019-01-28 09:45:00 qty: -2118.0 enter short filled\n",
      "trade: PEP 2019-01-28 10:00:00 qty: 915.0 prc: 109.335 order: PEP 2019-01-28 09:45:00 qty: 915.0 enter long filled\n",
      "trade: KO 2019-01-28 12:45:00 qty: 2118.0 prc: 47.134 order: KO 2019-01-28 12:30:00 qty: 2118.0 exit short filled\n",
      "trade: PEP 2019-01-28 12:45:00 qty: -915.0 prc: 108.21 order: PEP 2019-01-28 12:30:00 qty: -915.0 exit long filled\n",
      "order_qty: -2100.0 curr_equity: 997605.499 timestamp: 2019-01-30T09:45:00.000000000 risk_percent: 0.1 indicator: 47.515 signal_value: -2\n",
      "order_qty: 908.0 curr_equity: 997605.499 timestamp: 2019-01-30T09:45:00.000000000 risk_percent: 0.1 indicator: 109.83 signal_value: 2.0\n",
      "trade: KO 2019-01-30 10:00:00 qty: -2100.0 prc: 47.4875 order: KO 2019-01-30 09:45:00 qty: -2100.0 enter short filled\n",
      "trade: PEP 2019-01-30 10:00:00 qty: 908.0 prc: 110.1 order: PEP 2019-01-30 09:45:00 qty: 908.0 enter long filled\n"
     ]
    }
   ],
   "source": [
    "%%checkall\n",
    "from __future__ import annotations\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import types\n",
    "import sys\n",
    "from collections import defaultdict\n",
    "from pprint import pformat\n",
    "import math\n",
    "\n",
    "from pyqstrat.evaluator import compute_return_metrics, display_return_metrics, plot_return_metrics\n",
    "from pyqstrat.account import Account\n",
    "from pyqstrat.pq_types import ContractGroup, Contract, Order, Trade, ReasonCode, MarketOrder, RoundTripTrade\n",
    "from pyqstrat.plot import TimeSeries, trade_sets_by_reason_code, Subplot, Plot, LinePlotAttributes, FilledLinePlotAttributes\n",
    "from pyqstrat.pq_utils import series_to_array, assert_\n",
    "from types import SimpleNamespace\n",
    "import matplotlib as mpl\n",
    "from typing import Callable, Any, Union\n",
    "from collections.abc import Sequence\n",
    "\n",
    "StrategyContextType = SimpleNamespace\n",
    "\n",
    "PriceFunctionType = Callable[[Contract, np.ndarray, int, StrategyContextType], float]\n",
    "\n",
    "IndicatorType = Union[np.ndarray, \n",
    "                      pd.Series,\n",
    "                      Callable[[ContractGroup, np.ndarray, SimpleNamespace, StrategyContextType], np.ndarray]]\n",
    "\n",
    "SignalType = Callable[[ContractGroup, np.ndarray, SimpleNamespace, SimpleNamespace, StrategyContextType], np.ndarray]\n",
    "\n",
    "RuleType = Callable[\n",
    "    [ContractGroup, \n",
    "     int, \n",
    "     np.ndarray, \n",
    "     SimpleNamespace, \n",
    "     np.ndarray, \n",
    "     Account, \n",
    "     StrategyContextType],\n",
    "    Sequence[Order]]\n",
    "\n",
    "MarketSimulatorType = Callable[\n",
    "    [Sequence[Order], \n",
    "     int, \n",
    "     np.ndarray, \n",
    "     dict[ContractGroup, SimpleNamespace], \n",
    "     dict[ContractGroup, SimpleNamespace], \n",
    "     SimpleNamespace],\n",
    "    Sequence[Trade]]\n",
    "\n",
    "DateRangeType = Union[tuple[str, str], tuple[np.datetime64, np.datetime64]]\n",
    "\n",
    "# Placeholder for orders to create later\n",
    "OrderTupType = tuple[RuleType, ContractGroup, dict[str, Any]]\n",
    "\n",
    "PlotPropertiesType = dict[str, dict[str, Any]]\n",
    "\n",
    "NAT = np.datetime64('NaT')\n",
    "\n",
    "\n",
    "def _get_time_series_list(timestamps: np.ndarray, \n",
    "                          names: Sequence[str], \n",
    "                          values: SimpleNamespace, \n",
    "                          properties: PlotPropertiesType | None) -> list[TimeSeries]:\n",
    "    '''\n",
    "    For plotting, create a list of TimeSeries objects from the arrays in the values object\n",
    "    '''\n",
    "    ts_list = []\n",
    "    for name in names:\n",
    "        line_type: str | None = None\n",
    "        color: str | None = None\n",
    "        if properties is not None and name in properties:\n",
    "            if 'line_type' in properties[name]: line_type = properties[name]['line_type']\n",
    "            if 'color' in properties[name]: color = properties[name]['color']\n",
    "        y = getattr(values, name)\n",
    "        if not len(y): continue\n",
    "        if y.dtype.type in [np.str_, np.object_, np.datetime64]: continue\n",
    "        attrib = LinePlotAttributes(line_type=line_type, color=color)\n",
    "        ts = TimeSeries(name, timestamps, y, display_attributes=attrib)\n",
    "        ts_list.append(ts)\n",
    "    return ts_list\n",
    "\n",
    "\n",
    "class Strategy:\n",
    "    def __init__(self, \n",
    "                 timestamps: np.ndarray,\n",
    "                 contract_groups: Sequence[ContractGroup],\n",
    "                 price_function: PriceFunctionType,\n",
    "                 starting_equity: float = 1.0e6, \n",
    "                 pnl_calc_time: int = 15 * 60 + 1,\n",
    "                 trade_lag: int = 0,\n",
    "                 run_final_calc: bool = True, \n",
    "                 strategy_context: StrategyContextType | None = None) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            timestamps (np.array of np.datetime64): The \"heartbeat\" of the strategy.  We will evaluate trading rules and \n",
    "                simulate the market at these times.\n",
    "            contract_groups: The contract groups we will potentially trade.\n",
    "            price_function: A function that returns the price of a contract at a given timestamp\n",
    "            starting_equity: Starting equity in Strategy currency.  Default 1.e6\n",
    "            pnl_calc_time: Time of day used to calculate PNL.  Default 15 * 60 (3 pm)\n",
    "            trade_lag: Number of bars you want between the order and the trade.  For example, if you think it will take\n",
    "                5 seconds to place your order in the market, and your bar size is 1 second, set this to 5.  Set this to 0 if you\n",
    "                want to execute your trade at the same time as you place the order, for example, if you have daily bars.  Default 0.\n",
    "            run_final_calc: If set, calculates unrealized pnl and net pnl as well as realized pnl when strategy is done.\n",
    "                If you don't need unrealized pnl, turn this off for faster run time. Default True\n",
    "            strategy_context: A storage class where you can store key / value pairs relevant to this strategy.\n",
    "                For example, you may have a pre-computed table of correlations that you use in the indicator or trade rule functions.  \n",
    "                If not set, the __init__ function will create an empty member strategy_context object that you can access.\n",
    "        '''\n",
    "        self.name = 'main'  # Set by portfolio when running multiple strategies\n",
    "        self.timestamps = timestamps\n",
    "        assert_(len(contract_groups) > 0 and isinstance(contract_groups[0], ContractGroup))\n",
    "        self.contract_groups = contract_groups\n",
    "        if strategy_context is None: strategy_context = types.SimpleNamespace()\n",
    "        self.strategy_context = strategy_context\n",
    "        self.account = Account(contract_groups, timestamps, price_function, strategy_context, starting_equity, pnl_calc_time)\n",
    "        assert_(trade_lag >= 0, f'trade_lag cannot be negative: {trade_lag}')\n",
    "        self.trade_lag = trade_lag\n",
    "        self.run_final_calc = run_final_calc\n",
    "        self.indicators: dict[str, IndicatorType] = {}\n",
    "        self.signals: dict[str, SignalType] = {}\n",
    "        self.signal_values: dict[ContractGroup, SimpleNamespace] = defaultdict(types.SimpleNamespace)\n",
    "        self.rule_names: list[str] = []\n",
    "        self.rules: dict[str, RuleType] = {}\n",
    "        self.position_filters: dict[str, str | None] = {}\n",
    "        self.rule_signals: dict[str, tuple[str, Sequence]] = {}\n",
    "        self.market_sims: list[MarketSimulatorType] = []\n",
    "        self._trades: list[Trade] = []\n",
    "        self._orders: list[Order] = []\n",
    "        self._open_orders: dict[int, list[Order]] = defaultdict(list)\n",
    "        self.indicator_deps: dict[str, list[str]] = {}\n",
    "        self.indicator_cgroups: dict[str, list[ContractGroup]] = {}\n",
    "        self.indicator_values: dict[ContractGroup, SimpleNamespace] = defaultdict(types.SimpleNamespace)\n",
    "        self.signal_indicator_deps: dict[str, list[str]] = {}\n",
    "        self.signal_deps: dict[str, list[str]] = {}\n",
    "        self.signal_cgroups: dict[str, list[ContractGroup]] = {}\n",
    "        self.trades_iter: list[list] = [[] for x in range(len(timestamps))]  # For debugging, we don't really need this as a member variable\n",
    "        \n",
    "    def add_indicator(self, \n",
    "                      name: str, \n",
    "                      indicator: IndicatorType, \n",
    "                      contract_groups: Sequence[ContractGroup] | None = None, \n",
    "                      depends_on: Sequence[str] | None = None) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            name: Name of the indicator\n",
    "            indicator:  A function that takes strategy timestamps and other indicators and returns a numpy array\n",
    "              containing indicator values.  The return array must have the same length as the timestamps object.\n",
    "              Can also be a numpy array or a pandas Series in which case we just store the values.\n",
    "            contract_groups: Contract groups that this indicator applies to.  If not set, it applies to all contract groups. Default None.\n",
    "            depends_on: Names of other indicators that we need to compute this indicator. Default None.\n",
    "        '''\n",
    "        self.indicators[name] = indicator\n",
    "        self.indicator_deps[name] = [] if depends_on is None else list(depends_on)\n",
    "        if contract_groups is None: contract_groups = self.contract_groups\n",
    "        if isinstance(indicator, np.ndarray) or isinstance(indicator, pd.Series):\n",
    "            indicator_values = series_to_array(indicator)\n",
    "            for contract_group in contract_groups:\n",
    "                setattr(self.indicator_values[contract_group], name, indicator_values)\n",
    "        self.indicator_cgroups[name] = list(contract_groups)\n",
    "        \n",
    "    def add_signal(self,\n",
    "                   name: str,\n",
    "                   signal_function: SignalType,\n",
    "                   contract_groups: Sequence[ContractGroup] | None = None,\n",
    "                   depends_on_indicators: Sequence[str] | None = None,\n",
    "                   depends_on_signals: Sequence[str] | None = None) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            name (str): Name of the signal\n",
    "            signal_function (function):  A function that takes timestamps and a dictionary of indicator value arrays and \n",
    "                returns a numpy array\n",
    "                containing signal values.  The return array must have the same length as the input timestamps\n",
    "            contract_groups (list of :obj:`ContractGroup`, optional): Contract groups that this signal applies to.  \n",
    "                If not set, it applies to all contract groups.  Default None.\n",
    "            depends_on_indicators (list of str, optional): Names of indicators that we need to compute this signal. Default None.\n",
    "            depends_on_signals (list of str, optional): Names of other signals that we need to compute this signal. Default None.\n",
    "        '''\n",
    "        self.signals[name] = signal_function\n",
    "        self.signal_indicator_deps[name] = [] if depends_on_indicators is None else list(depends_on_indicators)\n",
    "        self.signal_deps[name] = [] if depends_on_signals is None else list(depends_on_signals)\n",
    "        if contract_groups is None: contract_groups = self.contract_groups\n",
    "        self.signal_cgroups[name] = list(contract_groups)\n",
    "        \n",
    "    def add_rule(self, \n",
    "                 name: str, \n",
    "                 rule_function: RuleType, \n",
    "                 signal_name: str, \n",
    "                 sig_true_values: Sequence | None = None, \n",
    "                 position_filter: str | None = None) -> None:\n",
    "        '''Add a trading rule.  Trading rules are guaranteed to run in the order in which you add them.  For example, if you set trade_lag to 0,\n",
    "               and want to exit positions and re-enter new ones in the same bar, make sure you add the exit rule before you add the entry rule to the \n",
    "               strategy.\n",
    "        \n",
    "        Args:\n",
    "            name: Name of the trading rule\n",
    "            rule_function: A trading rule function that returns a list of Orders\n",
    "            signal_name: The strategy will call the trading rule function when the signal with this name matches sig_true_values\n",
    "            sig_true_values: If the signal value at a bar is equal to one of these values, \n",
    "                the Strategy will call the trading rule function.  Default [TRUE]\n",
    "            position_filter: Can be \"zero\", \"nonzero\" or None.  Zero rules are only triggered when the corresponding contract positions are 0\n",
    "                Nonzero rules are only triggered when the corresponding contract positions are non-zero.  \n",
    "                If not set, we don't look at position before triggering the rule. Default None\n",
    "        '''\n",
    "        \n",
    "        if sig_true_values is None: sig_true_values = [True]\n",
    "            \n",
    "        if name in self.rule_names:\n",
    "            raise Exception(f'Rule {name} already exists')\n",
    "        # Rules should be run in order\n",
    "        self.rule_names.append(name)\n",
    "        self.rule_signals[name] = (signal_name, sig_true_values)\n",
    "        self.rules[name] = rule_function\n",
    "        if position_filter is not None:\n",
    "            assert_(position_filter in ['zero', 'nonzero'])\n",
    "        self.position_filters[name] = position_filter\n",
    "        \n",
    "    def add_market_sim(self, market_sim_function: MarketSimulatorType) -> None:\n",
    "        '''Add a market simulator.  A market simulator is a function that takes orders as input and returns trades.'''\n",
    "        self.market_sims.append(market_sim_function)\n",
    "        \n",
    "    def run_indicators(self, \n",
    "                       indicator_names: Sequence[str] | None = None, \n",
    "                       contract_groups: Sequence[ContractGroup] | None = None, \n",
    "                       clear_all: bool = False) -> None:\n",
    "        '''Calculate values of the indicators specified and store them.\n",
    "        \n",
    "        Args:\n",
    "            indicator_names: list of indicator names.  If None (default) run all indicators\n",
    "            contract_groups: Contract group to run this indicator for.  If None (default), we run it for all contract groups.\n",
    "            clear_all: If set, clears all indicator values before running.  Default False.\n",
    "        '''\n",
    "        \n",
    "        if indicator_names is None: indicator_names = list(self.indicators.keys())\n",
    "        if contract_groups is None: contract_groups = self.contract_groups\n",
    "            \n",
    "        if clear_all: self.indicator_values = defaultdict(types.SimpleNamespace)\n",
    "            \n",
    "        ind_names = []\n",
    "            \n",
    "        for ind_name, cgroup_list in self.indicator_cgroups.items():\n",
    "            if len(set(contract_groups).intersection(cgroup_list)): ind_names.append(ind_name)\n",
    "                \n",
    "        indicator_names = list(set(ind_names).intersection(indicator_names))\n",
    "         \n",
    "        for cgroup in contract_groups:\n",
    "            cgroup_ind_namespace = self.indicator_values[cgroup]\n",
    "            for indicator_name in indicator_names:\n",
    "                # First run all parents\n",
    "                parent_names = self.indicator_deps[indicator_name]\n",
    "                for parent_name in parent_names:\n",
    "                    if cgroup in self.indicator_values and hasattr(cgroup_ind_namespace, parent_name): continue\n",
    "                    self.run_indicators([parent_name], [cgroup])\n",
    "                    \n",
    "                # Now run the actual indicator\n",
    "                if cgroup in self.indicator_values and hasattr(cgroup_ind_namespace, indicator_name): continue\n",
    "                indicator_function = self.indicators[indicator_name]\n",
    "                     \n",
    "                parent_values = types.SimpleNamespace()\n",
    "\n",
    "                for parent_name in parent_names:\n",
    "                    setattr(parent_values, parent_name, getattr(cgroup_ind_namespace, parent_name))\n",
    "                    \n",
    "                if isinstance(indicator_function, np.ndarray) or isinstance(indicator_function, pd.Series):\n",
    "                    indicator_values = indicator_function\n",
    "                else:\n",
    "                    indicator_values = indicator_function(cgroup, self.timestamps, parent_values, self.strategy_context)\n",
    "\n",
    "                setattr(cgroup_ind_namespace, indicator_name, series_to_array(indicator_values))\n",
    "                \n",
    "    def run_signals(self, \n",
    "                    signal_names: Sequence[str] | None = None, \n",
    "                    contract_groups: Sequence[ContractGroup] | None = None, \n",
    "                    clear_all: bool = False) -> None:\n",
    "        '''Calculate values of the signals specified and store them.\n",
    "        \n",
    "        Args:\n",
    "            signal_names: list of signal names.  If None (default) run all signals\n",
    "            contract_groups: Contract groups to run this signal for. If None (default), we run it for all contract groups.\n",
    "            clear_all: If set, clears all signal values before running.  Default False.\n",
    "        '''\n",
    "        if signal_names is None: signal_names = list(self.signals.keys())\n",
    "        if contract_groups is None: contract_groups = self.contract_groups\n",
    "            \n",
    "        if clear_all: self.signal_values = defaultdict(types.SimpleNamespace)\n",
    "            \n",
    "        sig_names = []\n",
    "        \n",
    "        for sig_name, cgroup_list in self.signal_cgroups.items():\n",
    "            if len(set(contract_groups).intersection(cgroup_list)): \n",
    "                sig_names.append(sig_name)\n",
    "                \n",
    "        signal_names = list(set(sig_names).intersection(signal_names))\n",
    "        \n",
    "        for cgroup in contract_groups:\n",
    "            for signal_name in signal_names:\n",
    "                if cgroup not in self.signal_cgroups[signal_name]: continue\n",
    "                # First run all parent signals\n",
    "                parent_names = self.signal_deps[signal_name]\n",
    "                for parent_name in parent_names:\n",
    "                    if cgroup in self.signal_values and hasattr(self.signal_values[cgroup], parent_name): continue\n",
    "                    self.run_signals([parent_name], [cgroup])\n",
    "                # Now run the actual signal\n",
    "                if cgroup in self.signal_values and hasattr(self.signal_values[cgroup], signal_name): continue\n",
    "                signal_function = self.signals[signal_name]\n",
    "                parent_values = types.SimpleNamespace()\n",
    "                for parent_name in parent_names:\n",
    "                    sig_vals = getattr(self.signal_values[cgroup], parent_name)\n",
    "                    setattr(parent_values, parent_name, sig_vals)\n",
    "                    \n",
    "                # Get indicators needed for this signal\n",
    "                indicator_values = types.SimpleNamespace()\n",
    "                for indicator_name in self.signal_indicator_deps[signal_name]:\n",
    "                    setattr(indicator_values, indicator_name, getattr(self.indicator_values[cgroup], indicator_name))\n",
    "                    \n",
    "                signal_output = signal_function(cgroup, self.timestamps, indicator_values, parent_values, self.strategy_context)\n",
    "                setattr(self.signal_values[cgroup], signal_name, series_to_array(signal_output))\n",
    "\n",
    "    def _generate_order_iterations(self, \n",
    "                                   rule_names: Sequence[str] | None = None, \n",
    "                                   contract_groups: Sequence[ContractGroup] | None = None, \n",
    "                                   start_date: np.datetime64 = NAT, \n",
    "                                   end_date: np.datetime64 = NAT) -> None:\n",
    "        '''\n",
    "        >>> class MockStrat:\n",
    "        ...    def __init__(self):\n",
    "        ...        self.timestamps = timestamps\n",
    "        ...        self.trade_lag = 0\n",
    "        ...        self.account = self\n",
    "        ...        self.rules = {'rule_a': rule_a, 'rule_b': rule_b}\n",
    "        ...        self.market_sims = {ibm: market_sim_ibm, aapl: market_sim_aapl}\n",
    "        ...        self.rule_signals = {'rule_a': ('sig_a', [1]), 'rule_b': ('sig_b', [1, -1])}\n",
    "        ...        self.signal_values = {ibm: types.SimpleNamespace(sig_a=np.array([0., 1., 1.]), \n",
    "        ...                                                   sig_b = np.array([0., 0., 0.]) ),\n",
    "        ...                               aapl: types.SimpleNamespace(sig_a=np.array([0., 0., 0.]), \n",
    "        ...                                                    sig_b=np.array([0., -1., -1])\n",
    "        ...                                                   )}\n",
    "        ...        self.signal_cgroups = {'sig_a': [ibm, aapl], 'sig_b': [ibm, aapl]}\n",
    "        ...        self.indicator_values = {ibm: types.SimpleNamespace(), aapl: types.SimpleNamespace()}\n",
    "        >>>\n",
    "        >>> def market_sim_aapl(): pass\n",
    "        >>> def market_sim_ibm(): pass\n",
    "        >>> def rule_a(): pass\n",
    "        >>> def rule_b(): pass\n",
    "        >>> timestamps = np.array(['2018-01-01', '2018-01-02', '2018-01-03'], dtype = 'M8[D]')\n",
    "        >>> rule_names = ['rule_a', 'rule_b']\n",
    "        >>> ContractGroup.clear()\n",
    "        >>> ibm = ContractGroup.create('IBM')\n",
    "        >>> aapl = ContractGroup.create('AAPL')\n",
    "        >>> contract_groups = [ibm, aapl]\n",
    "        >>> start_date = np.datetime64('2018-01-01')\n",
    "        >>> end_date = np.datetime64('2018-02-05')\n",
    "        >>> strategy = MockStrat()\n",
    "        >>> Strategy._generate_order_iterations(strategy, rule_names, contract_groups, start_date, end_date)\n",
    "        >>> orders_iter = strategy.orders_iter\n",
    "        >>> assert_(len(orders_iter[0]) == 0)\n",
    "        >>> assert_(len(orders_iter[1]) == 2)\n",
    "        >>> assert_(orders_iter[1][0][1] == ibm)\n",
    "        >>> assert_(orders_iter[1][1][1] == aapl)\n",
    "        >>> assert_(len(orders_iter[2]) == 2)\n",
    "        '''\n",
    "        _start_date, _end_date = np.datetime64(start_date), np.datetime64(end_date)\n",
    "        if rule_names is None: rule_names = self.rule_names\n",
    "        if contract_groups is None: contract_groups = self.contract_groups\n",
    "\n",
    "        num_timestamps = len(self.timestamps)\n",
    "        \n",
    "        # list of lists, i -> list of order tuple\n",
    "        orders_iter: list[list[OrderTupType]] = [[] for x in range(num_timestamps)]\n",
    "\n",
    "        for rule_name in rule_names:\n",
    "            rule_function = self.rules[rule_name]\n",
    "            for cgroup in contract_groups:\n",
    "                signal_name, sig_true_values = self.rule_signals[rule_name]\n",
    "                if cgroup not in self.signal_cgroups[signal_name]:\n",
    "                    # We don't need to call this rule for this contract group\n",
    "                    continue\n",
    "                sig_values = getattr(self.signal_values[cgroup], signal_name)\n",
    "                timestamps = self.timestamps\n",
    "\n",
    "                null_value = False if sig_values.dtype == np.dtype('bool') else np.nan\n",
    "                \n",
    "                if not np.isnat(_start_date):\n",
    "                    start_idx: int = np.searchsorted(timestamps, _start_date)  # type: ignore\n",
    "                    sig_values[0:start_idx] = null_value\n",
    "                    \n",
    "                if not np.isnat(_end_date):\n",
    "                    end_idx: int = np.searchsorted(timestamps, _end_date)  # type: ignore\n",
    "                    sig_values[end_idx:] = null_value\n",
    "\n",
    "                indices = np.nonzero(np.isin(sig_values[:num_timestamps], sig_true_values))[0]\n",
    "                \n",
    "                # Don't run rules on last index since we cannot fill any orders\n",
    "                if len(indices) and indices[-1] == len(sig_values) - 1 and self.trade_lag > 0: indices = indices[:-1] \n",
    "                indicator_values = self.indicator_values[cgroup]\n",
    "                iteration_params = {'indicator_values': indicator_values, 'signal_values': sig_values, 'rule_name': rule_name}\n",
    "                for idx in indices: orders_iter[idx].append((rule_function, cgroup, iteration_params))\n",
    "\n",
    "        self.orders_iter = orders_iter\n",
    "    \n",
    "    def run_rules(self, \n",
    "                  rule_names: Sequence[str] | None = None, \n",
    "                  contract_groups: Sequence[ContractGroup] | None = None, \n",
    "                  start_date: np.datetime64 = NAT,\n",
    "                  end_date: np.datetime64 = NAT) -> None:\n",
    "        '''\n",
    "        Run trading rules.\n",
    "        \n",
    "        Args:\n",
    "            rule_names: list of rule names.  If None (default) run all rules\n",
    "            contract_groups: Contract groups to run this rule for.  If None (default), we run it for all contract groups.\n",
    "            start_date: Run rules starting from this date. Default None \n",
    "            end_date: Don't run rules after this date.  Default None\n",
    "        '''\n",
    "        self._generate_order_iterations(rule_names, contract_groups, start_date, end_date)\n",
    "        \n",
    "        # Now we know which rules, contract groups need to be applied for each iteration, go through each iteration and apply them\n",
    "        # in the same order they were added to the strategy\n",
    "        for i in range(len(self.orders_iter)):\n",
    "            self._run_iteration(i)\n",
    "            \n",
    "        if self.run_final_calc:\n",
    "            self.account.calc(self.timestamps[-1])\n",
    "        \n",
    "    def _run_iteration(self, i: int) -> None:\n",
    "        \n",
    "        self._sim_market(i)\n",
    "        # Treat all orders as IOC, i.e. if the order was not executed, then its cancelled.\n",
    "        self._open_orders[i] = []\n",
    "        \n",
    "        rules = self.orders_iter[i]\n",
    "        \n",
    "        for (rule_function, contract_group, params) in rules:\n",
    "            orders = self._get_orders(i, rule_function, contract_group, params)\n",
    "            self._orders += orders\n",
    "            self._open_orders[i + self.trade_lag] += orders\n",
    "            # If the lag is 0, then run rules one by one, and after each rule, run market sim to generate trades and update\n",
    "            # positions.  For example, if we have a rule to exit a position and enter a new one, we should make sure \n",
    "            # positions are updated after the first rule before running the second rule.  If the lag is not 0, \n",
    "            # run all rules and collect the orders, we don't need to run market sim after each rule\n",
    "            if self.trade_lag == 0: self._sim_market(i)\n",
    "        # If we failed to fully execute any orders in this iteration, add them to the next iteration so we get another chance to execute\n",
    "        open_orders = self._open_orders.get(i)\n",
    "        if open_orders is not None and len(open_orders):\n",
    "            self._open_orders[i + 1] += open_orders\n",
    "            \n",
    "    def run(self) -> None:\n",
    "        self.run_indicators()\n",
    "        self.run_signals()\n",
    "        self.run_rules()\n",
    "        \n",
    "    def _get_orders(self, idx: int, rule_function: RuleType, contract_group: ContractGroup, params: dict[str, Any]) -> Sequence[Order]:\n",
    "        try:\n",
    "            indicator_values, signal_values, rule_name = params['indicator_values'], params['signal_values'], params['rule_name']\n",
    "            position_filter = self.position_filters[rule_name]\n",
    "            if position_filter is not None:\n",
    "                curr_pos = self.account.position(contract_group, self.timestamps[idx])\n",
    "                if position_filter == 'zero' and not math.isclose(curr_pos, 0): return []\n",
    "                if position_filter == 'nonzero' and math.isclose(curr_pos, 0): return []\n",
    "            orders = rule_function(contract_group, idx, self.timestamps, indicator_values, signal_values, self.account,\n",
    "                                   self.strategy_context)\n",
    "        except Exception as e:\n",
    "            raise type(e)(\n",
    "                f'Exception: {str(e)} at rule: {type(rule_function)} contract_group: {contract_group} index: {idx}'\n",
    "            ).with_traceback(sys.exc_info()[2])\n",
    "        return orders\n",
    "        \n",
    "    def _sim_market(self, i: int) -> None:\n",
    "        '''\n",
    "        Go through all open orders and run market simulators to generate a list of trades and return any orders that were not filled.\n",
    "        '''\n",
    "        open_orders = self._open_orders.get(i)\n",
    "        if open_orders is None or len(open_orders) == 0: return\n",
    "            \n",
    "        for market_sim_function in self.market_sims:\n",
    "            try:\n",
    "                trades = market_sim_function(open_orders, i, self.timestamps, self.indicator_values, self.signal_values, self.strategy_context)\n",
    "                if len(trades): self.account.add_trades(trades)\n",
    "                self._trades += trades\n",
    "            except Exception as e:\n",
    "                raise type(e)(f'Exception: {str(e)} at index: {i} function: {market_sim_function}').with_traceback(sys.exc_info()[2])\n",
    "            \n",
    "        self._open_orders[i] = [order for order in open_orders if order.status != 'filled']\n",
    "            \n",
    "    def df_data(self, \n",
    "                contract_groups: Sequence[ContractGroup] | None = None, \n",
    "                add_pnl: bool = True, \n",
    "                start_date: str | np.datetime64 = NAT, \n",
    "                end_date: str | np.datetime64 = NAT) -> pd.DataFrame:\n",
    "        '''\n",
    "        Add indicators and signals to end of market data and return as a pandas dataframe.\n",
    "        \n",
    "        Args:\n",
    "            contract_groups (list of:obj:`ContractGroup`, optional): list of contract groups to include.  All if set to None (default)\n",
    "            add_pnl: If True (default), include P&L columns in dataframe\n",
    "            start_date: string or numpy datetime64. Default None\n",
    "            end_date: string or numpy datetime64: Default None\n",
    "        '''\n",
    "        _start_date, _end_date = np.datetime64(start_date), np.datetime64(end_date)\n",
    "        if contract_groups is None: contract_groups = self.contract_groups\n",
    "            \n",
    "        timestamps = self.timestamps\n",
    "        \n",
    "        if not np.isnat(_start_date): timestamps = timestamps[timestamps >= _start_date]\n",
    "        if not np.isnat(_end_date): timestamps = timestamps[timestamps <= _end_date]\n",
    "            \n",
    "        dfs = []\n",
    "             \n",
    "        for contract_group in contract_groups:\n",
    "            df = pd.DataFrame({'timestamp': self.timestamps})\n",
    "            if add_pnl: \n",
    "                df_pnl = self.df_pnl(contract_group)\n",
    "                \n",
    "            indicator_values = self.indicator_values[contract_group]\n",
    "            \n",
    "            for k in sorted(indicator_values.__dict__):\n",
    "                name = k\n",
    "                # Avoid name collisions\n",
    "                if name in df.columns: name = name + '.ind'\n",
    "                df.insert(len(df.columns), name, getattr(indicator_values, k))\n",
    "\n",
    "            signal_values = self.signal_values[contract_group]\n",
    "\n",
    "            for k in sorted(signal_values.__dict__):\n",
    "                name = k\n",
    "                if name in df.columns: name = name + '.sig'\n",
    "                df.insert(len(df.columns), name, getattr(signal_values, k))\n",
    "                \n",
    "            if add_pnl: df = pd.merge(df, df_pnl, on=['timestamp'], how='left')\n",
    "            # Add counter column for debugging\n",
    "            df.insert(len(df.columns), 'i', np.arange(len(df)))\n",
    "            \n",
    "            dfs.append(df)\n",
    "            \n",
    "        return pd.concat(dfs)\n",
    "    \n",
    "    def trades(self, \n",
    "               contract_group: ContractGroup | None = None, \n",
    "               start_date: np.datetime64 = NAT, \n",
    "               end_date: np.datetime64 = NAT) -> list[Trade]:\n",
    "        '''Returns a list of trades with the given contract group and with trade date between (and including) start date \n",
    "            and end date if they are specified.\n",
    "            If contract_group is None trades for all contract_groups are returned'''\n",
    "        return self.account.trades(contract_group, start_date, end_date)\n",
    "    \n",
    "    def roundtrip_trades(self,\n",
    "                         contract_group: ContractGroup | None = None, \n",
    "                         start_date: np.datetime64 = NAT, \n",
    "                         end_date: np.datetime64 = NAT) -> list[RoundTripTrade]:\n",
    "        '''Returns a list of trades with the given contract group and with trade date between (and including) start date \n",
    "            and end date if they are specified.\n",
    "            If contract_group is None trades for all contract_groups are returned'''\n",
    "        return self.account.roundtrip_trades(contract_group, start_date, end_date)\n",
    "    \n",
    "    def df_trades(self, \n",
    "                  contract_group: ContractGroup | None = None, \n",
    "                  start_date: np.datetime64 = NAT, \n",
    "                  end_date: np.datetime64 = NAT) -> pd.DataFrame:\n",
    "        '''Returns a dataframe with data from trades with the given contract group and with trade date between (and including)\n",
    "            start date and end date\n",
    "            if they are specified. If contract_group is None trades for all contract_groups are returned'''\n",
    "        return self.account.df_trades(contract_group, start_date, end_date)\n",
    "    \n",
    "    def df_roundtrip_trades(self, \n",
    "                            contract_group: ContractGroup | None = None, \n",
    "                            start_date: np.datetime64 = NAT, \n",
    "                            end_date: np.datetime64 = NAT) -> pd.DataFrame:\n",
    "        '''Returns a dataframe of round trip trades with the given contract group and with trade date \n",
    "            between (and including) start date and end date if they are specified. If contract_group is None trades for all \n",
    "            contract_groups are returned'''\n",
    "        return self.account.df_roundtrip_trades(contract_group, start_date, end_date)\n",
    "\n",
    "    def orders(self, \n",
    "               contract_group: ContractGroup | None = None, \n",
    "               start_date: np.datetime64 | str | None = None, \n",
    "               end_date: np.datetime64 | str | None = None) -> Sequence[Order]:\n",
    "        '''Returns a list of orders with the given contract group and with order date between (and including) start date and \n",
    "            end date if they are specified.\n",
    "            If contract_group is None orders for all contract_groups are returned'''\n",
    "        orders: list[Order] = []\n",
    "        _start_date: np.datetime64 = np.datetime64(start_date)\n",
    "        _end_date: np.datetime64 = np.datetime64(end_date)\n",
    "        if contract_group is None:\n",
    "            orders += [order for order in self._orders if (\n",
    "                np.isnat(_start_date) or np.datetime64(order.timestamp) >= _start_date) and (\n",
    "                np.isnat(_end_date) or np.datetime64(order.timestamp) <= _end_date)]\n",
    "        else:\n",
    "            for contract in contract_group.contracts:\n",
    "                orders += [order for order in self._orders if (contract is None or order.contract == contract) and (\n",
    "                    np.isnat(_start_date) or np.datetime64(order.timestamp) >= _start_date) and (\n",
    "                    np.isnat(_end_date) or np.datetime64(order.timestamp) <= _end_date)]\n",
    "        return orders\n",
    "    \n",
    "    def df_orders(self, \n",
    "                  contract_group: ContractGroup | None = None, \n",
    "                  start_date: np.datetime64 | str = NAT, \n",
    "                  end_date: np.datetime64 | str = NAT) -> pd.DataFrame:\n",
    "        '''Returns a dataframe with data from orders with the given contract group and with order date between (and including) \n",
    "            start date and end date\n",
    "            if they are specified. If contract_group is None orders for all contract_groups are returned'''\n",
    "        orders = self.orders(contract_group, start_date, end_date)\n",
    "        order_records = [(order.contract.symbol if order.contract else '',\n",
    "                          type(order).__name__, order.timestamp, order.qty, \n",
    "                          order.reason_code, \n",
    "                          (str(order.properties.__dict__) if order.properties.__dict__ else ''),\n",
    "                          (str(order.contract.properties.__dict__) \n",
    "                           if order.contract and order.contract.properties.__dict__ else '')) for order in orders]\n",
    "        df_orders = pd.DataFrame.from_records(order_records,\n",
    "                                              columns=['symbol', 'type', 'timestamp', 'qty', 'reason_code', 'order_props', 'contract_props'])\n",
    "        return df_orders\n",
    "    \n",
    "    def df_pnl(self, contract_group=None) -> pd.DataFrame:\n",
    "        '''Returns a dataframe with P&L columns.  If contract group is set to None (default), sums up P&L across all contract groups'''\n",
    "        return self.account.df_account_pnl(contract_group)\n",
    "    \n",
    "    def df_returns(self, \n",
    "                   contract_group: ContractGroup | None = None,\n",
    "                   sampling_frequency: str = 'D') -> pd.DataFrame:\n",
    "        '''Return a dataframe of returns and equity indexed by date.\n",
    "        \n",
    "        Args:\n",
    "            contract_group: The contract group to get returns for.  \n",
    "                If set to None (default), we return the sum of PNL for all contract groups\n",
    "            sampling_frequency: Downsampling frequency.  Default is None.  See pandas frequency strings for possible values\n",
    "        '''\n",
    "        pnl = self.df_pnl(contract_group)[['timestamp', 'net_pnl', 'equity']]\n",
    "\n",
    "        pnl.equity = pnl.equity.ffill()\n",
    "        pnl = pnl.set_index('timestamp').resample(sampling_frequency).last().reset_index()\n",
    "        pnl = pnl.dropna(subset=['equity'])\n",
    "        pnl['ret'] = pnl.equity.pct_change()\n",
    "        return pnl\n",
    "    \n",
    "    def plot(self, \n",
    "             contract_groups: Sequence[ContractGroup] | None = None, \n",
    "             primary_indicators: Sequence[str] | None = None,\n",
    "             primary_indicators_dual_axis: Sequence[str] | None = None,\n",
    "             secondary_indicators: Sequence[str] | None = None,\n",
    "             secondary_indicators_dual_axis: Sequence[str] | None = None,\n",
    "             indicator_properties: PlotPropertiesType | None = None,\n",
    "             signals: Sequence[str] | None = None,\n",
    "             signal_properties: PlotPropertiesType | None = None, \n",
    "             pnl_columns: Sequence[str] | None = None, \n",
    "             title: str | None = None, \n",
    "             figsize: tuple[int, int] = (20, 15), \n",
    "             date_range: DateRangeType = (NAT, NAT),\n",
    "             date_format: str | None = None, \n",
    "             sampling_frequency: str | None = None, \n",
    "             trade_marker_properties: PlotPropertiesType | None = None, \n",
    "             hspace: float = 0.15) -> None:\n",
    "        \n",
    "        '''\n",
    "        Plot indicators, signals, trades, position, pnl\n",
    "        \n",
    "        Args:\n",
    "            contract_groups: Contract groups to plot or None (default) for all contract groups. \n",
    "            primary indicators: list of indicators to plot in the main indicator section. Default None (plot everything)\n",
    "            primary indicators: list of indicators to plot in the secondary indicator section. Default None (don't plot anything)\n",
    "            indicator_properties: If set, we use the line color, line type indicated for the given indicators\n",
    "            signals: Signals to plot.  Default None (plot everything).\n",
    "            plot_equity: If set, we plot the equity curve.  Default is True\n",
    "            title: Title of plot. Default None\n",
    "            figsize: Figure size.  Default (20, 15)\n",
    "            date_range: Used to restrict the date range of the graph. Default None\n",
    "            date_format: Date format for tick labels on x axis.  If set to None (default), will be selected based on date range. \n",
    "                See matplotlib date format strings\n",
    "            sampling_frequency: Downsampling frequency.  The graph may get too busy if you have too many bars of data, \n",
    "                so you may want to downsample before plotting.  See pandas frequency strings for possible values. Default None.\n",
    "            trade_marker_properties: A dictionary of order reason code -> marker shape, marker size, marker color for plotting trades\n",
    "                with different reason codes. By default we use the dictionary from the :obj:`ReasonCode` class\n",
    "            hspace: Height (vertical) space between subplots.  Default is 0.15\n",
    "        '''\n",
    "        if contract_groups is None: contract_groups = self.contract_groups\n",
    "        if isinstance(contract_groups, ContractGroup): contract_groups = [contract_groups]\n",
    "        if pnl_columns is None: pnl_columns = ['equity']\n",
    "        \n",
    "        for contract_group in contract_groups:\n",
    "            primary_indicator_names = [ind_name for ind_name in self.indicator_values[contract_group].__dict__\n",
    "                                       if hasattr(self.indicator_values[contract_group], ind_name)]\n",
    "            if primary_indicators:\n",
    "                primary_indicator_names = list(set(primary_indicator_names).intersection(primary_indicators))\n",
    "            secondary_indicator_names: list[str] = []\n",
    "            if secondary_indicators:\n",
    "                secondary_indicator_names = list(secondary_indicators)\n",
    "            signal_names = [sig_name for sig_name in self.signals.keys() if hasattr(self.signal_values[contract_group], sig_name)]\n",
    "            if signals:\n",
    "                signal_names = list(set(signal_names).intersection(signals))\n",
    "                \n",
    "            primary_indicator_list = _get_time_series_list(self.timestamps, primary_indicator_names, \n",
    "                                                           self.indicator_values[contract_group], indicator_properties)\n",
    "            secondary_indicator_list = _get_time_series_list(self.timestamps, secondary_indicator_names, \n",
    "                                                             self.indicator_values[contract_group], indicator_properties)\n",
    "            signal_list = _get_time_series_list(self.timestamps, signal_names, self.signal_values[contract_group], signal_properties)\n",
    "            df_pnl_ = self.df_pnl(contract_group)\n",
    "            pnl_list = [TimeSeries(pnl_column, \n",
    "                                   timestamps=df_pnl_.timestamp.values, \n",
    "                                   values=df_pnl_[pnl_column].values) for pnl_column in pnl_columns]\n",
    "            \n",
    "            trades = [trade for trade in self._trades if trade.order.contract and trade.order.contract.contract_group == contract_group]\n",
    "            if trade_marker_properties:\n",
    "                trade_sets = trade_sets_by_reason_code(trades, trade_marker_properties, remove_missing_properties=True)\n",
    "            else:\n",
    "                trade_sets = trade_sets_by_reason_code(trades)\n",
    "                \n",
    "            primary_indicator_subplot = Subplot(\n",
    "                primary_indicator_list + trade_sets,  # type: ignore # mypy does not allow adding heterogeneous lists\n",
    "                secondary_y=primary_indicators_dual_axis,\n",
    "                height_ratio=0.5, \n",
    "                ylabel='Primary Indicators')\n",
    "            \n",
    "            if len(secondary_indicator_list):\n",
    "                secondary_indicator_subplot = Subplot(secondary_indicator_list, \n",
    "                                                      secondary_y=secondary_indicators_dual_axis,\n",
    "                                                      height_ratio=0.5, \n",
    "                                                      ylabel='Secondary Indicators')\n",
    "            signal_subplot = Subplot(signal_list, ylabel='Signals', height_ratio=0.167)\n",
    "            pnl_subplot = Subplot(pnl_list, ylabel='Equity', height_ratio=0.167, log_y=True, y_tick_format='${x:,.0f}')\n",
    "            position = df_pnl_.position.values\n",
    "            disp_attribs = FilledLinePlotAttributes()\n",
    "            pos_subplot = Subplot(\n",
    "                [TimeSeries('position', timestamps=df_pnl_.timestamp, values=position, display_attributes=disp_attribs)], \n",
    "                ylabel='Position', height_ratio=0.167)\n",
    "            \n",
    "            title_full = title\n",
    "            if len(contract_groups) > 1:\n",
    "                if title is None: title = ''\n",
    "                title_full = f'{title} {contract_group.name}'\n",
    "                \n",
    "            plot_list = []\n",
    "            if len(primary_indicator_list): plot_list.append(primary_indicator_subplot)\n",
    "            if len(secondary_indicator_list): plot_list.append(secondary_indicator_subplot)\n",
    "            if len(signal_list): plot_list.append(signal_subplot)\n",
    "            if len(position): plot_list.append(pos_subplot)\n",
    "            if len(pnl_list): plot_list.append(pnl_subplot)\n",
    "            \n",
    "            if not len(plot_list): return\n",
    "            \n",
    "            _date_range = (np.datetime64(date_range[0]), np.datetime64(date_range[1]))\n",
    "                \n",
    "            plot = Plot(plot_list, \n",
    "                        figsize=figsize, \n",
    "                        date_range=_date_range, \n",
    "                        date_format=date_format, \n",
    "                        sampling_frequency=sampling_frequency, \n",
    "                        title=title_full, \n",
    "                        hspace=hspace)\n",
    "            plot.draw()\n",
    "            \n",
    "    def evaluate_returns(self, \n",
    "                         contract_group: ContractGroup | None = None, \n",
    "                         plot: bool = True, \n",
    "                         display_summary: bool = True, \n",
    "                         float_precision: int = 4, \n",
    "                         return_metrics: bool = False) -> dict[str, Any] | None:\n",
    "        '''Returns a dictionary of common return metrics.\n",
    "        \n",
    "        Args:\n",
    "            contract_group (:obj:`ContractGroup`, optional): Contract group to evaluate or None (default) for all contract groups\n",
    "            plot (bool): If set to True, display plots of equity, drawdowns and returns.  Default False\n",
    "            float_precision (float, optional): Number of significant figures to show in returns.  Default 4\n",
    "            return_metrics (bool, optional): If set, we return the computed metrics as a dictionary\n",
    "        '''\n",
    "        returns = self.df_returns(contract_group)\n",
    "        ev = compute_return_metrics(returns.timestamp.values, returns.ret.values, self.account.starting_equity)\n",
    "        if display_summary:\n",
    "            display_return_metrics(ev.metrics(), float_precision=float_precision)\n",
    "        if plot: \n",
    "            plot_return_metrics(ev.metrics())\n",
    "        if return_metrics:\n",
    "            return ev.metrics()\n",
    "        return None\n",
    "    \n",
    "    def plot_returns(self, contract_group: ContractGroup | None = None) -> tuple[mpl.figure.Figure, mpl.axes.Axes] | None:\n",
    "        '''Display plots of equity, drawdowns and returns for the given contract group or for all contract groups if contract_group \n",
    "            is None (default)'''\n",
    "        if contract_group is None:\n",
    "            returns = self.df_returns()\n",
    "        else:\n",
    "            returns = self.df_returns(contract_group)\n",
    "\n",
    "        ev = compute_return_metrics(returns.timestamp.values, returns.ret.values, self.account.starting_equity)\n",
    "        return plot_return_metrics(ev.metrics())\n",
    "       \n",
    "    def __repr__(self):\n",
    "        return f'{pformat(self.indicators)} {pformat(self.rules)} {pformat(self.account)}'\n",
    "    \n",
    "\n",
    "def test_strategy() -> None:\n",
    "    import math\n",
    "    import numpy as np\n",
    "    import pandas as pd\n",
    "    import os\n",
    "    from types import SimpleNamespace\n",
    "    from pyqstrat.pq_types import Contract, ContractGroup, Trade\n",
    "\n",
    "    try:\n",
    "        # If we are running from unit tests\n",
    "        ko_file_path = os.path.dirname(os.path.realpath(__file__)) + '/notebooks/support/coke_15_min_prices.csv.gz'\n",
    "        pep_file_path = os.path.dirname(os.path.realpath(__file__)) + '/notebooks/support/pepsi_15_min_prices.csv.gz' \n",
    "    except NameError:\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",
    "    end_time = '2019-01-30 12:00'\n",
    "    \n",
    "    ko_prices = ko_prices.query(f'timestamp <= \"{end_time}\"')\n",
    "    pep_prices = pep_prices.query(f'timestamp <= \"{end_time}\"')\n",
    "\n",
    "    timestamps = ko_prices.timestamp.values\n",
    "    \n",
    "    ratio = ko_prices.c / pep_prices.c\n",
    "    \n",
    "    def zscore_indicator(contract_group: ContractGroup, \n",
    "                         timestamps: np.ndarray, \n",
    "                         indicators: SimpleNamespace,\n",
    "                         strategy_context: StrategyContextType) -> np.ndarray:  # 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",
    "    def pair_strategy_signal(contract_group: ContractGroup,\n",
    "                             timestamps: np.ndarray,\n",
    "                             indicators: SimpleNamespace, \n",
    "                             parent_signals: SimpleNamespace,\n",
    "                             strategy_context: StrategyContextType) -> np.ndarray: \n",
    "        # We don't need any indicators since the zscore is already part of the market data\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",
    "    def pair_entry_rule(contract_group: ContractGroup,\n",
    "                        i: int,\n",
    "                        timestamps: np.ndarray,\n",
    "                        indicators: SimpleNamespace,\n",
    "                        signal: np.ndarray,\n",
    "                        account: Account,\n",
    "                        strategy_context: StrategyContextType) -> Sequence[Order]:\n",
    "        timestamp = timestamps[i]\n",
    "        assert_(math.isclose(account.position(contract_group, timestamp), 0))\n",
    "        signal_value = signal[i]\n",
    "        risk_percent = 0.1\n",
    "\n",
    "        orders = []\n",
    "        \n",
    "        symbol = contract_group.name\n",
    "        contract = contract_group.get_contract(symbol)\n",
    "        if contract is None: contract = Contract.create(symbol, contract_group=contract_group)\n",
    "        \n",
    "        curr_equity = account.equity(timestamp)\n",
    "        order_qty = np.round(curr_equity * risk_percent / indicators.c[i] * np.sign(signal_value))\n",
    "        print(f'order_qty: {order_qty} curr_equity: {curr_equity} timestamp: {timestamp}'\n",
    "              f' risk_percent: {risk_percent} indicator: {indicators.c[i]} signal_value: {signal_value}')\n",
    "        reason_code = ReasonCode.ENTER_LONG if order_qty > 0 else ReasonCode.ENTER_SHORT\n",
    "        orders.append(MarketOrder(contract, timestamp, order_qty, reason_code=reason_code))\n",
    "        return orders\n",
    "            \n",
    "    def pair_exit_rule(contract_group: ContractGroup,\n",
    "                       i: int,\n",
    "                       timestamps: np.ndarray,\n",
    "                       indicators: SimpleNamespace,\n",
    "                       signal: np.ndarray,\n",
    "                       account: Account,\n",
    "                       strategy_context: StrategyContextType) -> Sequence[Order]:\n",
    "        timestamp = timestamps[i]\n",
    "        curr_pos = account.position(contract_group, timestamp)\n",
    "        assert_(not math.isclose(curr_pos, 0))\n",
    "        signal_value = signal[i]\n",
    "        orders = []\n",
    "        symbol = contract_group.name\n",
    "        contract = contract_group.get_contract(symbol)\n",
    "        if contract is None: contract = Contract.create(symbol, contract_group=contract_group)\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 = ReasonCode.EXIT_LONG if order_qty < 0 else ReasonCode.EXIT_SHORT\n",
    "            orders.append(MarketOrder(contract, timestamp, order_qty, reason_code=reason_code))\n",
    "        return orders\n",
    "\n",
    "    def market_simulator(orders: Sequence[Order],\n",
    "                         i: int,\n",
    "                         timestamps: np.ndarray,\n",
    "                         indicators: dict[ContractGroup, SimpleNamespace],\n",
    "                         signals: dict[ContractGroup, SimpleNamespace],\n",
    "                         strategy_context: StrategyContextType) -> Sequence[Trade]:\n",
    "        trades = []\n",
    "\n",
    "        timestamp = timestamps[i]\n",
    "\n",
    "        for order in orders:\n",
    "            trade_price = np.nan\n",
    "            \n",
    "            assert order.contract is not None\n",
    "            cgroup = order.contract.contract_group\n",
    "            ind = indicators[cgroup]\n",
    "            \n",
    "            o, h, l = ind.o[i], ind.h[i], ind.l[i]  # noqa: E741  # l is ambiguous\n",
    "\n",
    "            assert_(isinstance(order, MarketOrder), f'Unexpected order type: {order}')\n",
    "            trade_price = 0.5 * (o + h) if order.qty > 0 else 0.5 * (o + l)\n",
    "\n",
    "            if np.isnan(trade_price): continue\n",
    "\n",
    "            trade = Trade(order.contract, order, timestamp, order.qty, trade_price, commission=0, fee=0)\n",
    "            order.status = 'filled'\n",
    "            print(f'trade: {trade}')\n",
    "\n",
    "            trades.append(trade)\n",
    "\n",
    "        return trades\n",
    "    \n",
    "    def get_price(contract: Contract, timestamps: np.ndarray, i: int, strategy_context: StrategyContextType) -> float:\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 contract: {contract}')\n",
    "\n",
    "    Contract.clear()\n",
    "    ContractGroup.clear()\n",
    "        \n",
    "    ko_contract_group = ContractGroup.create('KO')\n",
    "    pep_contract_group = ContractGroup.create('PEP')\n",
    "\n",
    "    strategy_context = SimpleNamespace(ko_price=ko_prices.c.values, pep_price=pep_prices.c.values)\n",
    "\n",
    "    strategy = 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], contract_groups=[tup[0]])\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_entry_rule', pair_entry_rule, \n",
    "                      signal_name='pair_strategy_signal', sig_true_values=[-2, 2], position_filter='zero')\n",
    "    \n",
    "    strategy.add_rule('pair_exit_rule', pair_exit_rule, \n",
    "                      signal_name='pair_strategy_signal', sig_true_values=[-1, 1], position_filter='nonzero')\n",
    "    \n",
    "    strategy.add_market_sim(market_simulator)\n",
    "    \n",
    "    strategy.run_indicators()\n",
    "    strategy.run_signals()\n",
    "    strategy.run_rules()\n",
    "\n",
    "    metrics = strategy.evaluate_returns(plot=False, display_summary=False, return_metrics=True)\n",
    "    assert metrics is not None\n",
    "    assert round(metrics['gmean'], 6) == -0.062878\n",
    "    assert round(metrics['sharpe'], 4) == -9.7079  # -7.2709)\n",
    "    assert round(metrics['mdd_pct'], 6) == 0.002574  # -0.002841)\n",
    "\n",
    "\n",
    "def test_strategy_2() -> None:\n",
    "    '''Test of a dummy strategy'''\n",
    "    \n",
    "    def test_signal(contract_group: ContractGroup,\n",
    "                    timestamps: np.ndarray,\n",
    "                    indicators: SimpleNamespace, \n",
    "                    parent_signals: SimpleNamespace,\n",
    "                    strategy_context: StrategyContextType) -> np.ndarray: \n",
    "        # We don't need any indicators since the zscore is already part of the market data\n",
    "        return np.full(len(timestamps), True)\n",
    "    \n",
    "    def test_rule(contract_group: ContractGroup,\n",
    "                  i: int,\n",
    "                  timestamps: np.ndarray,\n",
    "                  indicators: SimpleNamespace,\n",
    "                  signal: np.ndarray,\n",
    "                  account: Account,\n",
    "                  strategy_context: StrategyContextType) -> Sequence[Order]:\n",
    "        \n",
    "        contract = contract_group.get_contract('TEST')\n",
    "        return [MarketOrder(contract, timestamps[i], 10, reason_code='ENTER')]\n",
    "            \n",
    "    def market_simulator(orders: Sequence[Order],\n",
    "                         i: int,\n",
    "                         timestamps: np.ndarray,\n",
    "                         indicators: dict[ContractGroup, SimpleNamespace],\n",
    "                         signals: dict[ContractGroup, SimpleNamespace],\n",
    "                         strategy_context: StrategyContextType) -> Sequence[Trade]:\n",
    "        trades = []\n",
    "        for order in orders:\n",
    "            assert order.contract is not None\n",
    "            trade = Trade(order.contract, order, timestamps[i], order.qty, 50)\n",
    "            trades.append(trade)\n",
    "        return trades\n",
    "    \n",
    "    def get_price(contract: Contract, timestamps: np.ndarray, i: int, strategy_context: StrategyContextType) -> float:\n",
    "        return 11\n",
    "                \n",
    "    timestamps = np.arange(np.datetime64('2018-01-05T08:00'), np.datetime64('2018-01-05T08:05'))\n",
    "    prices = np.array([8.9, 8.10, 8.11, 8.12, 8.13])\n",
    "    df = pd.DataFrame({'timestamp': timestamps, 'price': prices})\n",
    "    \n",
    "    ContractGroup.clear()\n",
    "    Contract.clear()\n",
    "    test_cg = ContractGroup.create('TEST')\n",
    "    contract = Contract.create('TEST', test_cg)\n",
    "    test_cg.add_contract(contract)\n",
    "    strategy = Strategy(timestamps, [test_cg], get_price)\n",
    "\n",
    "    strategy = Strategy(timestamps, [test_cg], get_price, trade_lag=0)\n",
    "    strategy.add_indicator('price', df.price.values)\n",
    "    strategy.add_signal('test_sig', test_signal, None, ['price'])\n",
    "    strategy.add_rule('test_rule', test_rule, signal_name='test_sig')\n",
    "\n",
    "    strategy.add_market_sim(market_simulator)\n",
    "    strategy.run()\n",
    "    \n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_strategy()\n",
    "    test_strategy_2()\n",
    "    import doctest\n",
    "    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
