{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-14T17:39:07.214410Z",
     "start_time": "2022-07-14T17:39:04.742538Z"
    }
   },
   "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"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 198568 characters]\n",
      "text/plain": [
       "<Figure size 1440x1440 with 12 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 16700 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%checkall\n",
    "from dataclasses import dataclass\n",
    "import collections\n",
    "import collections.abc\n",
    "import math\n",
    "from abc import abstractmethod\n",
    "from functools import reduce\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.dates as mdates\n",
    "import matplotlib.ticker as mtick\n",
    "import matplotlib.patches as mptch\n",
    "import matplotlib.gridspec as gridspec\n",
    "import matplotlib.path as path\n",
    "import matplotlib.cm as cm\n",
    "from matplotlib.colors import BoundaryNorm\n",
    "\n",
    "from mpl_toolkits.mplot3d import Axes3D  # noqa: F401  # not directly used but need to import to plot 3d \n",
    "from scipy.interpolate import griddata\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "from pyqstrat.pq_utils import series_to_array, strtup2date, has_display, resample_ts, resample_trade_bars\n",
    "from pyqstrat.pq_types import ReasonCode, Trade\n",
    "from typing import Sequence, Tuple, Mapping, Union, MutableMapping, List, Optional\n",
    "\n",
    "# set_defaults()\n",
    "    \n",
    "\n",
    "class DateFormatter(mtick.Formatter):\n",
    "    '''\n",
    "    Formats timestamps on plot axes.  See matplotlib Formatter\n",
    "    '''\n",
    "    def __init__(self, timestamps: np.ndarray, fmt: str) -> None:\n",
    "        self.timestamps = timestamps\n",
    "        self.fmt = fmt\n",
    "\n",
    "    def __call__(self, x, pos: int = 0):\n",
    "        'Return the label for time x at position pos'\n",
    "        ind = int(np.round(x))\n",
    "        if ind >= len(self.timestamps) or ind < 0: return ''\n",
    "        return mdates.num2date(self.timestamps[ind]).strftime(self.fmt)\n",
    "    \n",
    "\n",
    "class HorizontalLine:\n",
    "    '''Draws a horizontal line on a subplot'''\n",
    "    def __init__(self, y: float, name: str = None, line_type: str = 'dashed', color: str = None) -> None:\n",
    "        self.y = y\n",
    "        self.name = name\n",
    "        self.line_type = line_type\n",
    "        self.color = color\n",
    "    \n",
    "\n",
    "class VerticalLine:\n",
    "    '''Draws a vertical line on a subplot where x axis is not a date-time axis'''\n",
    "    def __init__(self, x: float, name: str = None, line_type: str = 'dashed', color: str = None) -> None:\n",
    "        self.x = x\n",
    "        self.name = name\n",
    "        self.line_type = line_type\n",
    "        self.color = color\n",
    "    \n",
    "\n",
    "class DateLine:\n",
    "    '''Draw a vertical line on a plot with a datetime x-axis'''\n",
    "    def __init__(self, date: np.datetime64, name: str = None, line_type: str = 'dashed', color: str = None) -> None:\n",
    "        self.date = date\n",
    "        self.name = name\n",
    "        self.line_type = line_type\n",
    "        self.color = color\n",
    "        \n",
    "        \n",
    "class DisplayAttributes:\n",
    "    pass\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class BoxPlotAttributes(DisplayAttributes):\n",
    "    '''\n",
    "    Attributes:\n",
    "        proportional_widths: if set to True, the width each box in the boxplot will be proportional \n",
    "            to the number of items in its corresponding array\n",
    "        show_means: Whether to display a marker where the mean is for each array\n",
    "        show_outliers: Whether to show markers for outliers that are outside the whiskers.  \n",
    "            Box is at Q1 = 25%, Q3 = 75% quantiles, whiskers are at Q1 - 1.5 * (Q3 - Q1), Q3 + 1.5 * (Q3 - Q1)\n",
    "        notched: Whether to show notches indicating the confidence interval around the median\n",
    "    '''\n",
    "    proportional_widths: bool = True\n",
    "    show_means: bool = True\n",
    "    show_all: bool = True\n",
    "    show_outliers: bool = False\n",
    "    notched: bool = False\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class LinePlotAttributes(DisplayAttributes):\n",
    "    line_type: Optional[str] = 'solid'\n",
    "    line_width: Optional[int] = None\n",
    "    color: Optional[str] = None\n",
    "    marker: Optional[str] = None\n",
    "    marker_size: Optional[int] = None\n",
    "    marker_color: Optional[str] = None\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class ScatterPlotAttributes(DisplayAttributes):\n",
    "    marker: str = 'X'\n",
    "    marker_size: int = 50\n",
    "    marker_color: str = 'red'\n",
    "        \n",
    "        \n",
    "@dataclass\n",
    "class SurfacePlotAttributes(DisplayAttributes):\n",
    "    '''\n",
    "    Attributes:\n",
    "        marker: Adds a marker to each point in x, y, z to show the actual data used for interpolation.  \n",
    "            You can set this to None to turn markers off.\n",
    "        interpolation: Can be ‘linear’, ‘nearest’ or ‘cubic’ for plotting z points between the ones passed in.  \n",
    "            See scipy.interpolate.griddata for details\n",
    "        cmap: Colormap to use (default matplotlib.cm.RdBu_r).  See matplotlib colormap for details\n",
    "    '''\n",
    "    marker: str = 'X'\n",
    "    marker_size: int = 50\n",
    "    marker_color: str = 'red'\n",
    "    interpolation: str = 'linear'\n",
    "    cmap: matplotlib.colors.Colormap = matplotlib.cm.RdBu_r                  \n",
    "        \n",
    "\n",
    "@dataclass\n",
    "class ContourPlotAttributes(DisplayAttributes):\n",
    "    marker: str = 'X'\n",
    "    marker_size: int = 50\n",
    "    marker_color: str = 'red'\n",
    "    interpolation: str = 'linear'\n",
    "    cmap: matplotlib.colors.Colormap = matplotlib.cm.RdBu_r\n",
    "    min_level: float = math.nan\n",
    "    max_level: float = math.nan\n",
    "\n",
    "        \n",
    "@dataclass\n",
    "class CandleStickPlotAttributes(DisplayAttributes):\n",
    "    colorup: str = 'darkgreen'\n",
    "    colordown: str = '#F2583E'\n",
    "\n",
    "        \n",
    "@dataclass\n",
    "class BarPlotAttributes(DisplayAttributes):\n",
    "    color: str = 'red'\n",
    "        \n",
    "\n",
    "@dataclass\n",
    "class FilledLinePlotAttributes(DisplayAttributes):\n",
    "    '''\n",
    "    colorup: Color for bars where close >= open.  Default \"darkgreen\"\n",
    "    colordown: Color for bars where open < close.  Default \"#F2583E\"\n",
    "    '''\n",
    "    positive_color: str = 'blue'\n",
    "    negative_color: str = 'red'\n",
    "        \n",
    "\n",
    "class PlotData:\n",
    "    name: str\n",
    "    display_attributes: DisplayAttributes\n",
    "        \n",
    "\n",
    "class TimePlotData(PlotData):\n",
    "    \n",
    "    timestamps: np.ndarray\n",
    "        \n",
    "    @abstractmethod\n",
    "    def reindex(self, timestamps: np.ndarray, fill: bool) -> None:\n",
    "        pass\n",
    "    \n",
    "\n",
    "class BucketedValues(PlotData):\n",
    "    '''\n",
    "    Data in a subplot where we summarize properties of a numpy array.  \n",
    "        For example, drawing a boxplot with percentiles.  x axis is a categorical\n",
    "    '''\n",
    "    def __init__(self, name: str, \n",
    "                 bucket_names: Sequence[str],\n",
    "                 bucket_values: Sequence[np.ndarray], \n",
    "                 display_attributes: DisplayAttributes = None) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            name: name used for this data in a plot legend\n",
    "            bucket_names: list of strings used on x axis labels\n",
    "            bucket_values: list of numpy arrays that are summarized in this plot\n",
    "        '''\n",
    "        assert isinstance(bucket_names, list) and isinstance(bucket_values, list) and len(bucket_names) == len(bucket_values)\n",
    "        self.display_attributes = BoxPlotAttributes()\n",
    "        self.name = name\n",
    "        self.bucket_names = bucket_names\n",
    "        self.bucket_values = series_to_array(bucket_values)\n",
    "        if display_attributes is None: display_attributes = BoxPlotAttributes()\n",
    "        self.display_attributes = display_attributes\n",
    "        \n",
    "\n",
    "class XYData(PlotData):\n",
    "    '''Data in a subplot that has x and y values that are both arrays of floats'''\n",
    "    def __init__(self, \n",
    "                 name: str,\n",
    "                 x: Union[np.ndarray, pd.Series],\n",
    "                 y: Union[np.ndarray, pd.Series],\n",
    "                 display_attributes: DisplayAttributes = None) -> None:\n",
    "        self.name = name\n",
    "        self.x = np.array(x) if isinstance(x, list) else series_to_array(x)\n",
    "        self.y = np.array(y) if isinstance(y, list) else series_to_array(y)\n",
    "        if display_attributes is None: display_attributes = LinePlotAttributes()\n",
    "        self.display_attributes = display_attributes\n",
    "        \n",
    "\n",
    "class XYZData(PlotData):\n",
    "    '''Data in a subplot that has x, y and z values that are all floats'''\n",
    "    def __init__(self, \n",
    "                 name: str, \n",
    "                 x: Union[np.ndarray, pd.Series],\n",
    "                 y: Union[np.ndarray, pd.Series],\n",
    "                 z: Union[np.ndarray, pd.Series],\n",
    "                 display_attributes: DisplayAttributes = None) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            name: Name to show in plot legend\n",
    "        '''\n",
    "        self.name = name\n",
    "        self.x = np.array(x) if isinstance(x, list) else series_to_array(x)\n",
    "        self.y = np.array(y) if isinstance(y, list) else series_to_array(y)\n",
    "        self.z = np.array(z) if isinstance(z, list) else series_to_array(z)\n",
    "        if display_attributes is None: display_attributes = ContourPlotAttributes()\n",
    "        self.display_attributes = display_attributes\n",
    "    \n",
    "\n",
    "class TimeSeries(TimePlotData):\n",
    "    '''Data in a subplot where x is an array of numpy datetimes and y is a numpy array of floats'''\n",
    "    def __init__(self, \n",
    "                 name: str, \n",
    "                 timestamps: Union[pd.Series, np.ndarray],\n",
    "                 values: Union[pd.Series, np.ndarray],\n",
    "                 display_attributes: DisplayAttributes = None) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            name: Name to show in plot legend\n",
    "        '''\n",
    "        self.name = name\n",
    "        self.timestamps = series_to_array(timestamps)\n",
    "        self.values = series_to_array(values)\n",
    "        if display_attributes is None: display_attributes = LinePlotAttributes()\n",
    "        self.display_attributes = display_attributes\n",
    "        \n",
    "    def reindex(self, timestamps: np.ndarray, fill: bool) -> None:\n",
    "        '''Reindex this series given a new array of timestamps, forward filling holes if fill is set to True'''\n",
    "        s = pd.Series(self.values, index=self.timestamps)\n",
    "        s = s.reindex(timestamps, method='ffill' if fill else None)\n",
    "        self.timestamps = s.index.values\n",
    "        self.values = s.values\n",
    "        \n",
    "\n",
    "class TradeBarSeries(TimePlotData):\n",
    "    '''\n",
    "    Data in a subplot that contains open, high, low, close, volume bars.  volume is optional.\n",
    "    '''\n",
    "    def __init__(self, \n",
    "                 name: str,\n",
    "                 timestamps: np.ndarray,\n",
    "                 o: Optional[np.ndarray],\n",
    "                 h: Optional[np.ndarray],\n",
    "                 l: Optional[np.ndarray],  # noqa: E741: ignore # l ambiguous\n",
    "                 c: Optional[np.ndarray],\n",
    "                 v: np.ndarray = None,\n",
    "                 vwap: np.ndarray = None,\n",
    "                 display_attributes: DisplayAttributes = None) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            name: Name to show in a legend\n",
    "        '''\n",
    "        self.name = name\n",
    "        self.timestamps = timestamps\n",
    "        self.o = o\n",
    "        self.h = h\n",
    "        self.l = l  # noqa: E741: ignore # l ambiguous\n",
    "        self.c = c\n",
    "        self.v = np.ones(len(self.timestamps), dtype=float) * np.nan if v is None else v\n",
    "        self.vwap = np.ones(len(self.timestamps), dtype=float) * np.nan if vwap is None else vwap\n",
    "        \n",
    "        if display_attributes is None: display_attributes = CandleStickPlotAttributes()\n",
    "        self.display_attributes = display_attributes\n",
    "        \n",
    "    def df(self) -> pd.DataFrame:\n",
    "        return pd.DataFrame({'o': self.o, 'h': self.h, 'l': self.l, 'c': self.c, 'v': self.v, 'vwap': self.vwap},  # type: ignore # l ambiguous\n",
    "                            index=self.timestamps)[['o', 'h', 'l', 'c', 'v', 'vwap']]\n",
    "        \n",
    "    def reindex(self, all_timestamps: np.ndarray, fill: bool) -> None:\n",
    "        df = self.df()\n",
    "        df = df.reindex(all_timestamps)\n",
    "        self.timestamps = all_timestamps\n",
    "        for col in df.columns:\n",
    "            setattr(self, col, df[col].values)\n",
    "                \n",
    "\n",
    "class TradeSet(TimePlotData):\n",
    "    '''Data for subplot that contains a set of trades along with marker properties for these trades'''\n",
    "    def __init__(self, \n",
    "                 name: str,\n",
    "                 trades: Sequence[Trade],\n",
    "                 display_attributes: DisplayAttributes = None) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            name: String to display in a subplot legend\n",
    "            trades: List of Trade objects to plot\n",
    "        '''\n",
    "        self.name = name\n",
    "        self.trades = trades\n",
    "        self.timestamps = np.array([trade.timestamp for trade in trades], dtype='M8[ns]')\n",
    "        self.values = np.array([trade.price for trade in trades], dtype=float)\n",
    "        if display_attributes is None:\n",
    "            display_attributes = ScatterPlotAttributes(marker='P', marker_color='red', marker_size=50)\n",
    "        self.display_attributes = display_attributes\n",
    "        \n",
    "    def reindex(self, all_timestamps: np.ndarray, fill: bool) -> None:\n",
    "        s = pd.Series(self.values, index=self.timestamps)\n",
    "        s = s.reindex(all_timestamps, method='ffill' if fill else None)\n",
    "        self.timestamps = s.index.values\n",
    "        self.values = s.values\n",
    "        \n",
    "    def __repr__(self) -> str:\n",
    "        s = ''\n",
    "        for trade in self.trades:\n",
    "            s += f'{trade.timestamp} {trade.qty} {trade.price}\\n'\n",
    "        return s\n",
    "    \n",
    "\n",
    "def draw_poly(ax: mpl.axes.Axes, \n",
    "              left: np.ndarray,\n",
    "              bottom: np.ndarray, \n",
    "              top: np.ndarray, \n",
    "              right: np.ndarray, \n",
    "              facecolor: str, \n",
    "              edgecolor: str, \n",
    "              zorder: int) -> None:\n",
    "    '''Draw a set of polygrams given parrallel numpy arrays of left, bottom, top, right points'''\n",
    "    XY = np.array([[left, left, right, right], [bottom, top, top, bottom]]).T\n",
    "\n",
    "    barpath = path.Path.make_compound_path_from_polys(XY)\n",
    "    \n",
    "    # Clean path to get rid of 0, 0 points.  Seems to be a matplotlib bug.  If we don't ylim lower bound is set to 0\n",
    "    v = []\n",
    "    c = []\n",
    "    for seg in barpath.iter_segments():\n",
    "        vertices, command = seg\n",
    "        if not (vertices[0] == 0. and vertices[1] == 0.):\n",
    "            v.append(vertices)\n",
    "            c.append(command)\n",
    "    cleaned_path = path.Path(v, c)\n",
    "\n",
    "    patch = mptch.PathPatch(cleaned_path, facecolor=facecolor, edgecolor=edgecolor, zorder=zorder)\n",
    "    ax.add_patch(patch)\n",
    "\n",
    "\n",
    "def draw_candlestick(ax: mpl.axes.Axes,\n",
    "                     index: np.ndarray, \n",
    "                     o: np.ndarray, \n",
    "                     h: np.ndarray, \n",
    "                     l: np.ndarray,  # noqa: E741: ignore # l ambiguous\n",
    "                     c: np.ndarray,\n",
    "                     v: Optional[np.ndarray], \n",
    "                     vwap: np.ndarray, \n",
    "                     colorup: str = 'darkgreen', \n",
    "                     colordown: str = '#F2583E') -> None:\n",
    "    '''Draw candlesticks given parrallel numpy arrays of o, h, l, c, v values.  v is optional.  \n",
    "        See TradeBarSeries class __init__ for argument descriptions.'''\n",
    "    width = 0.5\n",
    "    \n",
    "    # Have to do volume first because of a mpl bug with axes fonts if we use make_axes_locatable after plotting on top axis\n",
    "    if v is not None and not np.isnan(v).all(): \n",
    "        divider = make_axes_locatable(ax)\n",
    "        vol_ax = divider.append_axes('bottom', size='25%', sharex=ax, pad=0)\n",
    "        _c = np.nan_to_num(c)\n",
    "        _o = np.nan_to_num(o)\n",
    "        pos = _c >= _o\n",
    "        neg = _c < _o\n",
    "        vol_ax.bar(index[pos], v[pos], color=colorup, width=width)\n",
    "        vol_ax.bar(index[neg], v[neg], color=colordown, width=width)\n",
    "    \n",
    "    offset = width / 2.0\n",
    "    \n",
    "    mask = ~np.isnan(c) & ~np.isnan(o)\n",
    "    mask[mask] &= c[mask] < o[mask]\n",
    "    \n",
    "    left = index - offset\n",
    "    bottom = np.where(mask, o, c)\n",
    "    top = np.where(mask, c, o)\n",
    "    right = left + width\n",
    "    \n",
    "    draw_poly(ax, left[mask], bottom[mask], top[mask], right[mask], colordown, 'k', 100)\n",
    "    draw_poly(ax, left[~mask], bottom[~mask], top[~mask], right[~mask], colorup, 'k', 100)\n",
    "    draw_poly(ax, left + offset, l, h, left + offset, 'k', 'k', 1)\n",
    "    if vwap is not None:\n",
    "        ax.scatter(index, vwap, marker='o', color='orange', zorder=110)\n",
    "        \n",
    "\n",
    "def draw_boxplot(ax: mpl.axes.Axes, \n",
    "                 names: str, \n",
    "                 values: Sequence[np.ndarray], \n",
    "                 proportional_widths: bool = True, \n",
    "                 notched: bool = False, \n",
    "                 show_outliers: bool = True, \n",
    "                 show_means: bool = True, \n",
    "                 show_all: bool = True) -> None:\n",
    "    '''Draw a boxplot.  See BucketedValues class for explanation of arguments'''\n",
    "    outliers = None if show_outliers else ''\n",
    "    meanpointprops = dict(marker='D')\n",
    "    assert(isinstance(values, list) and isinstance(names, list) and len(values) == len(names))\n",
    "    widths = None\n",
    "    \n",
    "    if show_all:\n",
    "        all_values = np.concatenate(values)\n",
    "        values.append(all_values)\n",
    "        names.append('all')\n",
    "    \n",
    "    if proportional_widths:\n",
    "        counts = [len(v) for v in values]\n",
    "        total = float(sum(counts))\n",
    "        widths = [c / total for c in counts]  \n",
    "    \n",
    "    ax.boxplot(values, notch=notched, sym=outliers, showmeans=show_means, meanprops=meanpointprops, widths=widths)\n",
    "    ax.set_xticklabels(names)\n",
    "    \n",
    "    \n",
    "def draw_3d_plot(ax: mpl.axes.Axes,\n",
    "                 x: np.ndarray,\n",
    "                 y: np.ndarray,\n",
    "                 z: np.ndarray,\n",
    "                 plot_type: str = 'contour',\n",
    "                 marker: str = 'X',\n",
    "                 marker_size: int = 50, \n",
    "                 marker_color: str = 'red',\n",
    "                 interpolation: str = 'linear', \n",
    "                 cmap: matplotlib.colors.Colormap = matplotlib.cm.RdBu_r,                  \n",
    "                 min_level: float = math.nan,\n",
    "                 max_level: float = math.nan) -> None:\n",
    "\n",
    "    '''Draw a 3d plot.  See XYZData class for explanation of arguments\n",
    "    \n",
    "    >>> points = np.random.rand(1000, 2)\n",
    "    >>> x = np.random.rand(10)\n",
    "    >>> y = np.random.rand(10)\n",
    "    >>> z = x ** 2 + y ** 2\n",
    "    >>> if has_display():\n",
    "    ...    fig, ax = plt.subplots()\n",
    "    ...    draw_3d_plot(ax, x = x, y = y, z = z, plot_type = 'contour', interpolation = 'linear');\n",
    "    '''\n",
    "    xi = np.linspace(min(x), max(x))\n",
    "    yi = np.linspace(min(y), max(y))\n",
    "    X, Y = np.meshgrid(xi, yi)\n",
    "    Z = griddata((x, y), z, (xi[None, :], yi[:, None]), method=interpolation)\n",
    "    Z = np.nan_to_num(Z)\n",
    "\n",
    "    if plot_type == 'surface':\n",
    "        ax.plot_surface(X, Y, Z, cmap=cmap)\n",
    "        if marker is not None:\n",
    "            ax.scatter(x, y, z, marker=marker, s=marker_size, c=marker_color)\n",
    "        m = cm.ScalarMappable(cmap=cmap)\n",
    "        m.set_array(Z)\n",
    "        plt.colorbar(m, ax=ax)\n",
    "        \n",
    "    elif plot_type == 'contour':\n",
    "        # extract all colors from the  map\n",
    "        cmaplist = [cmap(i) for i in range(cmap.N)]\n",
    "        # create the new map\n",
    "        cmap = cmap.from_list('Custom cmap', cmaplist, cmap.N)\n",
    "        Z = np.ma.masked_array(Z, mask=~np.isfinite(Z))\n",
    "        if math.isnan(min_level): min_level = np.min(Z)\n",
    "        if math.isnan(max_level): max_level = np.max(Z)\n",
    "        # define the bins and normalize and forcing 0 to be part of the colorbar!\n",
    "        bounds = np.arange(min_level, max_level, (max_level - min_level) / cmap.N)\n",
    "        idx = np.searchsorted(bounds, 0)\n",
    "        bounds = np.insert(bounds, idx, 0)\n",
    "        norm = BoundaryNorm(bounds, cmap.N)\n",
    "        cs = ax.contourf(X, Y, Z, cmap=cmap, norm=norm)\n",
    "        \n",
    "        if marker is not None:\n",
    "            x = x[np.isfinite(z)]\n",
    "            y = y[np.isfinite(z)]\n",
    "            ax.scatter(x, y, marker=marker, s=marker_size, c=z[np.isfinite(z)], zorder=10, cmap=cmap)\n",
    "        LABEL_SIZE = 16\n",
    "        ax.tick_params(axis='both', which='major', labelsize=LABEL_SIZE)\n",
    "        ax.tick_params(axis='both', which='minor', labelsize=LABEL_SIZE)\n",
    "        cbar = plt.colorbar(cs, ax=ax)\n",
    "        cbar.ax.tick_params(labelsize=LABEL_SIZE) \n",
    "\n",
    "    else:\n",
    "        raise Exception(f'unknown plot type: {plot_type}')\n",
    "    \n",
    "\n",
    "def _adjust_axis_limit(lim: Tuple[float, float], values: Union[List, np.ndarray]) -> Tuple[float, float]:\n",
    "    '''If values + 10% buffer are outside current xlim or ylim, return expanded xlim or ylim for subplot'''\n",
    "    if isinstance(values, list):\n",
    "        values = np.array(values)\n",
    "    if values.dtype == np.bool_:\n",
    "        values = values.astype(float)\n",
    "    min_val, max_val = np.nanmin(values), np.nanmax(values)\n",
    "    val_range = max_val - min_val\n",
    "    lim_min = np.nanmin(values) - .1 * val_range\n",
    "    lim_max = np.nanmax(values) - .1 * val_range\n",
    "    return (min(lim[0], lim_min), max(lim[1], lim_max))\n",
    "\n",
    "\n",
    "def _plot_data(ax: mpl.axes.Axes, data: PlotData) -> Optional[List[mpl.lines.Line2D]]:\n",
    "    \n",
    "    lines = None  # Return line objects so we can add legends\n",
    "\n",
    "    disp = data.display_attributes\n",
    "    \n",
    "    if isinstance(data, XYData) or isinstance(data, TimeSeries):\n",
    "        x, y = (data.x, data.y) if isinstance(data, XYData) else (np.arange(len(data.timestamps)), data.values)\n",
    "        if isinstance(disp, LinePlotAttributes):\n",
    "            lines, = ax.plot(x, y, linestyle=disp.line_type, linewidth=disp.line_width, color=disp.color)\n",
    "            if disp.marker is not None:  # type: ignore\n",
    "                ax.scatter(x, y, marker=disp.marker, c=disp.marker_color, s=disp.marker_size, zorder=100)\n",
    "        elif isinstance(disp, ScatterPlotAttributes):\n",
    "            lines = ax.scatter(x, y, marker=disp.marker, c=disp.marker_color, s=disp.marker_size, zorder=100)\n",
    "        elif isinstance(disp, BarPlotAttributes):\n",
    "            lines = ax.bar(x, y, color=disp.color)  # type: ignore\n",
    "        elif isinstance(disp, FilledLinePlotAttributes):\n",
    "            x, y = np.nan_to_num(x), np.nan_to_num(y)\n",
    "            pos_values = np.where(y > 0, y, 0)\n",
    "            neg_values = np.where(y < 0, y, 0)\n",
    "            ax.fill_between(x, pos_values, color=disp.positive_color, step='post', linewidth=0.0)\n",
    "            ax.fill_between(x, neg_values, color=disp.negative_color, step='post', linewidth=0.0)\n",
    "        else:\n",
    "            raise Exception(f'unknown plot combination: {type(data)} {type(disp)}')\n",
    "            \n",
    "        # For scatter and filled line, xlim and ylim does not seem to get set automatically\n",
    "        if isinstance(disp, ScatterPlotAttributes) or isinstance(disp, FilledLinePlotAttributes):\n",
    "            xmin, xmax = _adjust_axis_limit(ax.get_xlim(), x)\n",
    "            if not np.isnan(xmin) and not np.isnan(xmax): ax.set_xlim((xmin, xmax))\n",
    "\n",
    "            ymin, ymax = _adjust_axis_limit(ax.get_ylim(), y)\n",
    "            if not np.isnan(ymin) and not np.isnan(ymax): ax.set_ylim((ymin, ymax))\n",
    "                \n",
    "    elif isinstance(data, TradeSet) and isinstance(disp, ScatterPlotAttributes):\n",
    "        lines = ax.scatter(np.arange(len(data.timestamps)), data.values, marker=disp.marker, c=disp.marker_color, s=disp.marker_size, zorder=100)\n",
    "    elif isinstance(data, TradeBarSeries) and isinstance(disp, CandleStickPlotAttributes):\n",
    "        if not (data.o is None or data.h is None or data.l is None or data.c is None):\n",
    "            draw_candlestick(ax, np.arange(len(data.timestamps)), data.o, data.h, data.l, data.c, \n",
    "                             data.v, data.vwap, colorup=disp.colorup, colordown=disp.colordown)\n",
    "    elif isinstance(data, BucketedValues) and isinstance(disp, BoxPlotAttributes):\n",
    "        draw_boxplot(\n",
    "            ax, data.bucket_names, data.bucket_values, disp.proportional_widths, disp.notched,  # type: ignore\n",
    "            disp.show_outliers, disp.show_means, disp.show_all)  # type: ignore\n",
    "    elif isinstance(data, XYZData) and (isinstance(disp, SurfacePlotAttributes) or isinstance(disp, ContourPlotAttributes)):\n",
    "        display_type: str = 'contour' if isinstance(disp, ContourPlotAttributes) else 'surface'\n",
    "        draw_3d_plot(ax, data.x, data.y, data.z, display_type, disp.marker, disp.marker_size, \n",
    "                     disp.marker_color, disp.interpolation, disp.cmap)\n",
    "    else:\n",
    "        raise Exception(f'unknown plot combination: {type(data)} {type(disp)}')\n",
    "\n",
    "    return lines\n",
    "\n",
    "\n",
    "def _draw_date_gap_lines(ax: mpl.axes.Axes, plot_timestamps: np.ndarray) -> None:\n",
    "    '''\n",
    "    Draw vertical lines wherever there are gaps between two timestamps.\n",
    "        i.e., the gap between two adjacent timestamps is more than the minimum gap in the series.\n",
    "    '''\n",
    "    timestamps = mdates.date2num(plot_timestamps)\n",
    "    freq = np.nanmin(np.diff(timestamps))\n",
    "    if freq <= 0: raise Exception('could not infer date frequency')\n",
    "    date_index = np.arange(len(timestamps))\n",
    "    date_diff = np.diff(timestamps)\n",
    "    \n",
    "    xs = []\n",
    "\n",
    "    for i in date_index:\n",
    "        if i < len(date_diff) and date_diff[i] > (freq + 0.000000001):\n",
    "            xs.append(i + 0.5)\n",
    "            \n",
    "    if len(xs) > 20:\n",
    "        return  # Too many lines will clutter the graph\n",
    "    \n",
    "    for x in xs:\n",
    "        ax.axvline(x, linestyle='dashed', color='0.5')\n",
    "            \n",
    "\n",
    "def draw_date_line(ax: mpl.axes.Axes, \n",
    "                   plot_timestamps: np.ndarray, \n",
    "                   date: np.datetime64, \n",
    "                   linestyle: str, \n",
    "                   color: Optional[str]) -> mpl.lines.Line2D:\n",
    "    '''Draw vertical line on a subplot with datetime x axis'''\n",
    "    closest_index = (np.abs(plot_timestamps - date)).argmin()\n",
    "    return ax.axvline(x=closest_index, linestyle=linestyle, color=color)\n",
    "\n",
    "\n",
    "def draw_horizontal_line(ax: mpl.axes.Axes, y: float, linestyle: str, color: Optional[str]) -> mpl.lines.Line2D:\n",
    "    '''Draw horizontal line on a subplot'''\n",
    "    return ax.axhline(y=y, linestyle=linestyle, color=color)\n",
    "\n",
    "\n",
    "def draw_vertical_line(ax: mpl.axes.Axes, x: float, linestyle: str, color: Optional[str]) -> mpl.lines.Line2D:\n",
    "    '''Draw vertical line on a subplot'''\n",
    "    return ax.axvline(x=x, linestyle=linestyle, color=color)\n",
    "           \n",
    "\n",
    "def get_date_formatter(plot_timestamps: np.ndarray, date_format: Optional[str]) -> DateFormatter:\n",
    "    '''Create an appropriate DateFormatter for x axis labels.  \n",
    "    If date_format is set to None, figures out an appropriate date format based on the range of timestamps passed in'''\n",
    "    num_timestamps = mdates.date2num(plot_timestamps)\n",
    "    if date_format is not None: return DateFormatter(num_timestamps, fmt=date_format)\n",
    "    date_range = num_timestamps[-1] - num_timestamps[0]\n",
    "    if date_range > 252:\n",
    "        date_format = '%d-%b-%Y'\n",
    "    elif date_range > 7:\n",
    "        date_format = '%b %d'\n",
    "    elif date_range > 1:\n",
    "        date_format = '%d %H:%M'\n",
    "    else:\n",
    "        date_format = '%H:%M:%S'\n",
    "        \n",
    "    formatter = DateFormatter(num_timestamps, fmt=date_format)\n",
    "    return formatter\n",
    "    \n",
    "\n",
    "class Subplot:\n",
    "    '''A top level plot contains a list of subplots, each of which contain a list of data objects to draw'''\n",
    "    def __init__(self, \n",
    "                 data_list: Union[PlotData, Sequence[PlotData]],\n",
    "                 secondary_y: Sequence[str] = None,\n",
    "                 title: str = None,\n",
    "                 xlabel: str = None,\n",
    "                 ylabel: str = None, \n",
    "                 zlabel: str = None,\n",
    "                 date_lines: Sequence[DateLine] = None,\n",
    "                 horizontal_lines: Sequence[HorizontalLine] = None,\n",
    "                 vertical_lines: Sequence[VerticalLine] = None,\n",
    "                 xlim: Union[Tuple[float, float], Tuple[np.datetime64, np.datetime64]] = None, \n",
    "                 ylim: Union[Tuple[float, float], Tuple[np.datetime64, np.datetime64]] = None, \n",
    "                 height_ratio: float = 1.0,\n",
    "                 display_legend: bool = True, \n",
    "                 legend_loc: str = 'best', \n",
    "                 log_y: bool = False, \n",
    "                 y_tick_format: str = None) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            data_list: A list of objects to draw.  Each element can contain XYData, XYZData, TimeSeries, TradeBarSeries, \n",
    "                BucketedValues or TradeSet\n",
    "            secondary_y: A list of objects to draw on the secondary y axis\n",
    "            title: Title to show for this subplot. Default None\n",
    "            zlabel: Only applicable to 3d subplots.  Default None\n",
    "            date_lines: A list of DateLine objects to draw as vertical lines. Only applicable when x axis is datetime.  \n",
    "                Default None\n",
    "            horizontal_lines: A list of HorizontalLine objects to draw on the plot.  Default None\n",
    "            vertical_lines: A list of VerticalLine objects to draw on the plot\n",
    "            xlim: x limits for the plot as a tuple of numpy datetime objects when x-axis is datetime, \n",
    "                or tuple of floats. Default None\n",
    "            ylim: y limits for the plot.  Tuple of floats.  Default None\n",
    "            height_ratio: If you have more than one subplot on a plot, use height ratio to determine how high each subplot should be.\n",
    "                For example, if you set height_ratio = 0.75 for the first subplot and 0.25 for the second, \n",
    "                the first will be 3 times taller than the second one. Default 1.0\n",
    "            display_legend: Whether to show a legend on the plot.  Default True\n",
    "            legend_loc: Location for the legend. Default 'best'\n",
    "            log_y: Whether the y axis should be logarithmic.  Default False\n",
    "            y_tick_format: Format string to use for y axis labels.  For example, you can decide to \n",
    "              use fixed notation instead of scientific notation or change number of decimal places shown.  Default None\n",
    "        '''\n",
    "        if not isinstance(data_list, collections.abc.Sequence): data_list = [data_list]\n",
    "            \n",
    "        self.time_plot = all([isinstance(data, TimePlotData) for data in data_list])\n",
    "        if self.time_plot and any([not isinstance(data, TimePlotData) for data in data_list]):\n",
    "            raise Exception('cannot add a non date subplot on a subplot which has time series plots')\n",
    "        if not self.time_plot and date_lines is not None: \n",
    "            raise Exception('date lines can only be specified on a time series subplot')\n",
    "        \n",
    "        self.is_3d = any([isinstance(data.display_attributes, SurfacePlotAttributes) for data in data_list])\n",
    "        if self.is_3d and any([not isinstance(data.display_attributes, SurfacePlotAttributes) for data in data_list]):\n",
    "            raise Exception('cannot combine 2d plot and 3d subplots on the same Subplot')\n",
    "\n",
    "        self.data_list = data_list\n",
    "        self.secondary_y = [] if secondary_y is None else secondary_y\n",
    "        self.date_lines = [] if date_lines is None else date_lines\n",
    "        self.horizontal_lines = [] if horizontal_lines is None else horizontal_lines\n",
    "        self.vertical_lines = [] if vertical_lines is None else vertical_lines\n",
    "        self.title = title\n",
    "        self.xlabel = xlabel\n",
    "        self.ylabel = ylabel\n",
    "        self.zlabel = zlabel\n",
    "        self.ylim = ylim\n",
    "        self.height_ratio = height_ratio\n",
    "        self.display_legend = display_legend\n",
    "        self.legend_loc = legend_loc\n",
    "        self.log_y = log_y\n",
    "        self.y_tick_format = y_tick_format\n",
    "        \n",
    "    def _resample(self, sampling_frequency: Optional[str]) -> None:\n",
    "        if sampling_frequency is None: return None\n",
    "        for data in self.data_list:\n",
    "            if isinstance(data, TimeSeries) or isinstance(data, TradeSet):\n",
    "                data.timestamps, data.values = resample_ts(data.timestamps, data.values, sampling_frequency)\n",
    "            elif isinstance(data, TradeBarSeries):\n",
    "                df_dict = {}\n",
    "                cols = ['timestamps', 'o', 'h', 'l', 'c', 'v', 'vwap']\n",
    "                for col in cols:\n",
    "                    val = getattr(data, col)\n",
    "                    if val is not None:\n",
    "                        df_dict[col] = val\n",
    "                df = pd.DataFrame(df_dict)\n",
    "                df = df.set_index('timestamps')\n",
    "                df = resample_trade_bars(df, sampling_frequency)\n",
    "                for col in cols:\n",
    "                    if col in df:\n",
    "                        setattr(data, col, df[col].values)\n",
    "            else:\n",
    "                raise Exception(f'unknown type: {data}')\n",
    "        \n",
    "    def get_all_timestamps(self, date_range: Tuple[Optional[np.datetime64], Optional[np.datetime64]]) -> np.ndarray:\n",
    "        timestamps_list = [data.timestamps for data in self.data_list if isinstance(data, TimePlotData)]\n",
    "        all_timestamps = np.array(reduce(np.union1d, timestamps_list))\n",
    "        if date_range is not None and date_range[0] is not None and date_range[1] is not None:\n",
    "            all_timestamps = all_timestamps[(all_timestamps >= date_range[0]) & (all_timestamps <= date_range[1])]\n",
    "        return all_timestamps\n",
    "    \n",
    "    def _reindex(self, all_timestamps: np.ndarray) -> None:\n",
    "        for data in self.data_list:\n",
    "            if not isinstance(data, TimePlotData): continue\n",
    "            disp = data.display_attributes\n",
    "            fill = not isinstance(data, TradeSet) and not (isinstance(disp, BarPlotAttributes) or isinstance(disp, ScatterPlotAttributes))\n",
    "            data.reindex(all_timestamps, fill=fill)\n",
    "            \n",
    "    def _draw(self, ax: mpl.axes.Axes, plot_timestamps: Optional[np.ndarray], date_formatter: Optional[DateFormatter]) -> None:\n",
    "        \n",
    "        if self.time_plot:\n",
    "            assert plot_timestamps is not None\n",
    "            self._reindex(plot_timestamps)\n",
    "            if date_formatter is not None: ax.xaxis.set_major_formatter(date_formatter)\n",
    "        \n",
    "        lines = []\n",
    "        \n",
    "        ax2 = None\n",
    "        if self.secondary_y is not None and len(self.secondary_y):\n",
    "            ax2 = ax.twinx()\n",
    "        \n",
    "        for data in self.data_list:\n",
    "            if ax2 and data.name in self.secondary_y:\n",
    "                line = _plot_data(ax2, data)\n",
    "            else:\n",
    "                line = _plot_data(ax, data)\n",
    "            lines.append(line)\n",
    "            \n",
    "        for date_line in self.date_lines:  # vertical lines on time plot\n",
    "            assert plot_timestamps is not None\n",
    "            line = draw_date_line(ax, plot_timestamps, date_line.date, date_line.line_type, date_line.color)\n",
    "            if date_line.name is not None: lines.append(line)\n",
    "                \n",
    "        for horizontal_line in self.horizontal_lines:\n",
    "            line = draw_horizontal_line(ax, horizontal_line.y, horizontal_line.line_type, horizontal_line.color)\n",
    "            if horizontal_line.name is not None: lines.append(line)\n",
    "                \n",
    "        for vertical_line in self.vertical_lines:\n",
    "            line = draw_vertical_line(ax, vertical_line.x, vertical_line.line_type, vertical_line.color)\n",
    "            if vertical_line.name is not None: lines.append(line)\n",
    "                \n",
    "        self.legend_names = [data.name for data in self.data_list]\n",
    "        self.legend_names += [date_line.name for date_line in self.date_lines if date_line.name is not None]\n",
    "        self.legend_names += [horizontal_line.name for horizontal_line in self.horizontal_lines if horizontal_line.name is not None]\n",
    "        self.legend_names += [vertical_line.name for vertical_line in self.vertical_lines if vertical_line.name is not None]\n",
    "        \n",
    "        if self.ylim: ax.set_ylim(self.ylim)\n",
    "        if (len(self.data_list) > 1 or len(self.date_lines)) and self.display_legend: \n",
    "            ax.legend([line for line in lines if line is not None],\n",
    "                      [self.legend_names[i] for i, line in enumerate(lines) if line is not None], loc=self.legend_loc)\n",
    "            \n",
    "        if self.log_y: \n",
    "            ax.set_yscale('log')\n",
    "            ax.yaxis.set_major_locator(mtick.AutoLocator())\n",
    "            ax.yaxis.set_minor_locator(mtick.NullLocator())\n",
    "        if self.y_tick_format:\n",
    "            ax.yaxis.set_major_formatter(mtick.StrMethodFormatter(self.y_tick_format))\n",
    "        \n",
    "        if self.title: ax.set_title(self.title)\n",
    "        if self.xlabel: ax.set_xlabel(self.xlabel)\n",
    "        if self.ylabel: ax.set_ylabel(self.ylabel)\n",
    "        if self.zlabel: ax.set_zlabel(self.zlabel)\n",
    "            \n",
    "        ax.autoscale_view()\n",
    "            \n",
    "\n",
    "class Plot:\n",
    "    '''Top level plot containing a list of subplots to draw'''\n",
    "    def __init__(self, \n",
    "                 subplot_list: Sequence[Subplot],  \n",
    "                 title: str = None, \n",
    "                 figsize: Tuple[float, float] = (15, 8), \n",
    "                 date_range: Union[Tuple[str, str], Tuple[Optional[np.datetime64], Optional[np.datetime64]]] = None,\n",
    "                 date_format: str = None, \n",
    "                 sampling_frequency: str = None, \n",
    "                 show_grid: bool = True, \n",
    "                 show_date_gaps: bool = True, \n",
    "                 hspace: Optional[float] = 0.15) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            subplot_list: List of Subplot objects to draw\n",
    "            title: Title for this plot.  Default None\n",
    "            figsize: Figure size.  Default (15, 8)\n",
    "            date_range: Tuple of strings or numpy datetime64 limiting timestamps to draw.  e.g. (\"2018-01-01 14:00\", \"2018-01-05\"). Default None\n",
    "            date_format: Date format to use for x-axis\n",
    "            sampling_frequency: Set this to downsample subplots that have a datetime x axis.  \n",
    "              For example, if you have minute bar data, you might want to subsample to hours if the plot is too crowded.\n",
    "              See pandas time frequency strings for possible values.  Default None\n",
    "            show_grid: If set to True, show a grid on the subplots.  Default True\n",
    "            show_date_gaps: If set to True, then when there is a gap between timestamps will draw a dashed vertical line. \n",
    "              For example, you may have minute bars and a gap between end of trading day and beginning of next day. \n",
    "              Even if set to True, this will turn itself off if there are too many gaps to avoid clutter.  Default True\n",
    "            hspace: Height (vertical) space between subplots.  Default 0.15\n",
    "        '''\n",
    "        if isinstance(subplot_list, Subplot): subplot_list = [subplot_list]\n",
    "        assert(len(subplot_list))\n",
    "        self.subplot_list = subplot_list\n",
    "        self.title = title\n",
    "        self.figsize = figsize\n",
    "        self.date_range = strtup2date(date_range)\n",
    "        self.date_format = date_format\n",
    "        self.sampling_frequency = sampling_frequency\n",
    "        self.show_date_gaps = show_date_gaps\n",
    "        self.show_grid = show_grid\n",
    "        self.hspace = hspace\n",
    "        \n",
    "    def _get_plot_timestamps(self) -> Optional[np.ndarray]:\n",
    "        timestamps_list = []\n",
    "        for subplot in self.subplot_list:\n",
    "            if not subplot.time_plot: continue\n",
    "            subplot._resample(self.sampling_frequency)\n",
    "            timestamps_list.append(subplot.get_all_timestamps(self.date_range))\n",
    "        if not len(timestamps_list): return None\n",
    "        plot_timestamps = np.array(reduce(np.union1d, timestamps_list))\n",
    "        return plot_timestamps\n",
    "\n",
    "    def draw(self, check_data_size: bool = True) -> Optional[Tuple[mpl.figure.Figure, mpl.axes.Axes]]:\n",
    "        '''Draw the subplots.\n",
    "        \n",
    "        Args:\n",
    "            check_data_size: If set to True, will not plot if there are > 100K points to avoid locking up your computer for a long time.\n",
    "              Default True\n",
    "        '''\n",
    "        if not has_display():\n",
    "            print('no display found, cannot plot')\n",
    "            return None\n",
    "        \n",
    "        plot_timestamps = self._get_plot_timestamps()\n",
    "        if check_data_size and plot_timestamps is not None and len(plot_timestamps) > 100000:\n",
    "            raise Exception(f'trying to plot large data set with {len(plot_timestamps)} points, reduce date range or turn check_data_size flag off')\n",
    "            \n",
    "        date_formatter = None\n",
    "        if plot_timestamps is not None: \n",
    "            date_formatter = get_date_formatter(plot_timestamps, self.date_format)\n",
    "        height_ratios = [subplot.height_ratio for subplot in self.subplot_list]\n",
    "        \n",
    "        fig = plt.figure(figsize=self.figsize)\n",
    "        gs = gridspec.GridSpec(len(self.subplot_list), 1, height_ratios=height_ratios, hspace=self.hspace)\n",
    "        axes = []\n",
    "        \n",
    "        for i, subplot in enumerate(self.subplot_list):\n",
    "            if subplot.is_3d:\n",
    "                ax = plt.subplot(gs[i], projection='3d')\n",
    "            else:\n",
    "                ax = plt.subplot(gs[i])\n",
    "            axes.append(ax)\n",
    "            \n",
    "        time_axes = [axes[i] for i, s in enumerate(self.subplot_list) if s.time_plot]\n",
    "        if len(time_axes):\n",
    "            time_axes[0].get_shared_x_axes().join(*time_axes)\n",
    "            \n",
    "        for i, subplot in enumerate(self.subplot_list):\n",
    "            subplot._draw(axes[i], plot_timestamps, date_formatter)\n",
    "            \n",
    "        if self.title: axes[0].set_title(self.title)\n",
    "\n",
    "        # We may have added new axes in candlestick plot so get list of axes again\n",
    "        ax_list = fig.axes\n",
    "        for ax in ax_list:\n",
    "            if self.show_grid: ax.grid(linestyle='dotted')\n",
    "                \n",
    "        for ax in ax_list:\n",
    "            if ax not in axes: time_axes.append(ax)\n",
    "                \n",
    "        for ax in time_axes:\n",
    "            if self.show_date_gaps and plot_timestamps is not None: _draw_date_gap_lines(ax, plot_timestamps)\n",
    "                \n",
    "        for ax in ax_list:\n",
    "            ax.autoscale_view()\n",
    "            \n",
    "        return fig, ax_list\n",
    "                \n",
    "\n",
    "def _group_trades_by_reason_code(trades: Sequence[Trade]) -> Mapping[str, List[Trade]]:\n",
    "    trade_groups: MutableMapping[str, List[Trade]] = collections.defaultdict(list)\n",
    "    for trade in trades:\n",
    "        trade_groups[trade.order.reason_code].append(trade)\n",
    "    return trade_groups\n",
    "\n",
    "\n",
    "def trade_sets_by_reason_code(trades: List[Trade], \n",
    "                              marker_props: Mapping[str, Mapping] = ReasonCode.MARKER_PROPERTIES, \n",
    "                              remove_missing_properties: bool = True) -> List[TradeSet]:\n",
    "    '''\n",
    "    Returns a list of TradeSet objects.  Each TradeSet contains trades with a different reason code.  The markers for each TradeSet\n",
    "    are set by looking up marker properties for each reason code using the marker_props argument:\n",
    "    \n",
    "    Args:\n",
    "        trades: We look up reason codes using the reason code on the corresponding orders\n",
    "        marker_props: Dictionary from reason code string -> dictionary of marker properties.  \n",
    "            See ReasonCode.MARKER_PROPERTIES for example.  Default ReasonCode.MARKER_PROPERTIES\n",
    "        remove_missing_properties: If set, we remove any reason codes that dont' have marker properties set.\n",
    "            Default True\n",
    "     '''\n",
    "    trade_groups = _group_trades_by_reason_code(trades)\n",
    "    tradesets = []\n",
    "    for reason_code, trades in trade_groups.items():\n",
    "        if reason_code in marker_props:\n",
    "            mp = marker_props[reason_code]\n",
    "            disp = ScatterPlotAttributes(marker=mp['symbol'], marker_color=mp['color'], marker_size=mp['size'])\n",
    "            tradeset = TradeSet(reason_code, trades, display_attributes=disp)\n",
    "        elif remove_missing_properties: \n",
    "            continue\n",
    "        else:\n",
    "            tradeset = TradeSet(reason_code, trades)\n",
    "        tradesets.append(tradeset)\n",
    "    return tradesets \n",
    "\n",
    "\n",
    "def test_plot() -> None:\n",
    "    \n",
    "    class MockOrder:\n",
    "        def __init__(self, reason_code: str) -> None:\n",
    "            self.reason_code = reason_code\n",
    "    \n",
    "    class MockTrade:\n",
    "        def __init__(self, timestamp: np.datetime64, qty: float, price: float, reason_code: str) -> None:\n",
    "            self.timestamp = timestamp\n",
    "            self.qty = qty\n",
    "            self.price = price\n",
    "            self.order = MockOrder(reason_code) \n",
    "            \n",
    "        def __repr__(self) -> str:\n",
    "            return f'{self.timestamp} {self.qty} {self.price}'\n",
    "        \n",
    "    np.random.seed(0)\n",
    "            \n",
    "    timestamps = np.array(['2018-01-08 15:00:00', '2018-01-09 15:00:00', '2018-01-10 15:00:00', '2018-01-11 15:00:00'], dtype='M8[ns]')\n",
    "    pnl_timestamps = np.array(['2018-01-08 15:00:00', '2018-01-09 14:00:00', '2018-01-10 15:00:00', '2018-01-15 15:00:00'], dtype='M8[ns]')\n",
    "    \n",
    "    positions = (pnl_timestamps, np.array([0., 5., 0., -10.]))\n",
    "    \n",
    "    trade_timestamps = np.array(['2018-01-09 14:00:00', '2018-01-10 15:00:00', '2018-01-15 15:00:00'], dtype='M8[ns]')\n",
    "    trade_price = [9., 10., 9.5]\n",
    "    trade_qty = [5, -5, -10]\n",
    "    reason_codes = [ReasonCode.ENTER_LONG, ReasonCode.EXIT_LONG, ReasonCode.ENTER_SHORT]\n",
    "    trades = [MockTrade(trade_timestamps[i], trade_qty[i], trade_price[i], reason_codes[i]) for i, d in enumerate(trade_timestamps)]\n",
    "    \n",
    "    disp = LinePlotAttributes(line_type='--')\n",
    "    \n",
    "    tb_series = TradeBarSeries(\n",
    "        'price', timestamps=timestamps,           \n",
    "        o=np.array([8.9, 9.1, 9.3, 8.6]),\n",
    "        h=np.array([9.0, 9.3, 9.4, 8.7]),\n",
    "        l=np.array([8.8, 9.0, 9.2, 8.4]),  # noqa: E741 # ambiguous l\n",
    "        c=np.array([8.95, 9.2, 9.35, 8.5]),\n",
    "        v=np.array([200, 100, 150, 300]),\n",
    "        vwap=np.array([8.9, 9.15, 9.3, 8.55]))\n",
    "\n",
    "    ind_subplot = Subplot([\n",
    "        TimeSeries('slow_support', timestamps=timestamps, values=np.array([8.9, 8.9, 9.1, 9.1]), display_attributes=disp),\n",
    "        TimeSeries('fast_support', timestamps=timestamps, values=np.array([8.9, 9.0, 9.1, 9.2]), display_attributes=disp),\n",
    "        TimeSeries('slow_resistance', timestamps=timestamps, values=np.array([9.2, 9.2, 9.4, 9.4]), display_attributes=disp),\n",
    "        TimeSeries('fast_resistance', timestamps=timestamps, values=np.array([9.2, 9.3, 9.4, 9.5]), display_attributes=disp),\n",
    "        TimeSeries('secondary_y_test', timestamps=timestamps, values=np.array([150, 160, 162, 135]), display_attributes=disp),\n",
    "        tb_series\n",
    "    ] + trade_sets_by_reason_code(trades),  # type: ignore  # mypy complains about adding heterogeneous lists \n",
    "        secondary_y=['secondary_y_test'], \n",
    "        ylabel=\"Price\", height_ratio=0.3) \n",
    "    \n",
    "    sig_subplot = Subplot(TimeSeries('trend', timestamps=timestamps, values=np.array([1, 1, -1, -1])), height_ratio=0.1, ylabel='Trend')\n",
    "    \n",
    "    equity_subplot = Subplot(\n",
    "        TimeSeries('equity', timestamps=pnl_timestamps, values=[1.0e6, 1.1e6, 1.2e6, 1.3e6]), \n",
    "        height_ratio=0.1, ylabel='Equity', \n",
    "        date_lines=[DateLine(date=np.datetime64('2018-01-09 14:00:00'), name='drawdown', color='red'),\n",
    "                    DateLine(date=np.datetime64('2018-01-10 15:00:00'), color='red')], \n",
    "        horizontal_lines=[HorizontalLine(y=0, name='zero', color='green')])\n",
    "    \n",
    "    pos_subplot = Subplot(\n",
    "        TimeSeries('position', timestamps=positions[0], values=positions[1], display_attributes=FilledLinePlotAttributes()), \n",
    "        height_ratio=0.1, ylabel='Position')\n",
    "    \n",
    "    annual_returns_subplot = Subplot(\n",
    "        BucketedValues('annual returns', ['2017', '2018'], \n",
    "                       bucket_values=[np.random.normal(0, 1, size=(250,)), np.random.normal(0, 1, size=(500,))]),\n",
    "        height_ratio=0.1, ylabel='Annual Returns')\n",
    "    \n",
    "    x = np.random.rand(10)\n",
    "    y = np.random.rand(10)\n",
    "    xy_subplot = Subplot(XYData('2d test', x, y, display_attributes=ScatterPlotAttributes(marker='X')), \n",
    "                         xlabel='x', ylabel='y', height_ratio=0.2, title='XY Plot')\n",
    "    \n",
    "    z = x ** 2 + y ** 2\n",
    "    xyz_subplot = Subplot(XYZData('3d test', x, y, z, display_attributes=SurfacePlotAttributes()), \n",
    "                          xlabel='x', ylabel='y', zlabel='z', height_ratio=0.3)\n",
    "    \n",
    "    xyz_contour = Subplot(XYZData('Contour test', x, y, z, display_attributes=ContourPlotAttributes()), \n",
    "                          xlabel='x', ylabel='y', height_ratio=0.3)\n",
    "    \n",
    "    subplot_list = [ind_subplot, sig_subplot, pos_subplot, equity_subplot, \n",
    "                    annual_returns_subplot, xy_subplot, xyz_contour, xyz_subplot]\n",
    "    plot = Plot(subplot_list, figsize=(20, 20), title='Plot Test', hspace=0.35)\n",
    "    plot.draw()\n",
    "    \n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_plot()\n",
    "    import doctest\n",
    "    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)"
   ]
  }
 ],
 "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
}
