{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-11-07T20:05:47.032440Z",
     "start_time": "2018-11-07T20:05:45.741825Z"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import math\n",
    "from pyqstrat.pq_utils import *\n",
    "from pyqstrat.strategy import Strategy\n",
    "from pyqstrat.account import Account\n",
    "from pyqstrat.portfolio import Portfolio\n",
    "from pyqstrat.evaluator import plot_return_metrics\n",
    "from pyqstrat.pq_types import Contract, Trade\n",
    "from pyqstrat.marketdata import MarketData, MarketDataCollection\n",
    "from pyqstrat.orders import StopLimitOrder, MarketOrder\n",
    "\n",
    "set_defaults()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-11-07T20:05:51.904135Z",
     "start_time": "2018-11-07T20:05:47.035457Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running\n",
      "done\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 74864 characters]\n",
      "text/plain": [
       "<Figure size 1080x576 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 196612 characters]\n",
      "text/plain": [
       "<Figure size 1440x1080 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def sim_order(order, i, date, md):\n",
    "    symbol = order.symbol\n",
    "    trade_price = np.nan\n",
    "    skid_fraction = 0.5\n",
    "\n",
    "    if not md.valid_row(i): return None\n",
    "\n",
    "    if isinstance(order, MarketOrder):\n",
    "        if order.qty > 0:\n",
    "            trade_price = 0.5 * (md.c[i] + md.h[i])\n",
    "        else:\n",
    "            trade_price = 0.5 * (md.c[i] + md.l[i])\n",
    "    elif order.qty > 0 and md.h[i] > order.trigger_price:\n",
    "        trade_price = skid_fraction * max(md.o[i], order.trigger_price, md.l[i]) + (1 - skid_fraction) * md.h[i]\n",
    "    elif order.qty < 0 and md.l[i] < order.trigger_price:\n",
    "        trade_price = skid_fraction * min(md.o[i], order.trigger_price, md.h[i]) + (1 - skid_fraction) * md.l[i]\n",
    "    else:\n",
    "        return None\n",
    "    sim_trade = Trade(symbol, date, order.qty, trade_price, order = order)\n",
    "    return sim_trade\n",
    "\n",
    "def market_simulator(strategy, orders, i, date, marketdata):\n",
    "    trades = []\n",
    "    for order in orders:\n",
    "        sim_trade = sim_order(order, i, date, marketdata)\n",
    "        if sim_trade is None: continue\n",
    "        if math.isclose(sim_trade.qty, order.qty): order.status = 'filled'\n",
    "\n",
    "        trades.append(sim_trade)\n",
    "    return trades\n",
    "\n",
    "def trade_rule(strategy, symbol, i, date, marketdata, indicator_values, signal_values, account):\n",
    "    heat = 0.05\n",
    "    reason_code = None\n",
    "\n",
    "    if not marketdata.valid_row(i): return []\n",
    "\n",
    "    curr_pos = account.position(symbol, date)\n",
    "\n",
    "    if i == len(marketdata.dates) - 2: # Last date so get out of position\n",
    "        if not math.isclose(curr_pos, 0): \n",
    "            return [MarketOrder(symbol, date, -curr_pos, reason_code = ReasonCode.BACKTEST_END)]\n",
    "        else:\n",
    "            return []\n",
    "\n",
    "    trend = signal_values[i]\n",
    "    fast_resistance, fast_support, slow_resistance, slow_support = (indicator_values['fast_resistance'][i], \n",
    "            indicator_values['fast_support'][i], indicator_values['slow_resistance'][i], indicator_values['slow_support'][i])\n",
    "\n",
    "    if trend == 1:\n",
    "        entry_limit = fast_resistance\n",
    "        stop_limit = fast_support\n",
    "    elif trend == -1:\n",
    "        entry_limit = fast_support\n",
    "        stop_limit = fast_resistance\n",
    "    else:\n",
    "        return []\n",
    "\n",
    "    if math.isclose(curr_pos, 0): # We got a trade in the previous bar so put in a stop limit order\n",
    "        if math.isclose(entry_limit, stop_limit): return []\n",
    "        curr_equity = account.equity(date)\n",
    "        order_qty = curr_equity * heat / (entry_limit - stop_limit)\n",
    "        trigger_price = entry_limit\n",
    "        reason_code = ReasonCode.ENTER_LONG if order_qty > 0 else ReasonCode.ENTER_SHORT\n",
    "    else:\n",
    "        order_qty = -curr_pos\n",
    "        trigger_price = stop_limit\n",
    "        reason_code = ReasonCode.EXIT_LONG if order_qty < 0 else ReasonCode.EXIT_SHORT\n",
    "\n",
    "    order_qty = round(order_qty)\n",
    "\n",
    "    if np.isnan(order_qty):\n",
    "        raise Exception(f'Got nan order qty date: {date} i: {i} curr_pos: {curr_pos} curr_equity: {curr_equity} entry_limit: {entry_limit} stop_limit: {stop_limit}')\n",
    "\n",
    "    if math.isclose(order_qty, 0): return []\n",
    "\n",
    "    order = StopLimitOrder(symbol, date, order_qty, trigger_price, reason_code = reason_code)\n",
    "\n",
    "    return [order]\n",
    "\n",
    "def get_support(lows, n): return pd.Series(lows).rolling(window = n, min_periods = 1).min().values\n",
    "\n",
    "def get_resistance(highs, n): return pd.Series(highs).rolling(window = n, min_periods = 1).max().values\n",
    "\n",
    "def get_trend(md, ind):\n",
    "    trend = pd.Series(np.where(pd.Series(md.h) > shift_np(ind['slow_resistance'], 1), 1, \n",
    "                      np.where(pd.Series(md.l) < shift_np(ind['slow_support'], 1), -1, \n",
    "                      np.nan)))\n",
    "    trend.fillna(method = 'ffill', inplace = True)\n",
    "    return trend.values\n",
    "\n",
    "def build_strategy(contract, marketdata_collection, fast_interval, slow_interval):\n",
    "    strategy = Strategy([contract], marketdata_collection)\n",
    "    strategy.add_indicator('slow_resistance', lambda md : get_resistance(md.h, slow_interval))\n",
    "    strategy.add_indicator('slow_support', lambda md : get_support(md.l, slow_interval))\n",
    "    strategy.add_indicator('fast_resistance', lambda md : get_resistance(md.h, fast_interval))\n",
    "    strategy.add_indicator('fast_support', lambda md : get_support(md.l, fast_interval))\n",
    "    strategy.add_signal('trend', get_trend)\n",
    "    strategy.add_market_sim(market_simulator)\n",
    "    strategy.add_rule('trade_rule', trade_rule, 'trend', np.array([-1, 1]))\n",
    "    return strategy\n",
    "    \n",
    "def test_portfolio():\n",
    "    \n",
    "    np.random.seed(0)\n",
    "    dates = np.arange(datetime.datetime(2018, 1, 1, 9, 0, 0), datetime.datetime(2018, 3, 1, 16, 0, 0), datetime.timedelta(minutes = 5))\n",
    "    dates = np.array([dt for dt in dates.astype(object) if dt.hour >= 9 and dt.hour <= 16]).astype('M8[m]')\n",
    "    rets = np.random.normal(size = len(dates)) / 1000\n",
    "    c_0 = 100\n",
    "    c = np.round(c_0 * np.cumprod(1 + rets), 2)\n",
    "    l = np.round(c * (1. - np.abs(np.random.random(size = len(dates)) / 1000.)), 2)\n",
    "    h = np.round(c * (1. + np.abs(np.random.random(size = len(dates)) / 1000.)), 2)\n",
    "    o = np.round(l + (h - l) * np.random.random(size = len(dates)), 2)\n",
    "    v = np.round(np.random.normal(size = len(dates)) * 100)\n",
    "    vwap = 0.5 * (h + l)\n",
    "    \n",
    "    portfolio = Portfolio()\n",
    "    \n",
    "    slow_interval = 0\n",
    "    \n",
    "    for days in [0.5, 1, 2]:\n",
    "        # 1 day from 9 - 4 pm has 7 hours which translate to 7 x 12 = 84 5 minute periods\n",
    "        fast_interval = round(days * 84) \n",
    "        slow_interval = round(5 * fast_interval)\n",
    "        contract = Contract('IBM')\n",
    "        marketdata_collection = MarketDataCollection(['IBM'], [MarketData(dates, c, o, h, l, v, vwap)])\n",
    "        portfolio.add_strategy(f'strat_{days}', build_strategy(contract, marketdata_collection, fast_interval, slow_interval))\n",
    "    \n",
    "    # Start at max slow days so all strategies start at the same time\n",
    "    print('running')\n",
    "    portfolio.run(start_date = dates[slow_interval])\n",
    "    print('done')\n",
    "    \n",
    "    strat1 = portfolio.strategies['strat_0.5']\n",
    "    portfolio.plot();\n",
    "    strat1.plot();\n",
    "    \n",
    "if __name__ == \"__main__\":\n",
    "    test_portfolio()\n"
   ]
  }
 ],
 "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
