{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "b8949d99-6bf5-4abb-a857-28c97f6667b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "import backtrader as bt\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import vectorbt as vbt\n",
    "\n",
    "df = pd.DataFrame(index=[datetime(2020, 1, i + 1) for i in range(9)])\n",
    "df['open'] = [1, 1, 2, 3, 4, 5, 6, 7, 8]\n",
    "df['high'] = df['open'] + 0.5\n",
    "df['low'] = df['open'] - 0.5\n",
    "df['close'] = df['open']\n",
    "data = bt.feeds.PandasData(dataname=df)\n",
    "size = np.array([5, 5, -5, -5, -5, -5, 5, 5, 0])\n",
    "\n",
    "\n",
    "class CommInfoFloat(bt.CommInfoBase):\n",
    "    \"\"\"Commission schema that keeps size as float.\"\"\"\n",
    "    params = (\n",
    "        ('stocklike', True),\n",
    "        ('commtype', bt.CommInfoBase.COMM_PERC),\n",
    "        ('percabs', True),\n",
    "      )\n",
    "    \n",
    "    def getsize(self, price, cash):\n",
    "        if not self._stocklike:\n",
    "            return self.p.leverage * (cash / self.get_margin(price))\n",
    "\n",
    "        return self.p.leverage * (cash / price)\n",
    "\n",
    "\n",
    "class CashValueAnalyzer(bt.analyzers.Analyzer):\n",
    "    \"\"\"Analyzer to extract cash and value.\"\"\"\n",
    "    def create_analysis(self):\n",
    "        self.rets = {}\n",
    "\n",
    "    def notify_cashvalue(self, cash, value):\n",
    "        self.rets[self.strategy.datetime.datetime()] = (cash, value)\n",
    "\n",
    "    def get_analysis(self):\n",
    "        return self.rets\n",
    "\n",
    "\n",
    "class TestStrategy(bt.Strategy):\n",
    "    def __init__(self):\n",
    "        self.i = 0\n",
    "        \n",
    "    def log(self, txt, dt=None):\n",
    "        dt = dt or self.data.datetime[0]\n",
    "        dt = bt.num2date(dt)\n",
    "        print('%s, %s' % (dt.isoformat(), txt))\n",
    "        \n",
    "    def notify_order(self, order):\n",
    "        if order.status in [bt.Order.Submitted, bt.Order.Accepted]:\n",
    "            return  # Await further notifications\n",
    "\n",
    "        if order.status == order.Completed:\n",
    "            if order.isbuy():\n",
    "                buytxt = 'BUY COMPLETE {}, size = {:.2f}, price = {:.2f}'.format(\n",
    "                    order.data._name, order.executed.size, order.executed.price)\n",
    "                self.log(buytxt, order.executed.dt)\n",
    "            else:\n",
    "                selltxt = 'SELL COMPLETE {}, size = {:.2f}, price = {:.2f}'.format(\n",
    "                    order.data._name, order.executed.size, order.executed.price)\n",
    "                self.log(selltxt, order.executed.dt)\n",
    "\n",
    "        elif order.status in [order.Expired, order.Canceled, order.Margin]:\n",
    "            self.log('%s ,' % order.Status[order.status])\n",
    "            pass  # Simply log\n",
    "\n",
    "        # Allow new orders\n",
    "        self.orderid = None\n",
    "        \n",
    "    def next(self):\n",
    "        if size[self.i] > 0:\n",
    "            self.buy(size=size[self.i])\n",
    "        elif size[self.i] < 0:\n",
    "            self.sell(size=-size[self.i])\n",
    "        self.i += 1\n",
    "\n",
    "def bt_simulate(shortcash):\n",
    "    cerebro = bt.Cerebro()\n",
    "    comminfo = CommInfoFloat(commission=0.01)\n",
    "    cerebro.broker.addcommissioninfo(comminfo)\n",
    "    cerebro.addstrategy(TestStrategy)\n",
    "    cerebro.addanalyzer(CashValueAnalyzer)\n",
    "    cerebro.broker.setcash(100.)\n",
    "    cerebro.broker.set_checksubmit(False)\n",
    "    cerebro.broker.set_shortcash(shortcash)\n",
    "    cerebro.adddata(data)\n",
    "    return cerebro.run()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "f3b0b9e6-c91e-4c1d-a587-b95266bc24b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2020-01-02T00:00:00, BUY COMPLETE , size = 5.00, price = 1.00\n",
      "2020-01-03T00:00:00, BUY COMPLETE , size = 5.00, price = 2.00\n",
      "2020-01-04T00:00:00, SELL COMPLETE , size = -5.00, price = 3.00\n",
      "2020-01-05T00:00:00, SELL COMPLETE , size = -5.00, price = 4.00\n",
      "2020-01-06T00:00:00, SELL COMPLETE , size = -5.00, price = 5.00\n",
      "2020-01-07T00:00:00, SELL COMPLETE , size = -5.00, price = 6.00\n",
      "2020-01-08T00:00:00, BUY COMPLETE , size = 5.00, price = 7.00\n",
      "2020-01-09T00:00:00, BUY COMPLETE , size = 5.00, price = 8.00\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{datetime.datetime(2020, 1, 1, 0, 0): (100.0, 100.0),\n",
       " datetime.datetime(2020, 1, 2, 0, 0): (94.95, 99.95),\n",
       " datetime.datetime(2020, 1, 3, 0, 0): (84.85000000000001, 104.85000000000001),\n",
       " datetime.datetime(2020, 1, 4, 0, 0): (99.7, 114.7),\n",
       " datetime.datetime(2020, 1, 5, 0, 0): (119.5, 119.5),\n",
       " datetime.datetime(2020, 1, 6, 0, 0): (144.25, 119.25),\n",
       " datetime.datetime(2020, 1, 7, 0, 0): (173.95, 113.94999999999999),\n",
       " datetime.datetime(2020, 1, 8, 0, 0): (138.6, 103.6),\n",
       " datetime.datetime(2020, 1, 9, 0, 0): (98.19999999999999, 98.19999999999999)}"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "strategy = bt_simulate(True)\n",
    "strategy.analyzers.cashvalueanalyzer.get_analysis()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "e91645b8-9efe-468a-9154-d0928c9518e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2020-01-01    100.00\n",
      "2020-01-02     94.95\n",
      "2020-01-03     84.85\n",
      "2020-01-04     99.70\n",
      "2020-01-05    119.50\n",
      "2020-01-06    144.25\n",
      "2020-01-07    173.95\n",
      "2020-01-08    138.60\n",
      "2020-01-09     98.20\n",
      "Name: close, dtype: float64\n",
      "2020-01-01    100.00\n",
      "2020-01-02     99.95\n",
      "2020-01-03    104.85\n",
      "2020-01-04    114.70\n",
      "2020-01-05    119.50\n",
      "2020-01-06    119.25\n",
      "2020-01-07    113.95\n",
      "2020-01-08    103.60\n",
      "2020-01-09     98.20\n",
      "Name: close, dtype: float64\n"
     ]
    }
   ],
   "source": [
    "portfolio = vbt.Portfolio.from_orders(df.close, [np.nan] + size[:-1].tolist(), fees=0.01)\n",
    "print(portfolio.cash(free=False))\n",
    "print(portfolio.value())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "bf4257a3-511a-405a-a1fd-9d7250935f5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2020-01-02T00:00:00, BUY COMPLETE , size = 5.00, price = 1.00\n",
      "2020-01-03T00:00:00, BUY COMPLETE , size = 5.00, price = 2.00\n",
      "2020-01-04T00:00:00, SELL COMPLETE , size = -5.00, price = 3.00\n",
      "2020-01-05T00:00:00, SELL COMPLETE , size = -5.00, price = 4.00\n",
      "2020-01-06T00:00:00, SELL COMPLETE , size = -5.00, price = 5.00\n",
      "2020-01-07T00:00:00, SELL COMPLETE , size = -5.00, price = 6.00\n",
      "2020-01-08T00:00:00, BUY COMPLETE , size = 5.00, price = 7.00\n",
      "2020-01-09T00:00:00, BUY COMPLETE , size = 5.00, price = 8.00\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{datetime.datetime(2020, 1, 1, 0, 0): (100.0, 100.0),\n",
       " datetime.datetime(2020, 1, 2, 0, 0): (94.95, 99.95),\n",
       " datetime.datetime(2020, 1, 3, 0, 0): (84.85000000000001, 104.85000000000001),\n",
       " datetime.datetime(2020, 1, 4, 0, 0): (99.7, 114.7),\n",
       " datetime.datetime(2020, 1, 5, 0, 0): (119.5, 119.5),\n",
       " datetime.datetime(2020, 1, 6, 0, 0): (94.25, 119.25),\n",
       " datetime.datetime(2020, 1, 7, 0, 0): (63.95, 113.95),\n",
       " datetime.datetime(2020, 1, 8, 0, 0): (83.60000000000001, 103.60000000000001),\n",
       " datetime.datetime(2020, 1, 9, 0, 0): (98.2, 98.2)}"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "strategy = bt_simulate(False)\n",
    "strategy.analyzers.cashvalueanalyzer.get_analysis()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "e936ecb7-4878-4d9f-97ba-675292e95a47",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2020-01-01    100.00\n",
      "2020-01-02     94.95\n",
      "2020-01-03     84.85\n",
      "2020-01-04     99.70\n",
      "2020-01-05    119.50\n",
      "2020-01-06     94.25\n",
      "2020-01-07     63.95\n",
      "2020-01-08     83.60\n",
      "2020-01-09     98.20\n",
      "Name: close, dtype: float64\n",
      "2020-01-01    100.00\n",
      "2020-01-02     99.95\n",
      "2020-01-03    104.85\n",
      "2020-01-04    114.70\n",
      "2020-01-05    119.50\n",
      "2020-01-06    119.25\n",
      "2020-01-07    113.95\n",
      "2020-01-08    103.60\n",
      "2020-01-09     98.20\n",
      "Name: close, dtype: float64\n"
     ]
    }
   ],
   "source": [
    "print(portfolio.cash(free=True))\n",
    "print(portfolio.value())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6690767c-70f6-4ff2-859e-979b3bd83f2f",
   "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": 5
}