{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-09T02:27:13.715035Z",
     "start_time": "2022-11-09T02:27:13.113284Z"
    }
   },
   "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"
     ]
    }
   ],
   "source": [
    "%%checkall\n",
    "import re\n",
    "import datetime\n",
    "from dateutil._common import weekday\n",
    "import dateutil.relativedelta as rd\n",
    "import numpy as np\n",
    "from pyqstrat.holiday_calendars import Calendar, get_date_from_weekday\n",
    "from pyqstrat.pq_utils import assert_\n",
    "\n",
    "FUTURE_CODES_INT = {'F': 1, 'G': 2, 'H': 3, 'J': 4, 'K': 5, 'M': 6, 'N': 7, 'Q': 8, 'U': 9, 'V': 10, 'X': 11, 'Z': 12}\n",
    "FUTURES_CODES_INVERTED: dict[int, str] = {v: k for k, v in FUTURE_CODES_INT.items()}\n",
    "\n",
    "FUTURE_CODES_STR = {'F': 'jan', 'G': 'feb', 'H': 'mar', 'J': 'apr', 'K': 'may', 'M': 'jun', \n",
    "                    'N': 'jul', 'Q': 'aug', 'U': 'sep', 'V': 'oct', 'X': 'nov', 'Z': 'dec'}\n",
    "\n",
    "\n",
    "def future_code_to_month(future_code: str) -> str:\n",
    "    '''\n",
    "    Given a future code such as \"X\", return the month abbreviation, such as \"nov\"\n",
    "    \n",
    "    Args:\n",
    "        future_code (str): the one letter future code\n",
    "        \n",
    "    >>> future_code_to_month('X')\n",
    "    'nov'\n",
    "    '''\n",
    "    assert_(len(future_code) == 1, f'Future code must be a single character: {future_code}')\n",
    "    if future_code not in FUTURE_CODES_STR: raise Exception(f'unknown future code: {future_code}')\n",
    "    return FUTURE_CODES_STR[future_code]\n",
    "    \n",
    "    if future_code not in FUTURE_CODES_INT: raise Exception(f'unknown future code: {future_code}')\n",
    "    return FUTURE_CODES_INT[future_code]\n",
    "\n",
    "\n",
    "def future_code_to_month_number(future_code: str) -> int:\n",
    "    '''\n",
    "    Given a future code such as \"X\", return the month number (from 1 - 12)\n",
    "    \n",
    "    Args:\n",
    "        future_code (str): the one letter future code\n",
    "        \n",
    "    >>> future_code_to_month_number('X')\n",
    "    11\n",
    "    '''\n",
    "    assert_(len(future_code) == 1, f'Future code must be a single character: {future_code}')\n",
    "    if future_code not in FUTURE_CODES_INT: raise Exception(f'unknown future code: {future_code}')\n",
    "    return FUTURE_CODES_INT[future_code]\n",
    "\n",
    "\n",
    "def get_future_code(month: int) -> str:\n",
    "    '''\n",
    "    Given a month number such as 3 for March, return the future code for it, e.g. H\n",
    "    >>> get_future_code(3)\n",
    "    'H'\n",
    "    '''\n",
    "    return FUTURES_CODES_INVERTED[month]\n",
    "\n",
    "\n",
    "class EminiFuture:\n",
    "    \n",
    "    calendar = Calendar.get_calendar(Calendar.NYSE)\n",
    "\n",
    "    @staticmethod\n",
    "    def get_current_symbol(curr_date: datetime.date) -> str:\n",
    "        '''\n",
    "        >>> assert(EminiFuture.get_current_symbol(datetime.date(2019, 3, 14)) == 'ESH9')\n",
    "        >>> assert(EminiFuture.get_current_symbol(datetime.date(2019, 3, 15)) == 'ESM9')\n",
    "        >>> assert(EminiFuture.get_current_symbol(datetime.date(2020, 3, 14)) == 'ESH0')\n",
    "        '''\n",
    "        year = curr_date.year\n",
    "        month = curr_date.month\n",
    "        day = curr_date.day\n",
    "        third_friday_ = EminiFuture.calendar.third_friday_of_month(month, year)\n",
    "        third_friday: datetime.date = third_friday_.astype(datetime.date)  # type: ignore\n",
    "        if month < 3 or (month == 3 and day < third_friday.day): month_str = 'H'\n",
    "        elif month < 6 or (month == 6 and day < third_friday.day): month_str = 'M'\n",
    "        elif month < 9 or (month == 9 and day < third_friday.day): month_str = 'U'\n",
    "        elif month < 12 or (month == 12 and day < third_friday.day): month_str = 'Z'\n",
    "        else:\n",
    "            month_str = 'H'\n",
    "            year += 1\n",
    "        base = 2010 if year < 2020 else 2020\n",
    "        fut_symbol = 'ES' + month_str + str(year - base)\n",
    "        return fut_symbol\n",
    "\n",
    "    @staticmethod\n",
    "    def get_previous_symbol(curr_future_symbol: str) -> str:\n",
    "        '''\n",
    "        >>> assert(EminiFuture.get_previous_symbol('ESH9') == 'ESZ8')\n",
    "        '''\n",
    "        month = curr_future_symbol[2]\n",
    "        year = int(curr_future_symbol[3])\n",
    "        prev_month = {'H': 'Z', 'M': 'H', 'U': 'M', 'Z': 'U'}[month]\n",
    "        prev_year = year if prev_month != 'Z' else year - 1\n",
    "        if prev_year == -1: prev_year == 9\n",
    "        return f'ES{prev_month}{prev_year}'\n",
    "\n",
    "    @staticmethod\n",
    "    def get_next_symbol(curr_future_symbol: str) -> str:\n",
    "        '''\n",
    "        >>> assert(EminiFuture.get_next_symbol('ESZ8') == 'ESH9')\n",
    "        '''\n",
    "        month = curr_future_symbol[2]\n",
    "        year = int(curr_future_symbol[3])\n",
    "        next_month = {'Z': 'H', 'H': 'M', 'M': 'U', 'U': 'Z'}[month]\n",
    "        next_year = year if next_month != 'H' else year + 1\n",
    "        if next_year == 10: next_year = 0\n",
    "        return f'ES{next_month}{next_year}'\n",
    "    \n",
    "    @staticmethod\n",
    "    def get_expiry(fut_symbol: str) -> np.datetime64:\n",
    "        '''\n",
    "        >>> assert(EminiFuture.get_expiry('ESH8') == np.datetime64('2018-03-16T08:30'))\n",
    "        '''\n",
    "        month_str = fut_symbol[-2: -1]\n",
    "        year_str = fut_symbol[-1:]\n",
    "        \n",
    "        month = future_code_to_month_number(month_str)\n",
    "        assert_(isinstance(month, int))\n",
    "        year = int(year_str)\n",
    "        year_base = 2020 if year < 5 else 2010\n",
    "        year = year_base + int(year_str)\n",
    "        expiry_date = EminiFuture.calendar.third_friday_of_month(month, year).astype(datetime.date)\n",
    "        return np.datetime64(expiry_date) + np.timedelta64(8 * 60 + 30, 'm')\n",
    "    \n",
    "\n",
    "class EminiOption:\n",
    "    \n",
    "    calendar = Calendar.get_calendar(Calendar.NYSE)\n",
    "\n",
    "    @staticmethod\n",
    "    def decode_symbol(name: str) -> tuple[weekday, int, int, int]:\n",
    "        '''\n",
    "        >>> EminiOption.decode_symbol('E1AF8')\n",
    "        (MO, 2018, 1, 1)\n",
    "        '''\n",
    "        if re.match('EW[1-4].[0-9]', name):  # Friday\n",
    "            year = int('201' + name[-1:])\n",
    "            if year in [2010, 2011]: year += 10\n",
    "            week = int(name[2:3])\n",
    "            month = future_code_to_month_number(name[3:4])\n",
    "            return rd.FR, year, month, week\n",
    "        if re.match('E[1-5]A.[0-9]', name):  # Monday\n",
    "            year = int('201' + name[-1:])\n",
    "            if year in [2010, 2011]: year += 10\n",
    "            week = int(name[1:2])\n",
    "            month = future_code_to_month_number(name[3:4])\n",
    "            return rd.MO, year, month, week\n",
    "        if re.match('E[1-5]C.[0-9]', name):  # Wednesday\n",
    "            year = int('201' + name[-1:])\n",
    "            if year in [2010, 2011]: year += 10\n",
    "            week = int(name[1:2])\n",
    "            month = future_code_to_month_number(name[3:4])\n",
    "            return rd.WE, year, month, week\n",
    "        if re.match('EW[A-Z][0-9]', name):  # End of month\n",
    "            year = int('201' + name[-1:])\n",
    "            if year in [2010, 2011]: year += 10\n",
    "            week = -1\n",
    "            month = future_code_to_month_number(name[2:3])\n",
    "            return rd.WE, year, month, week\n",
    "        else:\n",
    "            raise Exception(f'could not decode: {name}')\n",
    "            \n",
    "    @staticmethod\n",
    "    def get_expiry(symbol: str) -> np.datetime64:\n",
    "        '''\n",
    "        >>> EminiOption.get_expiry('EW2Z5')\n",
    "        numpy.datetime64('2015-12-11T15:00')\n",
    "        >>> EminiOption.get_expiry('E3AF7')\n",
    "        numpy.datetime64('2017-01-17T15:00')\n",
    "        >>> EminiOption.get_expiry('EWF0')\n",
    "        numpy.datetime64('2020-01-31T15:00')\n",
    "        '''\n",
    "        assert_(':' not in symbol, f'{symbol} contains: pass in option root instead')\n",
    "        weekday, year, month, week = EminiOption.decode_symbol(symbol)\n",
    "        expiry_ = get_date_from_weekday(weekday.weekday, year, month, week)\n",
    "        if weekday in [rd.WE, rd.FR]:\n",
    "            expiry = EminiOption.calendar.add_trading_days(expiry_, num_days=0, roll='backward')\n",
    "        else:\n",
    "            expiry = EminiOption.calendar.add_trading_days(expiry_, num_days=0, roll='forward')\n",
    "        # Option expirations changed on 9/20/2015 from 3:15 to 3 pm - \n",
    "        # See https://www.cmegroup.com/market-regulation/files/15-384.pdf\n",
    "        expiry += np.where(expiry < np.datetime64('2015-09-20'), np.timedelta64(15 * 60 + 15, 'm'), np.timedelta64(15, 'h'))\n",
    "        assert isinstance(expiry, np.datetime64)\n",
    "        return expiry\n",
    "       \n",
    "\n",
    "if __name__ == \"__main__\":\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
}
