{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-11-06T00:59:22.546798Z",
     "start_time": "2018-11-06T00:59:20.579680Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import IPython.display as dsp\n",
    "import matplotlib.dates as mdates\n",
    "from pyqstrat.pq_utils import *\n",
    "from pyqstrat.plot import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-11-06T01:10:24.272962Z",
     "start_time": "2018-11-06T01:10:22.389014Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inferred Frequency: 5.0 minutes\n"
     ]
    },
    {
     "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>date</th>\n",
       "      <th>o</th>\n",
       "      <th>h</th>\n",
       "      <th>l</th>\n",
       "      <th>c</th>\n",
       "      <th>v</th>\n",
       "      <th>vwap</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>count</th>\n",
       "      <td>5748</td>\n",
       "      <td>5748</td>\n",
       "      <td>5748</td>\n",
       "      <td>5748</td>\n",
       "      <td>5748</td>\n",
       "      <td>5748</td>\n",
       "      <td>5748</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>num_missing</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>pct_missing</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000174</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>min</th>\n",
       "      <td>2018-01-01 09:00:00</td>\n",
       "      <td>88.17</td>\n",
       "      <td>88.21</td>\n",
       "      <td>88.09</td>\n",
       "      <td>88.14</td>\n",
       "      <td>0</td>\n",
       "      <td>88.16</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>max</th>\n",
       "      <td>2018-03-01 15:55:00</td>\n",
       "      <td>102.1</td>\n",
       "      <td>102.1</td>\n",
       "      <td>1000</td>\n",
       "      <td>102.1</td>\n",
       "      <td>3832</td>\n",
       "      <td>102.1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                            date     o     h     l        c     v  vwap\n",
       "count                       5748  5748  5748  5748     5748  5748  5748\n",
       "num_missing                    0     0     0     0        1     0     0\n",
       "pct_missing                    0     0     0     0 0.000174     0     0\n",
       "min          2018-01-01 09:00:00 88.17 88.21 88.09    88.14     0 88.16\n",
       "max          2018-03-01 15:55:00 102.1 102.1  1000    102.1  3832 102.1"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Errors:\n"
     ]
    },
    {
     "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>o</th>\n",
       "      <th>h</th>\n",
       "      <th>l</th>\n",
       "      <th>c</th>\n",
       "      <th>v</th>\n",
       "      <th>vwap</th>\n",
       "      <th>error</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>date</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>2018-01-01 16:05:00</th>\n",
       "      <td>99.95</td>\n",
       "      <td>99.98</td>\n",
       "      <td>1000</td>\n",
       "      <td>99.96</td>\n",
       "      <td>631</td>\n",
       "      <td>99.96</td>\n",
       "      <td>bad low</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                        o     h    l     c   v  vwap    error\n",
       "date                                                         \n",
       "2018-01-01 16:05:00 99.95 99.98 1000 99.96 631 99.96  bad low"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warnings:\n"
     ]
    },
    {
     "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>o</th>\n",
       "      <th>h</th>\n",
       "      <th>l</th>\n",
       "      <th>c</th>\n",
       "      <th>v</th>\n",
       "      <th>vwap</th>\n",
       "      <th>ret</th>\n",
       "      <th>warning</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>date</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>2018-01-01 16:00:00</th>\n",
       "      <td>99.79</td>\n",
       "      <td>99.84</td>\n",
       "      <td>99.75</td>\n",
       "      <td>99.77</td>\n",
       "      <td>863</td>\n",
       "      <td>99.8</td>\n",
       "      <td>nan</td>\n",
       "      <td>l ret &gt; 10 * std: 0.11963</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2018-01-01 16:05:00</th>\n",
       "      <td>99.95</td>\n",
       "      <td>99.98</td>\n",
       "      <td>1000</td>\n",
       "      <td>99.96</td>\n",
       "      <td>631</td>\n",
       "      <td>99.96</td>\n",
       "      <td>9.025</td>\n",
       "      <td>l ret &gt; 10 * std: 0.11963</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                        o     h     l     c   v  vwap   ret  \\\n",
       "date                                                          \n",
       "2018-01-01 16:00:00 99.79 99.84 99.75 99.77 863  99.8   nan   \n",
       "2018-01-01 16:05:00 99.95 99.98  1000 99.96 631 99.96 9.025   \n",
       "\n",
       "                                       warning  \n",
       "date                                            \n",
       "2018-01-01 16:00:00  l ret > 10 * std: 0.11963  \n",
       "2018-01-01 16:05:00  l ret > 10 * std: 0.11963  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time distribution:\n",
      "dates len: 5748\n",
      "o len: 5748\n",
      "h len: 5748\n",
      "l len: 5748\n",
      "c len: 5748\n",
      "v len: 5748\n",
      "vwap len: 5748\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32684 characters]\n",
      "text/plain": [
       "<Figure size 1440x576 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 39028 characters]\n",
      "text/plain": [
       "<Figure size 1080x576 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def sort_ohlcv_key(a):\n",
    "    l = ['date', 'o', 'h', 'l', 'c', 'v', 'vwap']\n",
    "    if a in l:\n",
    "        return l.index(a)\n",
    "    else:\n",
    "        return len(l)\n",
    "    \n",
    "def sort_ohlcv(columns):\n",
    "    columns = sorted(list(columns)) # Use stable sort to sort columns that we don't know about alphabetically\n",
    "    return sorted(columns, key = sort_ohlcv_key)\n",
    "\n",
    "class MarketDataCollection:\n",
    "    '''\n",
    "    Used to store a set of market data linking symbol -> MarketData \n",
    "    '''\n",
    "    def __init__(self, symbols = None, marketdata_list = None):\n",
    "        '''\n",
    "        Args:\n",
    "            symbols (list of str, optional): Symbols we want to store market data for.  Default None\n",
    "            marketdata_list (list of MarketData): Corresponding MarketData object.  Default None\n",
    "        '''\n",
    "        if symbols is not None or marketdata_list is not None:\n",
    "            if symbols is None or marketdata_list is None:\n",
    "                raise Exception('symbols and marketdata must either both be None or not None')\n",
    "            if len(symbols) != len(marketdata_list):\n",
    "                raise Exception('symbols and marketdata_list must contain the same number of elements')\n",
    "            self.marketdata = dict(zip(symbols, marketdata_list))\n",
    "        else:\n",
    "            self.marketdata = {}\n",
    "        \n",
    "    def add_marketdata(self, symbol, marketdata):\n",
    "        self.marketdata[symbol] = marketdata\n",
    "        \n",
    "    def add_dates(self, dates):\n",
    "        for md in self.marketdata.values(): md.add_dates(dates)\n",
    "            \n",
    "    def dates(self):\n",
    "        if len(self.marketdata) == 0:\n",
    "            return np.array([], dtype = np.datetime64)\n",
    "        else:\n",
    "            return list(self.marketdata.values())[0].dates\n",
    "        \n",
    "    def items(self):\n",
    "        return self.marketdata.items()\n",
    "    \n",
    "class MarketData:\n",
    "    '''Used to store OHLCV bars, and any additional time series data you want to use to simulate orders and executions.\n",
    "        You must at least supply dates and close prices.  All other fields are optional.\n",
    "    \n",
    "    Attributes:\n",
    "        dates: A numpy datetime array with the datetime for each bar.  Must be monotonically increasing.\n",
    "        c:     A numpy float array with close prices for the bar.\n",
    "        o:     A numpy float array with open prices . Default None\n",
    "        h:     A numpy float array with high prices. Default None\n",
    "        l:     A numpy float array with high prices. Default None\n",
    "        v:     A numpy integer array with volume for the bar. Default None\n",
    "        vwap:  A numpy float array with the volume weighted average price for the bar.  Default None\n",
    "        additional_arrays: A dictionary of name -> numpy array you want to add.  Default None\n",
    "        resample_funcs: A dictionary of functions for resampling each additional array.  Default None.\n",
    "        fill_funcs: A dictionary of functions for filling empty rows when we add dates.  Default None.\n",
    "    '''\n",
    "    def __init__(self, dates, c, o = None, h = None, l = None, v = None, vwap = None, additional_arrays = None, resample_funcs = None, fill_values = None):\n",
    "        '''Zeroes in o, h, l, c are set to nan'''\n",
    "        assert(len(dates) > 1)\n",
    "        assert(len(c) == len(dates))\n",
    "        assert(o is None or len(o) == len(dates))\n",
    "        assert(h is None or len(h) == len(dates))\n",
    "        assert(l is None or len(l) == len(dates))\n",
    "        assert(v is None or len(v) == len(dates))\n",
    "        assert(vwap is None or len(vwap) == len(dates))\n",
    "        \n",
    "        additional_arrays = {} if additional_arrays is None else additional_arrays\n",
    "        \n",
    "        for k, arr in additional_arrays.items():\n",
    "            assert(len(arr) == len(dates))\n",
    "            setattr(self, k, arr)\n",
    "            \n",
    "        self.additional_col_names = list(additional_arrays.keys())\n",
    "        self.resample_funcs = {} if resample_funcs is None else resample_funcs\n",
    "        self.fill_values = {} if fill_values is None else fill_values\n",
    "        \n",
    "        if not np.all(np.diff(dates).astype(np.float) > 0): # check for monotonically increasing dates\n",
    "            raise Exception('marketdata dates must be unique monotonically increasing')\n",
    "            \n",
    "        self.dates = dates\n",
    "        self.o = zero_to_nan(o)\n",
    "        self.h = zero_to_nan(h)\n",
    "        self.l = zero_to_nan(l)\n",
    "        self.c = zero_to_nan(c)\n",
    "        self.v = v\n",
    "        self.vwap = vwap\n",
    "        self._set_valid_rows()\n",
    "        \n",
    "    def add_dates(self, dates):\n",
    "        '''\n",
    "        Adds new dates to a market data object.  If fill_values was specified we use that to fill in values for any columns \n",
    "        for new dates that are not the same as the old dates.\n",
    "        \n",
    "        Args:\n",
    "            dates (np.array of np.datetime64): New dates to add.  Does not have to be sorted or unique\n",
    "        \n",
    "        >>> dates = np.array(['2018-01-05', '2018-01-09', '2018-01-10'], dtype = 'M8[ns]')\n",
    "        >>> c = np.array([8.1, 8.2, 8.3])\n",
    "        >>> o = np.array([9, 10, 11])\n",
    "        >>> additional_arrays = {'x' : np.array([5.1, 5.3, 5.5])}\n",
    "        >>> fill_values = {'x' : 0}\n",
    "        >>> md = MarketData(dates, c, o, additional_arrays = additional_arrays, fill_values = fill_values)\n",
    "        >>> new_dates = np.array(['2018-01-07', '2018-01-09'], dtype = 'M8[ns]')\n",
    "        >>> md.add_dates(new_dates)\n",
    "        >>> print(md.dates)\n",
    "        ['2018-01-05T00:00:00.000000000' '2018-01-07T00:00:00.000000000'\n",
    "         '2018-01-09T00:00:00.000000000' '2018-01-10T00:00:00.000000000']\n",
    "        >>> np.set_printoptions(formatter = {'float' : lambda x : f'{x:.4f}'})  # After numpy 1.13 positive floats don't have a leading space for sign\n",
    "        >>> print(md.o, md.c, md.x)\n",
    "        [9.0000 nan 10.0000 11.0000] [8.1000 nan 8.2000 8.3000] [5.1000 0.0000 5.3000 5.5000]\n",
    "        '''\n",
    "        if dates is None or len(dates) == 0: return\n",
    "        dates = np.unique(dates)\n",
    "        new_dates = np.setdiff1d(dates, self.dates, assume_unique = True)\n",
    "        all_dates = np.concatenate([self.dates, new_dates])\n",
    "        col_list = ['o', 'h', 'l', 'c', 'vwap'] + self.additional_col_names\n",
    "        sort_index = all_dates.argsort()\n",
    "        for col in col_list:\n",
    "            v = getattr(self, col)\n",
    "            if v is None: continue\n",
    "            dtype = getattr(self, col).dtype\n",
    "            fill_value = self.fill_values[col] if col in self.fill_values else get_empty_np_value(dtype)\n",
    "            v = np.concatenate([v, np.full(len(new_dates), fill_value, dtype = dtype)])\n",
    "            v = v[sort_index]\n",
    "            setattr(self, col, v)\n",
    "        self.dates = np.sort(all_dates)\n",
    "        self._set_valid_rows\n",
    "        \n",
    "    def _get_fill_value(self, col_name):\n",
    "        dtype = getattr(self, col_name).dtype\n",
    "        return get_empty_np_value(dtype)\n",
    "        \n",
    "    def _set_valid_rows(self):\n",
    "        col_list = [col for col in [self.o, self.h, self.l, self.c, self.vwap] if col is not None]\n",
    "        nans = np.any(np.isnan(col_list), axis = 0)\n",
    "        self.valid_rows = ~nans\n",
    "    \n",
    "    def valid_row(self, i):\n",
    "        '''Return True if the row with index i has no nans in it.'''\n",
    "        return self.valid_rows[i]\n",
    "    \n",
    "    def get_additional_arrays(self):\n",
    "        ret = {}\n",
    "        for key in self.additional_col_names:\n",
    "            ret[key] = getattr(self, key)\n",
    "        return ret\n",
    "    \n",
    "    def resample(self, sampling_frequency, inplace = False):\n",
    "        '''\n",
    "        Downsample the OHLCV data into a new bar frequency\n",
    "        \n",
    "        Args:\n",
    "            sampling_frequency: See sampling frequency in pandas\n",
    "            inplace: If set to False, don't modify this object, return a new object instead.\n",
    "        '''\n",
    "        if sampling_frequency is None:\n",
    "            if inplace: return None\n",
    "            return self\n",
    "        \n",
    "        df = self.df()\n",
    "        # Rename index from date to dates since our internal variable is called \"dates\" but the df() function returns a column \"date\"\n",
    "        df.index.name = 'dates'\n",
    "\n",
    "        df = resample_ohlc(df, sampling_frequency, self.resample_funcs)\n",
    "              \n",
    "        if inplace:\n",
    "            md = self\n",
    "        else:\n",
    "            # Create a dummy object, will replace everything (except additional arrays and resample funcs) later\n",
    "            md = MarketData(self.dates, self.c, self.o, self.h, self.l, self.v, self.vwap, self.get_additional_arrays(), self.resample_funcs)\n",
    "            \n",
    "        for col in df.columns: setattr(md, col, df[col].values)\n",
    "        md._set_valid_rows()\n",
    "        \n",
    "        return None if inplace else md\n",
    "    \n",
    "    def errors(self, display = True):\n",
    "        '''Returns a dataframe indicating any highs that are lower than opens, closes, lows or lows that are higher than other columns\n",
    "        Also includes any ohlcv values that are negative\n",
    "        '''\n",
    "        df = self.df()\n",
    "        errors_list = []\n",
    "        if 'h' in df.columns:\n",
    "            bad_highs = df[(df.h < df.c) | (df.h < df.o)]\n",
    "            if len(bad_highs):                 \n",
    "                bad_highs.insert(len(df.columns), 'error', 'bad high')\n",
    "                errors_list.append(bad_highs)\n",
    "        if 'l' in df.columns:\n",
    "            bad_lows = df[(df.l > df.c) | (df.l > df.o)]\n",
    "            if len(bad_lows): \n",
    "                bad_lows.insert(len(df.columns), 'error', 'bad low')\n",
    "                errors_list.append(bad_lows)\n",
    "\n",
    "        neg_values_mask = (df.c < 0)\n",
    "        for col in ['o', 'h', 'l', 'c', 'v', 'vwap']:\n",
    "            if col in df.columns:\n",
    "                neg_values_mask |= (df[col] < 0)\n",
    "        neg_values = df[neg_values_mask]\n",
    "        if len(neg_values): \n",
    "            neg_values.insert(len(df.columns), 'error', 'negative values')\n",
    "            errors_list.append(neg_values)\n",
    "            \n",
    "        if not len(errors_list): return None\n",
    "            \n",
    "        df = pd.concat(errors_list)\n",
    "        df = df[sort_ohlcv(df.columns)]\n",
    "        \n",
    "        if display: dsp.display(df)\n",
    "        return df\n",
    "    \n",
    "    def warnings(self, warn_std = 10, display = True):\n",
    "        '''Returns a dataframe indicating any values where the bar over bar change is more than warn_std standard deviations.\n",
    "        \n",
    "        Args:\n",
    "            warn_std: Number of standard deviations to use as a threshold (default 10)\n",
    "            display:  Whether to print out the warning dataframe as well as returning it\n",
    "        '''\n",
    "        df = self.df()\n",
    "        warnings_list = []\n",
    "\n",
    "        for col in ['o', 'h', 'l', 'c', 'vwap']:\n",
    "            if col in df.columns:\n",
    "                data = df[col]\n",
    "                ret = np.abs(df[col].pct_change())\n",
    "                std = ret.std()\n",
    "                mask = ret > warn_std * std\n",
    "                df_tmp = df[mask]\n",
    "                if len(df_tmp):\n",
    "                    double_mask = mask | mask.shift(-1) # Add the previous row so we know the two values computing a return\n",
    "                    df_tmp = df[double_mask]\n",
    "                    df_tmp.insert(len(df_tmp.columns), 'ret', ret[mask])\n",
    "                    df_tmp.insert(len(df_tmp.columns), 'warning', f'{col} ret > {warn_std} * std: {std:.5g}')\n",
    "                    warnings_list.append(df_tmp)\n",
    "\n",
    "        if not len(warnings_list): return None\n",
    "        df = pd.concat(warnings_list)\n",
    "        df = df[sort_ohlcv(df.columns)]\n",
    "        if display: dsp.display(df)\n",
    "        return df\n",
    "                              \n",
    "    def overview(self, display = True):\n",
    "        '''Returns a dataframe showing basic information about the data, including count, number and percent missing, min, max\n",
    "        \n",
    "        Args:\n",
    "            display:  Whether to print out the warning dataframe as well as returning it\n",
    "        '''\n",
    "        df = self.df().reset_index()\n",
    "        df_overview = pd.DataFrame({'count': len(df), 'num_missing' : df.isnull().sum(), 'pct_missing': df.isnull().sum() / len(df), 'min' : df.min(), 'max' : df.max()})\n",
    "        df_overview = df_overview.T\n",
    "        df_overview = df_overview[sort_ohlcv(df_overview.columns)]\n",
    "        if display: dsp.display(df_overview)\n",
    "        return df_overview\n",
    "       \n",
    "    def time_distribution(self, frequency = '15 minutes', display = True, plot = True, figsize = None):\n",
    "        '''\n",
    "        Return a dataframe with the time distribution of the bars\n",
    "        \n",
    "        Args:\n",
    "            frequency: The width of each bin (default \"15 minutes\").  You can use hours or days as well.\n",
    "            display:   Whether to display the data in addition to returning it.\n",
    "            plot:      Whether to plot the data in addition to returning it.\n",
    "            figsize:   If plot is set, optional figure size for the plot (default (20,8))\n",
    "        '''\n",
    "        group_col = None\n",
    "        \n",
    "        n = int(frequency.split(' ')[0])\n",
    "        freq = frequency.split(' ')[1]\n",
    "        \n",
    "        df = self.df().reset_index()\n",
    "        \n",
    "        if freq == 'minutes' or freq == 'mins' or freq == 'min':\n",
    "            group_col = [df.date.dt.hour, df.date.dt.minute // n * n]\n",
    "            names = ['hour', 'minute']\n",
    "        elif freq == 'hours' or freq == 'hrs' or freq == 'hr':\n",
    "            group_col = [df.date.dt.weekday_name, df.date.dt.hour // n * n]\n",
    "            names = ['weekday', 'hour']\n",
    "        elif freq == 'weekdays' or freq == 'days' or freq == 'day':\n",
    "            group_col = df.date.dt.weekday_name // n * n\n",
    "            names = ['weekday']\n",
    "        else:\n",
    "            raise Exception(f'unknown time freq: {freq}')\n",
    "            \n",
    "        count = df.groupby(group_col)['c'].count()\n",
    "        tdf = pd.DataFrame({'close_count': count, 'count_pct' : count / df.c.count()})[['close_count', 'count_pct']]\n",
    "            \n",
    "        if 'v' in df.columns:\n",
    "            vsum = df.groupby(group_col)['v'].sum()\n",
    "            vdf = pd.DataFrame({'volume' : vsum, 'volume_pct' : vsum / df.v.sum()})[['volume', 'volume_pct']]\n",
    "            tdf = pd.concat([vdf, tdf], axis = 1)\n",
    "            \n",
    "        tdf.index.names = names\n",
    "            \n",
    "        if display:\n",
    "            dsp.display(tdf)\n",
    "    \n",
    "        if plot:\n",
    "            if not figsize: figsize = (20, 8)\n",
    "            cols = ['close_count', 'volume'] if 'v' in df.columns else ['close_count']\n",
    "            if not has_display():\n",
    "                print('no display found, cannot plot time distribution')\n",
    "                return tdf\n",
    "            tdf[cols].plot(figsize = figsize, kind = 'bar', subplots = True, title = 'Time Distribution')\n",
    "            \n",
    "        return tdf\n",
    "    \n",
    "    def freq_str(self):\n",
    "        \n",
    "        freq = infer_frequency(self.dates)\n",
    "        if freq < 1:\n",
    "            freq_str = f'{round(freq * 24. * 60, 2)} minutes'\n",
    "        else:\n",
    "            freq_str = f'{freq} days'\n",
    "        return freq_str\n",
    "            \n",
    "    def describe(self, warn_std = 10, time_distribution_frequency = '15 min', print_time_distribution = False):\n",
    "        '''\n",
    "        Describe the bars.  Shows an overview, errors and warnings for the bar data.  This is a good function to use \n",
    "        before running any backtests on a set of bar data.\n",
    "        \n",
    "        Args:\n",
    "            warn_std: See warning function\n",
    "            time_distribution_frequency: See time_distribution function\n",
    "            print_time_distribution: Whether to print the time distribution in addition to plotting it.\n",
    "        '''\n",
    "        print(f'Inferred Frequency: {self.freq_str()}')\n",
    "        self.overview()\n",
    "        print('Errors:')\n",
    "        self.errors()\n",
    "        print('Warnings:')\n",
    "        self.warnings(warn_std = warn_std)\n",
    "        print('Time distribution:')\n",
    "        self.time_distribution(display = print_time_distribution, frequency = time_distribution_frequency)\n",
    "        \n",
    "    def is_ohlc(self):\n",
    "        '''\n",
    "        Returns True if we have all ohlc columns and none are empty\n",
    "        '''\n",
    "        return not (self.o is None or self.h is None or self.l is None or self.c is None)\n",
    "\n",
    "    def plot(self, figsize = (15,8), date_range = None, sampling_frequency = None, title = 'Price / Volume'):\n",
    "        '''\n",
    "        Plot a candlestick or line plot depending on whether we have ohlc data or just close prices\n",
    "        \n",
    "        Args:\n",
    "            figsize: Size of the figure (default (15,8))\n",
    "            date_range: A tuple of strings or numpy datetimes for plotting a smaller sample of the data, e.g. (\"2018-01-01\", \"2018-01-06\")\n",
    "            sampling_frequency: Downsample before plotting.  See pandas frequency strings for possible values.\n",
    "            title: Title of the graph, default \"Price / Volume\"\n",
    "        '''\n",
    "        date_range = strtup2date(date_range)\n",
    "        if self.is_ohlc():\n",
    "            data = OHLC('price', self.dates, self.o, self.h, self.l, self.c, self.v, self.vwap)\n",
    "        else:\n",
    "            data = TimeSeries('price', self.dates, self.c)\n",
    "        subplot = Subplot(data)\n",
    "        plot = Plot([subplot], figsize = figsize, date_range = date_range, sampling_frequency = sampling_frequency, title = title)\n",
    "        plot.draw()\n",
    "                              \n",
    "    def df(self, start_date = None, end_date = None):\n",
    "        df = pd.DataFrame({'date' : self.dates, 'c' : self.c}).set_index('date')\n",
    "        for tup in [('o', self.o), ('h', self.h), ('l', self.l), ('v', self.v), ('vwap', self.vwap)]:\n",
    "            if tup[1] is not None: df.insert(0, tup[0], tup[1])\n",
    "        if start_date: df = df[df.index.values >= start_date]\n",
    "        if end_date: df = df[df.index.values <= end_date]\n",
    "        return df\n",
    "    \n",
    "    \n",
    "def roll_futures(md, date_func, condition_func, expiries = None, return_full_df = False):\n",
    "    '''Construct a continuous futures dataframe with one row per datetime given rolling logic\n",
    "    \n",
    "    Args:\n",
    "        md: A dataframe containing the columns 'date', 'series', and any other market data, for example, ohlcv data. Date can contain time for sub-daily bars. \n",
    "          The series column must contain a different string name for each futures series, e.g. SEP2018, DEC2018, etc.\n",
    "        date_func: A function that takes the market data object as an input and returns a numpy array of booleans\n",
    "          True indicates that the future should be rolled on this date if the condition specified in condition_func is met.\n",
    "          This function can assume that we have all the columns in the original market data object plus the same columns suffixed with _next for the potential series\n",
    "          to roll over to.\n",
    "        condition_func: A function that takes the market data object as input and returns a numpy array of booleans.\n",
    "          True indicates that we should try to roll the future at that row.\n",
    "        expiries: An optional dataframe with 2 columns, 'series' and 'expiry'.  This should have one row per future series indicating that future's expiry date.\n",
    "          If you don't pass this in, the function will assume that the expiry column is present in the original dataframe.\n",
    "        return_full_df: If set, will return the datframe without removing extra dates so you can use your own logic for rolling, including the _next columns and \n",
    "          the roll flag\n",
    "          \n",
    "    Returns:\n",
    "        A pandas DataFrame with one row per date, which contains the columns in the original md DataFrame and the same columns suffixed with _next \n",
    "          representing the series we want to roll to.  There is also a column called roll_flag which is set to True whenever \n",
    "          the date and roll condition functions are met.\n",
    "          \n",
    "    >>> md = pd.DataFrame({'date' : np.concatenate((np.arange(np.datetime64('2018-03-11'), np.datetime64('2018-03-16')),\n",
    "    ...                                            np.arange(np.datetime64('2018-03-11'), np.datetime64('2018-03-16')))),\n",
    "    ...                    'c' : [10, 10.1, 10.2, 10.3, 10.4] + [10.35, 10.45, 10.55, 10.65, 10.75],\n",
    "    ...                    'v' : [200, 200, 150, 100, 100] + [100, 50, 200, 250, 300],\n",
    "    ...                    'series' : ['MAR2018'] * 5 + ['JUN2018'] * 5})[['date','series', 'c', 'v']]\n",
    "    >>> expiries = pd.Series(np.array(['2018-03-15', '2018-06-15'], dtype = 'M8[D]'), index = ['MAR2018', 'JUN2018'], name = \"expiry\")\n",
    "    >>> date_func = lambda md : md.expiry - md.date <= np.timedelta64(3, 'D')\n",
    "    >>> condition_func = lambda md : md.v_next > md.v\n",
    "\n",
    "    >>> df = roll_futures(md, date_func, condition_func, expiries)\n",
    "    >>> df[df.series == 'MAR2018'].date.max() == np.datetime64('2018-03-14')\n",
    "    True\n",
    "    >>> df[df.series == 'JUN2018'].date.max() == np.datetime64('2018-03-15')\n",
    "    True\n",
    "    '''\n",
    "    if 'date' not in md.columns or 'series' not in md.columns:\n",
    "        raise Exception('date or series not found in columns: {md.columns}')\n",
    "        \n",
    "    if expiries is not None:\n",
    "        expiries = expiries.to_frame(name = 'expiry')\n",
    "        md = pd.merge(md, expiries, left_on = ['series'], right_index = True, how = 'left')\n",
    "    else:\n",
    "        if 'expiry' not in md.columns: raise Exception('expiry column must be present in market data if expiries argument is not specified')\n",
    "        expiries = md[['series', 'expiry']].drop_duplicates().sort_values(by = 'expiry').set_index('s')\n",
    "\n",
    "    expiries = pd.merge(expiries, expiries.shift(-1), left_index = True, right_index = True, how = 'left', suffixes = ['', '_next'])\n",
    "\n",
    "    orig_cols = [col for col in md.columns if col not in ['date']]\n",
    "    md1 = pd.merge(md, expiries[['expiry', 'expiry_next']], on = ['expiry'], how = 'left')\n",
    "    md = pd.merge(md1, md, left_on = ['date', 'expiry_next'], right_on = ['date', 'expiry'], how = 'left', suffixes = ['', '_next'])\n",
    "\n",
    "    md.sort_values(by = ['expiry', 'date'], inplace = True)\n",
    "\n",
    "    roll_flag = date_func(md) & condition_func(md) \n",
    "\n",
    "    df_roll = pd.DataFrame({'series' : md.series, 'date' : md.date, 'roll_flag' : roll_flag})\n",
    "    df_roll = df_roll[df_roll.roll_flag].groupby('series', as_index = False).first()\n",
    "    md = pd.merge(md, df_roll, on = ['series', 'date'], how = 'left')\n",
    "    md.roll_flag = md.roll_flag.fillna(False)\n",
    "    \n",
    "    cols = ['date'] + orig_cols + [col + '_next' for col in orig_cols] + ['roll_flag']\n",
    "    md = md[cols]\n",
    "    \n",
    "    if return_full_df: return md\n",
    "    \n",
    "    df_list = []\n",
    "    for series, g in md.groupby('expiry'):\n",
    "        roll_flag = g.roll_flag\n",
    "        true_values = roll_flag[roll_flag]\n",
    "        if len(true_values):\n",
    "            first_true_index = true_values.index[0]\n",
    "            roll_flag = roll_flag[first_true_index:]\n",
    "            false_after_true_values = roll_flag[~roll_flag]\n",
    "            if len(false_after_true_values):\n",
    "                first_false_after_true_idx = false_after_true_values.index[0]\n",
    "                g = g.loc[:first_false_after_true_idx]\n",
    "        df_list.append(g)\n",
    "\n",
    "    full_df = pd.concat(df_list)\n",
    "    full_df = full_df.sort_values(by = ['expiry', 'date']).drop_duplicates(subset=['date'])\n",
    "\n",
    "    return full_df\n",
    "\n",
    "\n",
    "def test_marketdata():\n",
    "    from datetime import datetime, timedelta\n",
    "    np.random.seed(0)\n",
    "    dates = np.arange(datetime(2018, 1, 1, 9, 0, 0), datetime(2018, 3, 1, 16, 0, 0), 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.abs(np.round(np.random.normal(size = len(dates)) * 1000))\n",
    "    vwap = 0.5 * (l + h)\n",
    "    c[18] = np.nan\n",
    "    l[85] = 1000\n",
    "    md = MarketData(dates, c, o, h, l, v, vwap)\n",
    "    md.describe()\n",
    "    md.plot(date_range = ('2018-01-02', '2018-01-02 12:00'))\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_marketdata()"
   ]
  }
 ],
 "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
}
