{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-16T19:47:25.632709Z",
     "start_time": "2019-10-16T19:47:25.430165Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running typecheck\n",
      "Success: no issues found in 1 source file\n",
      "\n",
      "running flake8\n",
      "flake8 success\n",
      "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>o</th>\n",
       "      <th>h</th>\n",
       "      <th>l</th>\n",
       "      <th>c</th>\n",
       "      <th>v</th>\n",
       "      <th>vwap</th>\n",
       "      <th>date</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>count</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",
       "      <td>5748</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>num_missing</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",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>pct_missing</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000173973556</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>min</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",
       "      <td>2018-01-01 09:00:00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>max</td>\n",
       "      <td>102.09</td>\n",
       "      <td>102.11</td>\n",
       "      <td>1000</td>\n",
       "      <td>102.07</td>\n",
       "      <td>3832</td>\n",
       "      <td>102.07</td>\n",
       "      <td>2018-03-01 15:55:00</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                 o      h     l              c     v   vwap  \\\n",
       "count         5748   5748  5748           5748  5748   5748   \n",
       "num_missing      0      0     0              1     0      0   \n",
       "pct_missing      0      0     0 0.000173973556     0      0   \n",
       "min          88.17  88.21 88.09          88.14     0  88.16   \n",
       "max         102.09 102.11  1000         102.07  3832 102.07   \n",
       "\n",
       "                            date  \n",
       "count                       5748  \n",
       "num_missing                    0  \n",
       "pct_missing                    0  \n",
       "min          2018-01-01 09:00:00  \n",
       "max          2018-03-01 15:55:00  "
      ]
     },
     "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",
       "      <td>2018-01-01 16:05:00</td>\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",
       "      <td>2018-01-01 16:00:00</td>\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.795</td>\n",
       "      <td>nan</td>\n",
       "      <td>l ret &gt; 10 * std: 0.11963</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2018-01-01 16:05:00</td>\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.02506266</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.795        nan   \n",
       "2018-01-01 16:05:00 99.95 99.98  1000 99.96 631  99.96 9.02506266   \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"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/sal/anaconda3/envs/py37/lib/python3.7/site-packages/pandas/plotting/_matplotlib/tools.py:307: MatplotlibDeprecationWarning: \n",
      "The rowNum attribute was deprecated in Matplotlib 3.2 and will be removed two minor releases later. Use ax.get_subplotspec().rowspan.start instead.\n",
      "  layout[ax.rowNum, ax.colNum] = ax.get_visible()\n",
      "/Users/sal/anaconda3/envs/py37/lib/python3.7/site-packages/pandas/plotting/_matplotlib/tools.py:307: MatplotlibDeprecationWarning: \n",
      "The colNum attribute was deprecated in Matplotlib 3.2 and will be removed two minor releases later. Use ax.get_subplotspec().colspan.start instead.\n",
      "  layout[ax.rowNum, ax.colNum] = ax.get_visible()\n",
      "/Users/sal/anaconda3/envs/py37/lib/python3.7/site-packages/pandas/plotting/_matplotlib/tools.py:313: MatplotlibDeprecationWarning: \n",
      "The rowNum attribute was deprecated in Matplotlib 3.2 and will be removed two minor releases later. Use ax.get_subplotspec().rowspan.start instead.\n",
      "  if not layout[ax.rowNum + 1, ax.colNum]:\n",
      "/Users/sal/anaconda3/envs/py37/lib/python3.7/site-packages/pandas/plotting/_matplotlib/tools.py:313: MatplotlibDeprecationWarning: \n",
      "The colNum attribute was deprecated in Matplotlib 3.2 and will be removed two minor releases later. Use ax.get_subplotspec().colspan.start instead.\n",
      "  if not layout[ax.rowNum + 1, ax.colNum]:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TestResults(failed=0, attempted=16)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 39420 characters]\n",
      "text/plain": [
       "<Figure size 1440x576 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 41496 characters]\n",
      "text/plain": [
       "<Figure size 1080x576 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%checkall\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import IPython.display as dsp\n",
    "from pyqstrat.pq_utils import zero_to_nan, get_empty_np_value, infer_frequency, resample_trade_bars, has_display, strtup2date\n",
    "from pyqstrat.plot import TradeBarSeries, TimeSeries, Subplot, Plot\n",
    "\n",
    "from typing import Optional, Sequence, Tuple, Union, Callable\n",
    "\n",
    "\n",
    "def _sort_trade_bars_key(a: str) -> int:\n",
    "    sorted_cols = ['timestamp', 'o', 'h', 'l', 'c', 'v', 'vwap']\n",
    "    if a in sorted_cols:\n",
    "        return sorted_cols.index(a)\n",
    "    else:\n",
    "        return len(sorted_cols)\n",
    "    \n",
    "\n",
    "def sort_trade_bars(columns: Sequence[str]) -> Sequence[str]:\n",
    "    '''Given a list of column names, sort them in olhcv order'''\n",
    "    columns = sorted(list(columns))  # Use stable sort to sort columns that we don't know about alphabetically\n",
    "    return sorted(columns, key=_sort_trade_bars_key)\n",
    "    \n",
    "\n",
    "class TradeBars:\n",
    "    '''Used to store OHLCV bars.  You must at least supply timestamps and close prices.  All other fields are optional.\n",
    "    \n",
    "    Attributes:\n",
    "        timestamp: 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",
    "    '''\n",
    "    def __init__(self, \n",
    "                 timestamps: np.ndarray, \n",
    "                 c: np.ndarray, \n",
    "                 o: Optional[np.ndarray] = None, \n",
    "                 h: Optional[np.ndarray] = None, \n",
    "                 l: Optional[np.ndarray] = None,  # noqa: E741: ignore # l ambiguous\n",
    "                 v: Optional[np.ndarray] = None, \n",
    "                 vwap: Optional[np.ndarray] = None) -> None:\n",
    "        '''Zeroes in o, h, l, c are set to nan'''\n",
    "        assert(len(timestamps) > 1)\n",
    "        assert(len(c) == len(timestamps))\n",
    "        assert(o is None or len(o) == len(timestamps))\n",
    "        assert(h is None or len(h) == len(timestamps))\n",
    "        assert(l is None or len(l) == len(timestamps))\n",
    "        assert(v is None or len(v) == len(timestamps))\n",
    "        assert(vwap is None or len(vwap) == len(timestamps))\n",
    "        \n",
    "        if not np.all(np.diff(timestamps).astype(np.float) > 0):  # check for monotonically increasing timestamps\n",
    "            raise Exception('timestamps must be unique monotonically increasing')\n",
    "        self.timestamps, self.o, self.h, self.l, self.c, self.v, self.vwap = timestamps, o, h, l, c, v, vwap\n",
    "            \n",
    "        for field in ['timestamps', 'h', 'l', 'c', 'v', 'vwap']:\n",
    "            v = getattr(self, field)\n",
    "            if isinstance(v, pd.Series):\n",
    "                setattr(self, field, v.values)\n",
    "                \n",
    "        for field in ['o', 'h', 'l', 'c']:\n",
    "            setattr(self, field, zero_to_nan(getattr(self, field)))\n",
    "        \n",
    "        self._set_valid_rows()\n",
    "        \n",
    "    def add_timestamps(self, timestamps: np.ndarray) -> None:\n",
    "        '''\n",
    "        Adds new timestamps to a market data object.\n",
    "        \n",
    "        Args:\n",
    "            timestamps (np.array of np.datetime64): New timestamps to add.  Does not have to be sorted or unique\n",
    "        \n",
    "        >>> timestamps = 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",
    "        >>> trade_bar = TradeBars(timestamps, c, o)\n",
    "        >>> new_timestamps = np.array(['2018-01-07', '2018-01-09'], dtype = 'M8[ns]')\n",
    "        >>> trade_bar.add_timestamps(new_timestamps)\n",
    "        >>> print(trade_bar.timestamps)\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(trade_bar.o, trade_bar.c)\n",
    "        [9.0000 nan 10.0000 11.0000] [8.1000 nan 8.2000 8.3000]\n",
    "        '''\n",
    "        if timestamps is None or len(timestamps) == 0: return\n",
    "        timestamps = np.unique(timestamps)\n",
    "        new_timestamps = np.setdiff1d(timestamps, self.timestamps, assume_unique=True)\n",
    "        all_timestamps = np.concatenate([self.timestamps, new_timestamps])\n",
    "        col_list = ['o', 'h', 'l', 'c', 'vwap']\n",
    "        sort_index = all_timestamps.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 = get_empty_np_value(dtype)\n",
    "            v = np.concatenate([v, np.full(len(new_timestamps), fill_value, dtype=dtype)])\n",
    "            v = v[sort_index]\n",
    "            setattr(self, col, v)\n",
    "        self.timestamps = np.sort(all_timestamps)\n",
    "        self._set_valid_rows\n",
    "        \n",
    "    def _get_fill_value(self, col_name: str) -> np.generic:\n",
    "        dtype = getattr(self, col_name).dtype\n",
    "        return get_empty_np_value(dtype)\n",
    "        \n",
    "    def _set_valid_rows(self) -> None:\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: int) -> bool:\n",
    "        '''Return True if the row with index i has no nans in it.'''\n",
    "        return self.valid_rows[i]\n",
    "    \n",
    "    def resample(self, sampling_frequency: str) -> Optional['TradeBars']:\n",
    "        '''\n",
    "        Downsample the trade bars data into a new bar frequency\n",
    "        \n",
    "        Args:\n",
    "            sampling_frequency: See sampling frequency in pandas\n",
    "        '''\n",
    "        if sampling_frequency is None:\n",
    "            return self\n",
    "        \n",
    "        df = self.df()\n",
    "        # Rename timestamps to timestamp\n",
    "        df.index.name = 'timestamp'\n",
    "\n",
    "        df = resample_trade_bars(df, sampling_frequency)\n",
    "        o = df.o if 'o' in df.columns else None\n",
    "        h = df.h if 'h' in df.columns else None\n",
    "        _l = df.l if 'l' in df.columns else None\n",
    "        v = df.v if 'v' in df.columns else None\n",
    "        vwap = df.vwap if 'vwap' in df.columns else None\n",
    "              \n",
    "        trade_bar = TradeBars(df.timestamp, df.c, o, h, _l, v, vwap)\n",
    "            \n",
    "        trade_bar._set_valid_rows()\n",
    "        \n",
    "        return trade_bar\n",
    "    \n",
    "    def errors(self, display: bool = True) -> Optional[pd.DataFrame]:\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_trade_bars(df.columns)]\n",
    "        \n",
    "        if display: dsp.display(df)\n",
    "        return df\n",
    "    \n",
    "    def warnings(self, warn_std: int = 10, display: bool = True) -> pd.DataFrame:\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",
    "                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_trade_bars(df.columns)]\n",
    "        if display: dsp.display(df)\n",
    "        return df\n",
    "                              \n",
    "    def overview(self, display: bool = True) -> pd.DataFrame:\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), \n",
    "                                    'num_missing': df.isnull().sum(), \n",
    "                                    'pct_missing': df.isnull().sum() / len(df), \n",
    "                                    'min': df.min(), \n",
    "                                    'max': df.max()})\n",
    "        df_overview = df_overview.T\n",
    "        df_overview = df_overview[sort_trade_bars(df_overview.columns)]\n",
    "        if display: dsp.display(df_overview)\n",
    "        return df_overview\n",
    "       \n",
    "    def time_distribution(self, \n",
    "                          frequency: str = '15 minutes', \n",
    "                          display: bool = True, \n",
    "                          plot: bool = True, \n",
    "                          figsize: Optional[Tuple[int, int]] = None) -> pd.DataFrame:\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) -> str:\n",
    "        \n",
    "        freq = infer_frequency(self.timestamps)\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, \n",
    "                 warn_std: int = 10, \n",
    "                 time_distribution_frequency: str = '15 min', \n",
    "                 print_time_distribution: bool = False) -> None:\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 has_ohlc(self) -> bool:\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,\n",
    "             figsize: Tuple[int, int] = (15, 8),\n",
    "             date_range: Optional[Union[Tuple[str, str], Tuple[np.datetime64, np.datetime64]]] = None,\n",
    "             sampling_frequency: str = None,\n",
    "             title: str = 'Price / Volume') -> None:\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",
    "        if date_range and isinstance(date_range[0], str):\n",
    "            date_range = strtup2date(date_range)\n",
    "        data: Union[TradeBarSeries, TimeSeries]\n",
    "        if self.has_ohlc():\n",
    "            data = TradeBarSeries('price', self.timestamps, self.o, self.h, self.l, self.c, self.v, self.vwap)\n",
    "        else:\n",
    "            data = TimeSeries('price', self.timestamps, 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, \n",
    "           start_date: Optional[np.datetime64] = None, \n",
    "           end_date: Optional[np.datetime64] = None) -> pd.DataFrame:\n",
    "        df = pd.DataFrame({'date': self.timestamps, '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(fut_prices: pd.DataFrame, \n",
    "                 date_func: Callable[[pd.DataFrame], np.ndarray], \n",
    "                 condition_func: Callable[[pd.DataFrame], np.ndarray], \n",
    "                 expiries: pd.DataFrame = None,\n",
    "                 return_full_df: bool = False) -> pd.DataFrame:\n",
    "    '''Construct a continuous futures dataframe with one row per datetime given rolling logic\n",
    "    \n",
    "    Args:\n",
    "        fut_prices: A dataframe containing the columns 'date', 'series', and any other market data, \n",
    "            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 future prices 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 \n",
    "          columns suffixed with _next for the potential series to roll over to.\n",
    "        condition_func: A function that takes the future prices 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 \n",
    "            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 timestamps so you can use your own logic for rolling, \n",
    "            including the _next columns and 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",
    "          \n",
    "    >>> fut_prices = pd.DataFrame({'timestamp': 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})[['timestamp','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 fut_prices: fut_prices.expiry - fut_prices.timestamp <= np.timedelta64(3, 'D')\n",
    "    >>> condition_func = lambda fut_prices: fut_prices.v_next > fut_prices.v\n",
    "    >>> df = roll_futures(fut_prices, date_func, condition_func, expiries)\n",
    "    >>> print(df[df.series == 'MAR2018'].timestamp.max() == np.datetime64('2018-03-14'))\n",
    "    True\n",
    "    >>> print(df[df.series == 'JUN2018'].timestamp.max() == np.datetime64('2018-03-15'))\n",
    "    True\n",
    "    '''\n",
    "    if 'timestamp' not in fut_prices.columns or 'series' not in fut_prices.columns:\n",
    "        raise Exception('timestamp or series not found in columns: {fut_prices.columns}')\n",
    "        \n",
    "    if expiries is not None:\n",
    "        expiries = expiries.to_frame(name='expiry')\n",
    "        fut_prices = pd.merge(fut_prices, expiries, left_on=['series'], right_index=True, how='left')\n",
    "    else:\n",
    "        if 'expiry' not in fut_prices.columns: raise Exception('expiry column must be present in market data if expiries argument is not specified')\n",
    "        expiries = fut_prices[['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 fut_prices.columns if col not in ['timestamp']]\n",
    "    fut_prices1 = pd.merge(fut_prices, expiries[['expiry', 'expiry_next']], on=['expiry'], how='left')\n",
    "    fut_prices = pd.merge(fut_prices1, fut_prices, left_on=['timestamp', 'expiry_next'],\n",
    "                          right_on=['timestamp', 'expiry'], how='left', suffixes=['', '_next'])\n",
    "\n",
    "    fut_prices = fut_prices.sort_values(by=['expiry', 'timestamp'])\n",
    "\n",
    "    roll_flag = date_func(fut_prices) & condition_func(fut_prices) \n",
    "\n",
    "    df_roll = pd.DataFrame({'series': fut_prices.series, 'timestamp': fut_prices.timestamp, 'roll_flag': roll_flag})\n",
    "    df_roll = df_roll[df_roll.roll_flag].groupby('series', as_index=False).first()\n",
    "    fut_prices = pd.merge(fut_prices, df_roll, on=['series', 'timestamp'], how='left')\n",
    "    fut_prices.roll_flag = fut_prices.roll_flag.fillna(False)\n",
    "    \n",
    "    cols = ['timestamp'] + orig_cols + [col + '_next' for col in orig_cols] + ['roll_flag']\n",
    "    fut_prices = fut_prices[cols]\n",
    "    \n",
    "    if return_full_df: return fut_prices\n",
    "    \n",
    "    df_list = []\n",
    "    for series, g in fut_prices.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', 'timestamp']).drop_duplicates(subset=['timestamp'])\n",
    "\n",
    "    return full_df\n",
    "\n",
    "\n",
    "def test_trade_bars() -> None:\n",
    "    from datetime import datetime, timedelta\n",
    "    np.random.seed(0)\n",
    "    timestamps = np.arange(datetime(2018, 1, 1, 9, 0, 0), datetime(2018, 3, 1, 16, 0, 0), timedelta(minutes=5))\n",
    "    timestamps = np.array([dt for dt in timestamps.astype(object) if dt.hour >= 9 and dt.hour <= 16]).astype('M8[m]')\n",
    "    rets = np.random.normal(size=len(timestamps)) / 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(timestamps)) / 1000.)), 2)  # PEP8 thinks l is hard to distinguish\n",
    "    h = np.round(c * (1. + np.abs(np.random.random(size=len(timestamps)) / 1000.)), 2)\n",
    "    o = np.round(_l + (h - _l) * np.random.random(size=len(timestamps)), 2)\n",
    "    v = np.abs(np.round(np.random.normal(size=len(timestamps)) * 1000))\n",
    "    vwap = 0.5 * (_l + h)\n",
    "    c[18] = np.nan\n",
    "    _l[85] = 1000\n",
    "    trade_bar = TradeBars(timestamps, c, o, h, _l, v, vwap)\n",
    "    trade_bar.describe()\n",
    "    trade_bar.plot(date_range=('2018-01-02', '2018-01-02 12:00'))\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_trade_bars()\n",
    "    import doctest\n",
    "    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
