{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) 2021, salesforce.com, inc.  \n",
    "All rights reserved.  \n",
    "SPDX-License-Identifier: BSD-3-Clause  \n",
    "For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# This notebook will be used to gather real-world data and perform data processing in order to use it in the covid-19 simulation.\n",
    "\n",
    "### All the downloaded data will be formatted into pandas dataframes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Below is the list of COVID-19 data sources used in this notebook\n",
    "\n",
    "1. **US state government policies** (Oxford Covid-19 Government Response Tracker (OxCGRT))\n",
    "\n",
    "    https://github.com/OxCGRT/USA-covid-policy\n",
    "\n",
    "\n",
    "2. **US federal government direct payments** (Committee for a Responsible Federal Budget)\n",
    "\n",
    "    https://www.covidmoneytracker.org/\n",
    "    \n",
    "    https://docs.google.com/spreadsheets/d/1Nr_J5wLfUT4IzqSXkYbdOXrRgEkBxhX0/edit#gid=682404301\n",
    "    \n",
    "\n",
    "3. **US deaths data** (COVID-19 Data Repository by the Center for Systems Science and Engineering (CSSE) at Johns Hopkins University)\n",
    "\n",
    "    https://github.com/CSSEGISandData/COVID-19\n",
    "\n",
    "\n",
    "4. **US vaccinations** (Our World in Data)\n",
    "    \n",
    "    https://ourworldindata.org/covid-vaccinations\n",
    "    \n",
    "    \n",
    "5. **US unemployment** (Bureau of Labor and Statistics)\n",
    "\n",
    "    https://www.bls.gov/lau/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime, timedelta\n",
    "import json\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import pickle\n",
    "import scipy\n",
    "from scipy.signal import convolve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Classes to fetch the real-world data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ai_economist.datasets.covid19_datasets.us_policies import DatasetCovidPoliciesUS\n",
    "from ai_economist.datasets.covid19_datasets.us_deaths import DatasetCovidDeathsUS\n",
    "from ai_economist.datasets.covid19_datasets.us_vaccinations import DatasetCovidVaccinationsUS\n",
    "from ai_economist.datasets.covid19_datasets.us_unemployment import DatasetCovidUnemploymentUS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set a base directory where you would like to download real world data. The latest data will be downloaded into a folder within the base directory, named using the current date"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_DATA_DIR_PATH = \"/tmp/covid19_data\"  # SPECIFY A BASE DIRECTORY TO STORE ALL THE DOWNLOADED DATA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DOWNLOAD_LATEST_DATA = True  # Download the latest data or use whatever is saved earlier \n",
    "CURRENT_DATE = datetime.now()\n",
    "DATE_FORMAT = \"%Y-%m-%d\"\n",
    "date_string = CURRENT_DATE.strftime(DATE_FORMAT).replace('/','-')\n",
    "data_dir = os.path.join(BASE_DATA_DIR_PATH, date_string)\n",
    "\n",
    "print(\"All the data will be downloaded to the directory: '{}'.\".format(data_dir))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up dictionary to write model constants\n",
    "model_constants = {}\n",
    "model_constants_filename = \"model_constants.json\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gather real-world data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. COVID-19 US State Government Policies\n",
    "### Source: Oxford Covid-19 Government Response Tracker (OxCGRT) \n",
    "(https://github.com/OxCGRT/USA-covid-policy)\n",
    "\n",
    "**NOTE:** All data will use the same format as **policy_df** (below) and use the same date index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "covid_policies_us = DatasetCovidPoliciesUS(\n",
    "    data_dir=data_dir,\n",
    "    download_latest_data=DOWNLOAD_LATEST_DATA\n",
    ")\n",
    "\n",
    "# Which of the policy indicators to treat as the open/close level\n",
    "STRINGENCY_POLICY_KEY = 'StringencyIndex'\n",
    "# Number of levels to discretize the stringency policy into. \n",
    "# In the context of reinforcement learning, this also determines the action space of the agents.\n",
    "NUM_STRINGENCY_LEVELS = 10\n",
    "\n",
    "policies_us_df = covid_policies_us.process_policy_data(\n",
    "    stringency_policy_key=STRINGENCY_POLICY_KEY,\n",
    "    num_stringency_levels=NUM_STRINGENCY_LEVELS\n",
    ")\n",
    "\n",
    "print(\"Policy data are available between {} and {}\".format(policies_us_df[\"Date\"].min(), \n",
    "                                                           policies_us_df[\"Date\"].max()))\n",
    "\n",
    "policy_df = policies_us_df.pivot(\n",
    "    index=\"Date\", columns=\"RegionName\", values=STRINGENCY_POLICY_KEY\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the common date index that all the dataframes will use\n",
    "COMMON_DATE_INDEX = policy_df.index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the list of states (in order) all the dataframes will use\n",
    "US_STATE_ORDER = policy_df.columns.values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the stringency level for a specified US state\n",
    "state = \"California\"\n",
    "policy_df[state].plot(figsize=(15,5), x='Date', title=\"Stringency Level for {}\".format(state), grid=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. COVID-19 Federal government subsidies (direct payments) to the states\n",
    "### Source: Committee For A Responsible Federal Budget\n",
    "https://www.covidmoneytracker.org/\n",
    "\n",
    "### Direct payments provided by the Federal Government so far are recorded in this google spreadsheet\n",
    "https://docs.google.com/spreadsheets/d/1Nr_J5wLfUT4IzqSXkYbdOXrRgEkBxhX0/edit#gid=682404301\n",
    "### Read as (date: direct payment amount)\n",
    "2020-04-15: 274B\n",
    "\n",
    "2020-12-27: 142B\n",
    "\n",
    "2021-03-11: 386B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subsidy_df = pd.DataFrame(policy_df.index).set_index(\"Date\")\n",
    "subsidy_df[\"USA\"] = 0.0\n",
    "\n",
    "subsidy_df.loc[\"2020-04-15\", \"USA\"] = 274e9\n",
    "subsidy_df.loc[\"2020-12-27\", \"USA\"] = 142e9\n",
    "subsidy_df.loc[\"2021-03-11\", \"USA\"] = 386e9"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. COVID-19 Deaths data\n",
    "### Source: COVID-19 Data Repository by the Center for Systems Science and Engineering (CSSE) at Johns Hopkins University \n",
    "(https://github.com/CSSEGISandData/COVID-19)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deaths_us_df = DatasetCovidDeathsUS(\n",
    "    data_dir=data_dir,\n",
    "    download_latest_data=DOWNLOAD_LATEST_DATA\n",
    ").df\n",
    "\n",
    "print(\"COVID-19 death data for the US is available between {} and {}\".format(\n",
    "    deaths_us_df.columns[12], deaths_us_df.columns[-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retain just the states in US_STATE_ORDER\n",
    "deaths_us_df = deaths_us_df[deaths_us_df.Province_State.isin(US_STATE_ORDER)]\n",
    "\n",
    "# We will visualize this later in the notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. COVID-19 Vaccination Data\n",
    "### Source: Our World in Data\n",
    "(https://ourworldindata.org/covid-vaccinations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vaccinations_us_df = DatasetCovidVaccinationsUS(\n",
    "    data_dir=data_dir,\n",
    "    download_latest_data=DOWNLOAD_LATEST_DATA\n",
    ").df\n",
    "\n",
    "vaccination_dates = sorted(vaccinations_us_df.date.unique())\n",
    "print(\"Vaccination data is available between {} and {}\".format(min(vaccination_dates), max(vaccination_dates)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vaccinated_df = vaccinations_us_df.pivot(\n",
    "    index=\"date\", columns=\"location\", values=\"people_fully_vaccinated\"\n",
    ")[US_STATE_ORDER]\n",
    "\n",
    "vaccinated_df.index = pd.to_datetime(vaccinated_df.index)\n",
    "vaccinated_df = vaccinated_df.reindex(COMMON_DATE_INDEX).fillna(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the vaccinations for a specified US state\n",
    "# Warning: the last value may not be updated (may show it to be 0)\n",
    "\n",
    "state = \"California\"\n",
    "vaccinated_df[state].plot(figsize=(15,5), x='Date', title=\"Vaccinations for {}\".format(state), grid=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using deaths and vaccinations to compute the susceptible-infected-recovered (SIR) numbers\n",
    "\n",
    "Our SIR data will only treat **deaths** as ground-truth.\n",
    "\n",
    "Given death data and some assumed constants about the _death rate_ and _recovery rate_ , we can apply some \"SIR algebra\" (i.e. solve for unknowns using the SIR equations) to _infer_ quantities like total \"recovered\", number of infected people, and ultimately **Beta**, which is the rate of transmission times the number of people an infected person comes into contact with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For data representation, we will want to build a dataframe for...\n",
    "# ... deaths...\n",
    "deaths_df = pd.DataFrame(COMMON_DATE_INDEX, columns=['Date']).set_index('Date')\n",
    "smoothed_deaths_df = pd.DataFrame(COMMON_DATE_INDEX, columns=['Date']).set_index('Date')\n",
    "# ... (inferred) SIR states...\n",
    "susceptible_df = pd.DataFrame(COMMON_DATE_INDEX, columns=['Date']).set_index('Date')\n",
    "infected_df    = pd.DataFrame(COMMON_DATE_INDEX, columns=['Date']).set_index('Date')\n",
    "recovered_df   = pd.DataFrame(COMMON_DATE_INDEX, columns=['Date']).set_index('Date')\n",
    "# ... and (inferred) Beta.\n",
    "beta_df = pd.DataFrame(COMMON_DATE_INDEX, columns=['Date']).set_index('Date')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# STD of the Gaussian smoothing window applied to the death data.\n",
    "SIR_SMOOTHING_STD = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fill the death dataframe from (smoothed) raw data\n",
    "\n",
    "def smooth(x, gauss_std=10):\n",
    "    \"\"\"\n",
    "    gauss_std: standard deviation of the Gaussian smoothing window applied to the death data.\n",
    "    \"\"\"\n",
    "    if gauss_std <= 0:\n",
    "        return x\n",
    "    # To invalidate the near-edge results, bookend the input x with nans\n",
    "    x = np.concatenate([[np.nan], np.array(x), [np.nan]])\n",
    "    \n",
    "    kernel = scipy.stats.norm.pdf(\n",
    "        np.linspace(-3*gauss_std, 3*gauss_std, 1+6*gauss_std),\n",
    "        scale=gauss_std\n",
    "    )\n",
    "    normer = np.ones_like(x)\n",
    "    smoothed_x = convolve(x, kernel, mode='same') / convolve(normer, kernel, mode='same')\n",
    "    \n",
    "    # Remove the indices added by the nan padding\n",
    "    return smoothed_x[1:-1]\n",
    "\n",
    "for us_state_name in US_STATE_ORDER:\n",
    "    state_deaths = deaths_us_df[deaths_us_df['Province_State']==us_state_name]\n",
    "    cumulative_state_deaths = []\n",
    "    for d in COMMON_DATE_INDEX:\n",
    "        date_string = '{d.month}/{d.day}/{y}'.format(d=d, y=d.year % 2000)\n",
    "        if date_string in state_deaths:\n",
    "            cumulative_state_deaths.append(\n",
    "                state_deaths[date_string].sum()\n",
    "            )\n",
    "        else:\n",
    "            cumulative_state_deaths.append(\n",
    "                np.nan\n",
    "            )\n",
    "            \n",
    "    # Store raw numbers (for direct comparison)\n",
    "    deaths_df[us_state_name] = cumulative_state_deaths\n",
    "    \n",
    "    # Store smoothed numbers (for beta analysis)\n",
    "    smoothed_cumulative_state_deaths = smooth(cumulative_state_deaths, gauss_std=SIR_SMOOTHING_STD)\n",
    "    smoothed_deaths_df[us_state_name] = smoothed_cumulative_state_deaths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_deaths = deaths_us_df[deaths_us_df['Province_State']==\"California\"]\n",
    "cumulative_state_deaths = []\n",
    "for d in COMMON_DATE_INDEX:\n",
    "    date_string = '{d.month}/{d.day}/{y}'.format(d=d, y=d.year % 2000)\n",
    "    if date_string in state_deaths:\n",
    "        cumulative_state_deaths.append(\n",
    "            state_deaths[date_string].sum()\n",
    "        )\n",
    "    else:\n",
    "        cumulative_state_deaths.append(\n",
    "            np.nan\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the deaths for a specified US state\n",
    "state = \"California\"\n",
    "\n",
    "# Some values near the ends may be \"missing\" because of smoothing\n",
    "deaths_df[state].plot(figsize=(15,5), x='Date', ylim=[0, 65000]);\n",
    "smoothed_deaths_df[state].plot(figsize=(15,5), x='Date', title=\"COVID deaths in {}\".format(state), ylim=[0, 65000], grid=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Death rate: fraction of infected persons who die\n",
    "SIR_MORTALITY = 0.02\n",
    "\n",
    "# Recovery rate: the inverse of expected time someone remains infected\n",
    "SIR_GAMMA = 1 / 14"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the core \"SIR algebra\" used to infer S, I, R, and Beta at each date.\n",
    "\n",
    "def infer_sir_and_beta(us_state_name):\n",
    "    state_population = deaths_us_df[deaths_us_df['Province_State']==us_state_name]['Population'].sum()\n",
    "    \n",
    "    # Helpful to do this math in normalized numbers\n",
    "    dead = np.array(smoothed_deaths_df[us_state_name]) / state_population\n",
    "    vaccinated = np.array(vaccinated_df[us_state_name]) / state_population\n",
    "    \n",
    "    # Dead is the fraction of \"recovered\" that did not survive\n",
    "    # Also, the vaccinated lot is part of the recovered\n",
    "    recovered = dead / SIR_MORTALITY + vaccinated\n",
    "    \n",
    "    # The daily change in recovered (ignoring the vaccinated) is a fraction of the infected population on the previous day\n",
    "    infected = np.nan * np.zeros_like(dead)\n",
    "    infected[:-1] = (recovered[1:] - recovered[:-1] - (vaccinated[1:] - vaccinated[:-1])) / SIR_GAMMA\n",
    "    \n",
    "    # S+I+R must always = 1\n",
    "    susceptible = 1 - infected - recovered\n",
    "    \n",
    "    # Here's where things get interesting. The change in infected is due to...\n",
    "    change_in_i = infected[1:] - infected[:-1]\n",
    "    # ... infected people that transition to the recovered state (decreases I)...\n",
    "    expected_change_from_recovery = -infected[:-1] * SIR_GAMMA\n",
    "    # ... and susceptible people that transition to the infected state (increases I).\n",
    "    new_infections = change_in_i - expected_change_from_recovery\n",
    "    \n",
    "    # With these pieces, we can solve for Beta.\n",
    "    beta_ = new_infections / (infected[:-1] * susceptible[:-1] + 1e-6)\n",
    "    beta_ = np.clip(beta_, 0, 1)\n",
    "    # Apply a threshold in terms of normalized daily deaths (if too low, beta estimates are bad)\n",
    "    normalized_daily_deaths = dead[1:]-dead[:-1]\n",
    "    ndd_lookback = np.zeros_like(new_infections)\n",
    "    lookback_window = 3*SIR_SMOOTHING_STD\n",
    "    ndd_cutoff = 1e-8\n",
    "    ndd_lookback[lookback_window:] = normalized_daily_deaths[:-lookback_window]\n",
    "    beta_[np.logical_not(ndd_lookback > 1e-8)] = np.nan\n",
    "    \n",
    "    beta = np.nan * np.zeros_like(dead)\n",
    "    beta[:-1] = beta_\n",
    "    \n",
    "    # Undo normalization\n",
    "    susceptible *= state_population\n",
    "    infected    *= state_population\n",
    "    recovered   *= state_population\n",
    "    \n",
    "    return susceptible, infected, recovered, beta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fill the SIR and Beta dataframes with their inferred values\n",
    "for st in US_STATE_ORDER:\n",
    "    susceptible_df[st], infected_df[st], recovered_df[st], beta_df[st] = infer_sir_and_beta(us_state_name=st)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Visualize the SIR and BETA for a specified US state\n",
    "# Warning: some values near the ends may be \"missing\" because of smoothing\n",
    "\n",
    "state = \"California\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "susceptible_df[state].plot(figsize=(15,3), x='Date', title=\"(Inferred) Susceptible Population in {}\".format(state), grid=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "infected_df[state].plot(figsize=(15,3), x='Date', title=\"(Inferred) Infected Population in {}\".format(state), grid=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recovered_df[state].plot(figsize=(15,3), x='Date', title=\"(Inferred) Recovered Population in {}\".format(state), grid=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_df[state].plot(figsize=(15,3), x='Date', title=\"(Inferred) SIR Beta in {}\".format(state), grid=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. COVID-19 Unemployment data\n",
    "### Source: Bureau of Labor and Statistics\n",
    "\n",
    "https://www.bls.gov/lau/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "monthly_unemployment_us = DatasetCovidUnemploymentUS(\n",
    "    data_dir=data_dir,\n",
    "    download_latest_data=DOWNLOAD_LATEST_DATA).data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_monthly_unemployment = monthly_unemployment_us['California']\n",
    "unemp_year_keys = sorted(sample_monthly_unemployment.keys())\n",
    "unemp_starting_month_key = sorted(sample_monthly_unemployment[unemp_year_keys[0]].keys())[0]\n",
    "unemp_ending_month_key = sorted(sample_monthly_unemployment[unemp_year_keys[-1]].keys())[-1]\n",
    "unemp_starting_date = datetime.strptime(\n",
    "    str(unemp_year_keys[0]) + '-' + str(unemp_ending_month_key+1) + '-1', DATE_FORMAT)\n",
    "unemp_ending_date = datetime.strptime(\n",
    "    str(unemp_year_keys[-1]) + '-' + str(unemp_ending_month_key+1) + '-1', DATE_FORMAT) - timedelta(1)\n",
    "\n",
    "print(\"Unemployment data is available between {} and {}\".format(datetime.strftime(unemp_starting_date, DATE_FORMAT),\n",
    "                                                                datetime.strftime(unemp_ending_date, DATE_FORMAT)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert this to a daily unemployment dataframe\n",
    "\n",
    "unemployment_df = pd.DataFrame(COMMON_DATE_INDEX, columns=['Date']).set_index('Date')\n",
    "\n",
    "for us_state_name in monthly_unemployment_us.keys():\n",
    "    unemployment_df[us_state_name] = [\n",
    "        monthly_unemployment_us[us_state_name][x.year].get(x.month, np.nan)\n",
    "        for x in unemployment_df.index\n",
    "    ]\n",
    "unemployment_df = unemployment_df[US_STATE_ORDER]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Visualize the unemployment rate for a specified US state\n",
    "# There is likely going to be some unemployment data missing at the tail end, \n",
    "# as the unemployment data isn't updated as frequently as the other data.\n",
    "\n",
    "state = \"California\"\n",
    "unemployment_df[state].plot(figsize=(15,5), x='Date', title=\"Unemployment for {} (%)\".format(state), grid=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Unemployment rate -> unemployed (the number of unemployed people)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "us_state_to_pop_dict = {}\n",
    "for us_state in US_STATE_ORDER:\n",
    "    us_state_to_pop_dict[us_state] = deaths_us_df[deaths_us_df.Province_State==us_state].Population.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unemployed_df = unemployment_df.multiply([us_state_to_pop_dict[col]/100.0 for col in unemployment_df.columns])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save some of the data processing constants for use within the environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_constants_dict = {}\n",
    "\n",
    "model_constants_dict[\"DATE_FORMAT\"] = DATE_FORMAT\n",
    "model_constants_dict[\"STRINGENCY_POLICY_KEY\"] = STRINGENCY_POLICY_KEY\n",
    "model_constants_dict[\"NUM_STRINGENCY_LEVELS\"] = int(NUM_STRINGENCY_LEVELS)\n",
    "model_constants_dict[\"SIR_SMOOTHING_STD\"] = SIR_SMOOTHING_STD\n",
    "model_constants_dict[\"SIR_MORTALITY\"] = SIR_MORTALITY\n",
    "model_constants_dict[\"SIR_GAMMA\"] = SIR_GAMMA\n",
    "model_constants_dict[\"US_STATE_IDX_TO_STATE_NAME\"] = {\n",
    "    us_state_idx: us_state for us_state_idx, us_state in enumerate(US_STATE_ORDER)\n",
    "}\n",
    "model_constants_dict[\"US_STATE_POPULATION\"] = [int(us_state_to_pop_dict[us_state]) for us_state in US_STATE_ORDER]\n",
    "model_constants_dict[\"US_POPULATION\"] = int(sum([us_state_to_pop_dict[us_state] for us_state in US_STATE_ORDER]))\n",
    "\n",
    "# 2019: https://data.worldbank.org/indicator/NY.GDP.PCAP.CD?locations=US&view=chart\n",
    "model_constants_dict[\"GDP_PER_CAPITA\"] = 65300  # TODO: Load this in from model_constants.json.\n",
    "\n",
    "model_constants_filename = \"model_constants.json\"\n",
    "with open(os.path.join(data_dir, model_constants_filename), \"w\") as fp: \n",
    "    json.dump(model_constants_dict, fp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save all the processed dataframes in order to use for model fitting notebook (fit_model_parameters.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataframes = {\n",
    "    \"policy\": policy_df,\n",
    "    \"subsidy\": subsidy_df,\n",
    "    \"deaths\": deaths_df,\n",
    "    \"vaccinated\": vaccinated_df,\n",
    "    \"smoothed_deaths\": smoothed_deaths_df,\n",
    "    \"susceptible\": susceptible_df,\n",
    "    \"infected\": infected_df,\n",
    "    \"recovered\": recovered_df,\n",
    "    \"beta\": beta_df,\n",
    "    \"unemployment\": unemployment_df,\n",
    "    \"unemployed\": unemployed_df,\n",
    "}\n",
    "\n",
    "for k, df in dataframes.items():\n",
    "    if k == \"subsidy\":  # This is at the USA level, not at the US states level\n",
    "        continue\n",
    "    assert (df.columns.to_list() == US_STATE_ORDER).all()\n",
    "\n",
    "with open(os.path.join(data_dir, 'dataframes.pkl'), 'wb') as F:\n",
    "    pickle.dump(dataframes, F)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Also save all the data as numpy arrays for use within the covid19 simulation environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "real_world_data = {}\n",
    "for key in dataframes:\n",
    "    real_world_data[key] = dataframes[key].values\n",
    "    \n",
    "# Save the real-world data as a .npz for use within the environment\n",
    "np.savez(os.path.join(data_dir, \"real_world_data.npz\"), **real_world_data)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finally, in order to use this gathered real-world data when you run the covid19 simulation, you will need to also"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Run the \"fit_model_parameters.ipynb\" notebook with the base data directory specified below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"BASE_DATA_DIR_PATH = '{}'\".format(BASE_DATA_DIR_PATH))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Set \"path_to_data_and_fitted_params\" in the env config also to the data directory below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"path_to_data_and_fitted_params = '{}'\".format(data_dir))"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
