{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-10-22T21:23:02.599431Z",
     "start_time": "2019-10-22T21:23:01.164338Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "flake8 success\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from functools import reduce\n",
    "from pyqstrat.evaluator import compute_return_metrics, display_return_metrics, plot_return_metrics\n",
    "from pyqstrat.strategy import Strategy\n",
    "from pyqstrat.pq_utils import str2date\n",
    "\n",
    "from typing import Sequence, MutableMapping, Mapping, Tuple\n",
    "\n",
    "\n",
    "class Portfolio:\n",
    "    '''A portfolio contains one or more strategies that run concurrently so you can test running strategies that are uncorrelated together.'''\n",
    "    def __init__(self, name: str = 'main') -> None:\n",
    "        '''Args:\n",
    "            name: String used for displaying this portfolio\n",
    "        '''\n",
    "        self.name = name\n",
    "        self.strategies: MutableMapping[str, Strategy] = {}\n",
    "        \n",
    "    def add_strategy(self, name: str, strategy: Strategy) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            name: Name of the strategy\n",
    "            strategy: Strategy instance\n",
    "        '''\n",
    "        self.strategies[name] = strategy\n",
    "        strategy.name = name\n",
    "        \n",
    "    def run_indicators(self, strategy_names: Sequence[str] = None) -> None:\n",
    "        '''Compute indicators for the strategies specified\n",
    "        \n",
    "        Args:\n",
    "            strategy_names: By default this is set to None and we use all strategies.\n",
    "        '''\n",
    "        if strategy_names is None: strategy_names = list(self.strategies.keys())\n",
    "        if len(strategy_names) == 0: raise Exception('a portfolio must have at least one strategy')\n",
    "        for name in strategy_names: self.strategies[name].run_indicators()\n",
    "                \n",
    "    def run_signals(self, strategy_names: Sequence[str] = None) -> None:\n",
    "        '''Compute signals for the strategies specified.  Must be called after run_indicators\n",
    "        \n",
    "        Args:\n",
    "            strategy_names: By default this is set to None and we use all strategies.\n",
    "        '''\n",
    "        if strategy_names is None: strategy_names = list(self.strategies.keys())\n",
    "        if len(strategy_names) == 0: raise Exception('a portfolio must have at least one strategy')\n",
    "        for name in strategy_names: self.strategies[name].run_signals()\n",
    "            \n",
    "    def _generate_order_iterations(self, \n",
    "                                   strategies: Sequence[Strategy], \n",
    "                                   start_date: np.datetime64, \n",
    "                                   end_date: np.datetime64) -> Tuple[np.ndarray, Sequence[Tuple[Strategy, np.ndarray]]]:\n",
    "        '''\n",
    "        >>> class Strategy:\n",
    "        ...    def __init__(self, num): \n",
    "        ...        self.timestamps = [\n",
    "        ...            np.array(['2018-01-01', '2018-01-02', '2018-01-03'], dtype='M8[D]'),\n",
    "        ...            np.array(['2018-01-02', '2018-01-03', '2018-01-04'], dtype='M8[D]')][num]\n",
    "        ...        self.num = num\n",
    "        ...    def _generate_order_iterations(self, start_date, end_date):\n",
    "        ...        pass\n",
    "        ...    def __repr__(self):\n",
    "        ...        return f'{self.num}'\n",
    "        >>> all_timestamps, orders_iter = Portfolio._generate_order_iterations(None, [Strategy(0), Strategy(1)], None, None)\n",
    "        >>> assert(all(all_timestamps == np.array(['2018-01-01', '2018-01-02', '2018-01-03','2018-01-04'], dtype = 'M8[D]')))\n",
    "        >>> assert(all(orders_iter[0][1] == np.array([0, 1, 2, 3])))\n",
    "        >>> assert(all(orders_iter[1][1] == np.array([0, 0, 1, 2])))\n",
    "        '''\n",
    "        if strategies is None: strategies = self.strategies.values\n",
    "            \n",
    "        timestamps_list = [strategy.timestamps for strategy in strategies]\n",
    "        \n",
    "        all_timestamps = np.array(reduce(np.union1d, timestamps_list))\n",
    "        \n",
    "        if start_date is not None:\n",
    "            all_timestamps = all_timestamps[(all_timestamps >= start_date)]\n",
    "        if end_date is not None:\n",
    "            all_timestamps = all_timestamps[(all_timestamps <= end_date)]\n",
    "        \n",
    "        iterations = []\n",
    "        \n",
    "        for strategy in strategies:\n",
    "            indices = np.searchsorted(strategy.timestamps, all_timestamps)\n",
    "            iterations.append((strategy, indices))\n",
    "            strategy._generate_order_iterations(start_date=start_date, end_date=end_date)\n",
    "            \n",
    "        return all_timestamps, iterations\n",
    "                \n",
    "    def run_rules(self, strategy_names: Sequence[str] = None, start_date: np.datetime64 = None, end_date: np.datetime64 = None) -> None:\n",
    "        '''Run rules for the strategies specified.  Must be called after run_indicators and run_signals.  \n",
    "          See run function for argument descriptions\n",
    "        '''\n",
    "        start_date, end_date = str2date(start_date), str2date(end_date)\n",
    "        if strategy_names is None: strategy_names = list(self.strategies.keys())\n",
    "        if len(strategy_names) == 0: raise Exception('a portfolio must have at least one strategy')\n",
    "\n",
    "        strategies = [self.strategies[key] for key in strategy_names]\n",
    "        \n",
    "        min_date = min([strategy.timestamps[0] for strategy in strategies])\n",
    "        if start_date: min_date = max(min_date, start_date)\n",
    "        max_date = max([strategy.timestamps[-1] for strategy in strategies])\n",
    "        if end_date: max_date = min(max_date, end_date)\n",
    "            \n",
    "        all_timestamps, iterations = self._generate_order_iterations(strategies, start_date, end_date)\n",
    "        \n",
    "        for i, timestamp in enumerate(all_timestamps):\n",
    "            for (strategy, indices) in iterations:\n",
    "                # index into strategy timestamps\n",
    "                idx = indices[i]\n",
    "                if idx != len(strategy.timestamps) and strategy.timestamps[idx] == timestamp:\n",
    "                    strategy._run_iteration(idx)\n",
    "                    \n",
    "        # Make sure we calc to the end for each strategy\n",
    "        for strategy in strategies:\n",
    "            strategy.account.calc(strategy.timestamps[-1])\n",
    "                \n",
    "    def run(self, strategy_names: Sequence[str] = None, start_date: np.datetime64 = None, end_date: np.datetime64 = None) -> None:\n",
    "        '''\n",
    "        Run indicators, signals and rules.\n",
    "        \n",
    "        Args:\n",
    "            strategy_names: A list of strategy names.  By default this is set to None and we use all strategies.\n",
    "            start_date: Run rules starting from this date.  \n",
    "              Sometimes we have a few strategies in a portfolio that need different lead times before they are ready to trade\n",
    "              so you can set this so they are all ready by this date.  Default None\n",
    "            end_date: Don't run rules after this date.  Default None\n",
    "         '''\n",
    "        start_date, end_date = str2date(start_date), str2date(end_date)\n",
    "        self.run_indicators()\n",
    "        self.run_signals()\n",
    "        self.run_rules(strategy_names, start_date, end_date)\n",
    "        \n",
    "    def df_returns(self, sampling_frequency: str = 'D', strategy_names: Sequence[str] = None) -> pd.DataFrame:\n",
    "        '''\n",
    "        Return dataframe containing equity and returns with a date index.  Equity and returns are combined from all strategies passed in.\n",
    "        \n",
    "        Args:\n",
    "            sampling_frequency: Date frequency for rows.  Default 'D' for daily so we will have one row per day\n",
    "            strategy_names: By default this is set to None and we use all strategies.\n",
    "        '''\n",
    "        if strategy_names is None: strategy_names = list(self.strategies.keys())\n",
    "        if len(strategy_names) == 0: raise Exception('portfolio must have at least one strategy')\n",
    "        equity_list = []\n",
    "        for name in strategy_names:\n",
    "            equity = self.strategies[name].df_returns(sampling_frequency=sampling_frequency)[['timestamp', 'equity']]\n",
    "            equity.columns = ['timestamp', name]\n",
    "            equity = equity.set_index('timestamp')\n",
    "            equity_list.append(equity)\n",
    "        df = pd.concat(equity_list, axis=1)\n",
    "        df['equity'] = df.sum(axis=1)\n",
    "        df['ret'] = df.equity.pct_change()\n",
    "        return df.reset_index()\n",
    "        \n",
    "    def evaluate_returns(self, \n",
    "                         sampling_frequency: str = 'D', \n",
    "                         strategy_names: Sequence[str] = None, \n",
    "                         plot: bool = True, \n",
    "                         float_precision: int = 4) -> Mapping:\n",
    "        '''Returns a dictionary of common return metrics.\n",
    "        \n",
    "        Args:\n",
    "            sampling_frequency: Date frequency.  Default 'D' for daily so we downsample to daily returns before computing metrics\n",
    "            strategy_names: By default this is set to None and we use all strategies.\n",
    "            plot: If set to True, display plots of equity, drawdowns and returns.  Default False\n",
    "            float_precision: Number of significant figures to show in returns.  Default 4\n",
    "        '''\n",
    "        returns = self.df_returns(sampling_frequency, strategy_names)\n",
    "        ev = compute_return_metrics(returns.timestamp.values, returns.ret.values, returns.equity.values[0])\n",
    "        display_return_metrics(ev.metrics(), float_precision=float_precision)\n",
    "        if plot: plot_return_metrics(ev.metrics())\n",
    "        return ev.metrics()\n",
    "    \n",
    "    def plot(self, \n",
    "             sampling_frequency: str = 'D', \n",
    "             strategy_names: Sequence[str] = None) -> None:\n",
    "        '''Display plots of equity, drawdowns and returns\n",
    "        \n",
    "        Args:\n",
    "            sampling_frequency: Date frequency.  Default 'D' for daily so we downsample to daily returns before computing metrics\n",
    "            strategy_names: A list of strategy names.  By default this is set to None and we use all strategies.\n",
    "        '''\n",
    "        returns = self.df_returns(sampling_frequency, strategy_names)\n",
    "        timestamps = returns.timestamp.values\n",
    "        ev = compute_return_metrics(timestamps, returns.ret.values, returns.equity.values[0])\n",
    "        plot_return_metrics(ev.metrics())\n",
    "        \n",
    "    def __repr__(self) -> str:\n",
    "        return f'{self.name} {self.strategies.keys()}'\n",
    "    \n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    import doctest\n",
    "    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
