{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2442bd44-853a-4b07-81aa-bf06b917709b",
   "metadata": {},
   "source": [
    "# Synthetic data generation script for BY-COVID WP5 baseline use case\n",
    "Data model specification is available at https://docs.google.com/spreadsheets/d/17561rLcEC1SeV94JUXHbcFAef2EmrsOQCB_6q0MjIM0/edit#gid=0\n",
    "\n",
    "> **_NOTE 1:_** Please, be aware that all main parameters has been stablished based on known COVID case distribution and COVID vaccine coverage for [Aragon, Spain](https://transparencia.aragon.es/COVID19) \n",
    "\n",
    "> **_NOTE 2:_** Each variable simulated is accompained by figures checking their distribution\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e915c393-69e7-475e-b61c-8dc2ac5c1d72",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from faker import Faker\n",
    "from datetime import datetime, timedelta\n",
    "from dateutil.relativedelta import relativedelta\n",
    "import pandas as pd\n",
    "from datetime import date\n",
    "from scipy.interpolate import InterpolatedUnivariateSpline\n",
    "\n",
    "fake = Faker()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "61648fd1-74ce-40bd-8e48-1e4a5f7d43f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# General population information\n",
    "\n",
    "pop_total = 10000 # total population to be simulated\n",
    "perc_vaccinated = 0.9 # percentage of population fully vaccinated (vaccination schedule completed)\n",
    "pop_vaccinated = int(pop_total * perc_vaccinated) # total vaccinated population\n",
    "pop_novaccinated = pop_total-pop_vaccinated # total unvaccinated population\n",
    "\n",
    "#Modified:  Percentage of population with COVID-19 (Sciensano) = 0.37\n",
    "perc_pop_with_covid = 0.37 # acummulated percentage of the population with covid during the period of study\n",
    "#Modified:  Percentage of population with reinfection of COVID-19 (Sciensano) = 0.026\n",
    "previous_infection_perc = 0.026 # accumulated percentage of the population with a reinfection of covid\n",
    "\n",
    "pop_with_covid = int(perc_pop_with_covid * pop_total)\n",
    "#Modified:  Percentage of population that died = 0.01210076\n",
    "perc_exitus = 0.025363077  # accumulated percentage of deaths (due to all causes) in confirmed cases \n",
    "# It might be the case that you want to input the accummulated percentage of deaths in general population within the period instead\n",
    "pop_exitus  = int(pop_with_covid * perc_exitus) # deceased population\n",
    "\n",
    "#Modified: Percentage of cases with reinfection of COVID-19 = 0.070\n",
    "perc_previous = 0.070 # accumulated percentage of the covid cases with a reinfection of covid\n",
    "pop_with_previous = int(perc_previous*pop_with_covid) # Population with a previous confirmed case of covid\n",
    "\n",
    "\n",
    "# Relevant dates\n",
    "covid_starting_date = datetime.strptime('2020-02-28', '%Y-%m-%d')\n",
    "vaccination_starting_date = datetime.strptime('2020-12-28' , '%Y-%m-%d')\n",
    "end_date = datetime.strptime('2022-05-24', '%Y-%m-%d')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "deae6a9e-c22a-44ce-a166-008193adac0d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 7640 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 7436 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 5916 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# age_nm\n",
    "# Limits age_nm >= 5 | age_nm <= 115\n",
    "# Behave: normal distribution\n",
    "\n",
    "mu = 55\n",
    "sigma = 15\n",
    "age_nm = np.random.normal(mu, sigma, pop_total)\n",
    "age_nm = [5 if x<5 else round(x) for x in age_nm] #correct out-of-limit values\n",
    "age_nm = [115 if x>115 else round(x) for x in age_nm] #correct out-of-limit values\n",
    "plt.hist(age_nm, bins='auto')\n",
    "plt.show()\n",
    "\n",
    "\n",
    "# sex_cd\n",
    "# Options: 0 (not known), 1 (male), 2 (female), 9 (not applicable)\n",
    "# Behave: prob(0) = 0.02, prob(1) = 0.49, prob(2) = 0.49, prob(9) = 0\n",
    "sex_cd = np.random.choice(np.arange(0, 4),pop_total, p=[0.02,0.49,0.49,0])\n",
    "plt.hist(sex_cd, bins='auto')\n",
    "plt.show()\n",
    "\n",
    "\n",
    "# socecon_lvl_cd\n",
    "# Socioeconomic categories are based (in this case) on income quintiles\n",
    "# Limits socecon_lvl_cd >= 1 | socecon_lvl_cd <= 5\n",
    "# Behave: normal distribution\n",
    "socecon_lvl_cd = []  \n",
    "mu = 4\n",
    "sigma = 1\n",
    "    \n",
    "socecon_lvl_cd = np.random.normal(mu, sigma, pop_total)\n",
    "socecon_lvl_cd = [2 if x<1 else x for x in socecon_lvl_cd]\n",
    "socecon_lvl_cd = [3 if x>6 else x for x in socecon_lvl_cd]\n",
    "socecon_lvl_cd = [int(x) for x in socecon_lvl_cd]\n",
    "plt.hist(socecon_lvl_cd, bins='auto')\n",
    "plt.show()\n",
    "\n",
    "\n",
    "# residence_area_cd\n",
    "# NUTS level 3 values for Huesca, Teruel and Zaragoza (731 , 732, 733) with probabilities \n",
    "residence_area_cd = np.random.choice([731 , 732, 733,None],pop_total, p=[0.14, 0.14,0.70,0.02])\n",
    "\n",
    "\n",
    "# country_cd\n",
    "# Spain (ESP)\n",
    "country_cd = np.random.choice([\"ESP\",None],pop_total, p=[0.97,0.03])\n",
    "# ESP\n",
    "\n",
    "\n",
    "# foreign_bl\n",
    "# Options: TRUE, FALSE\n",
    "# prob(TRUE) = 0.15, prob(FALSE) = 0.85\n",
    "foreign_bl = np.random.choice([False,True],pop_total, p=[0.85, 0.15])\n",
    "\n",
    "\n",
    "# essential_worker_bl\n",
    "# Options: TRUE, FALSE\n",
    "# prob(TRUE) = 0.03, prob(FALSE) = 0.97\n",
    "essential_worker_bl = np.random.choice([False,True],pop_total, p=[0.97, 0.03])\n",
    "\n",
    "\n",
    "# institutionalized_bl\n",
    "# Options: TRUE, FALSE\n",
    "# prob(TRUE) = 0.03, prob(FALSE) = 0.97\n",
    "institutionalized_bl = np.random.choice([False,True],pop_total, p=[0.97, 0.03])\n",
    "\n",
    "\n",
    "# Be careful! Note that a record may be assigned with a value that does not make sense,\n",
    "# e.g. essential worker with an age of 6 years. \n",
    "# Clean up these mismatches at a later time.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "13ae5a56-f1fb-4621-bbdd-11be0501058d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 8580 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6596 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 7408 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 7316 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "# Generate vaccination dates for the different vaccine brands (i.e. laboratories)\n",
    "# Generate time distributions between first and second dose. !!!Note that each vaccine brand (laboratory) has its own distribution. \n",
    "# Options:  'BP' (BioNTech-Pfizer), 'MD' (Moderna), 'AZ' (Astrazeneca), 'JJ' (Johnson-Johnson), 'NV' (Novavax) (or empty)\n",
    "\n",
    "from collections import Counter\n",
    "\n",
    "# Distribution of laboratories, percentage of the population vaccinated by each laboratory\n",
    "# prob('BP' (BioNTech-Pfizer)) = 0.70, \n",
    "# prob('MD' (Moderna)) = 0.15, \n",
    "# prob('AZ' (Astrazeneca)) = 0.10, \n",
    "# prob('JJ' (Johnson-Johnson)) = 0.05,\n",
    "# prob('NV' (Novavax) (or empty)) = 0\n",
    "\n",
    "lab_distribution = np.random.choice(np.arange(0, 5),pop_vaccinated, p=[0.70, 0.15, 0.10, 0.05, 0.00])\n",
    "plt.hist(lab_distribution, bins='auto')\n",
    "plt.show()\n",
    "\n",
    "# population per vaccine brand (i.e. lab)\n",
    "populationPerLab = dict(Counter(lab_distribution))\n",
    "populationBP = populationPerLab[0]\n",
    "populationMD = populationPerLab[1]\n",
    "populationAZ = populationPerLab[2]\n",
    "populationJJ = populationPerLab[3]\n",
    "\n",
    "\n",
    "#BioNTech - 21 days interval between first and second dose\n",
    "mu=21\n",
    "sigma=1\n",
    "days_from_first_dosis_bp = np.random.normal(mu, sigma, populationBP)\n",
    "days_from_first_dosis_bp = [round(x) for x in days_from_first_dosis_bp]\n",
    "days_from_first_dosis_bp = [21 if x<21 else x for x in days_from_first_dosis_bp]\n",
    "plt.hist(days_from_first_dosis_bp, bins='auto')\n",
    "plt.show()\n",
    "\n",
    "#Moderna - 28 days interval between first and second dose\n",
    "mu=28\n",
    "sigma=1\n",
    "days_from_first_dosis_md = np.random.normal(mu, sigma, populationMD)\n",
    "days_from_first_dosis_md = [round(x) for x in days_from_first_dosis_md]\n",
    "days_from_first_dosis_md = [21 if x<21 else x for x in days_from_first_dosis_md]\n",
    "plt.hist(days_from_first_dosis_md, bins='auto')\n",
    "plt.show()\n",
    "\n",
    "#Astrazeneca 28 to 84 days interval between first and second dose\n",
    "mu=58\n",
    "sigma=6\n",
    "days_from_first_dosis_az = np.random.normal(mu, sigma, populationAZ)\n",
    "days_from_first_dosis_az = [round(x) for x in days_from_first_dosis_az]\n",
    "days_from_first_dosis_az = [21 if x<21 else x for x in days_from_first_dosis_az]\n",
    "plt.hist(days_from_first_dosis_az, bins='auto')\n",
    "plt.show()\n",
    "\n",
    "# JJ Single dose\n",
    "# No time interval is necessary between first and second doses.\n",
    "\n",
    "\n",
    "#shuffle values\n",
    "np.random.shuffle(days_from_first_dosis_bp)\n",
    "np.random.shuffle(days_from_first_dosis_md)\n",
    "np.random.shuffle(days_from_first_dosis_az)\n",
    "\n",
    "# In what period since the vaccination started is it likely that you will be vaccinated?\n",
    "# This distribution is necessary to move the first and second doses to the appropriate period\n",
    "period_bp = np.random.binomial(12, 0.35, size=populationBP)\n",
    "period_md = np.random.binomial(12, 0.35, size=populationMD)\n",
    "period_az = np.random.binomial(12, 0.35, size=populationAZ)\n",
    "period_jj = np.random.binomial(12, 0.35, size=populationJJ)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ee18a2f9-c19e-4243-aecf-e18dd62e32e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate dates of first and second doses in real periods and create the dataset with fully vaccinated people\n",
    "# Methodology:\n",
    "# - Generate a random date in a period greater than the study period (1)\n",
    "# - Generate the date of the first dose taking into account the days between the first and second dose calculated in the previous cell (2)\n",
    "# - Generate real date from the beginning of vaccination, in the period where you have to be vaccinated (3)\n",
    "# - Calculate the number of months between the fictitious dates and the real month where you are due for vaccination (4)\n",
    "# - Move fictitious dates of the first and second dose to the real period (5)\n",
    "# - Add person to dataset with fictitious index, laboratory, first dose dt , second dose dt and number of doses (6)\n",
    "\n",
    "#'BP' (BioNTech-Pfizer)\n",
    "fake_dates = [fake.date_between_dates(date_start=datetime.strptime('2024-01-01', '%Y-%m-%d'), \n",
    "                                      date_end=datetime.strptime('2024-12-31', '%Y-%m-%d')) for x in range(0,populationBP)] # (1)\n",
    "bp_population = []\n",
    "for i in range(0,len(fake_dates)):\n",
    "    date_2_dosis = fake_dates[i] # (1)\n",
    "    days_between = days_from_first_dosis_bp[i] # (2)\n",
    "    date_1_dosis = date_2_dosis - timedelta(days=days_between) # (2)\n",
    "    real_period = period_bp[i] # (3)\n",
    "    ref_date = vaccination_starting_date  + relativedelta(months=real_period) # (3)\n",
    "    num_months = (date_1_dosis.year - ref_date.year) * 12 + (date_1_dosis.month - ref_date.month) # (4)\n",
    "    date_1_dosis = date_1_dosis - relativedelta(months=num_months) # (5)\n",
    "    date_2_dosis = date_2_dosis - relativedelta(months=num_months) # (5)\n",
    "    bp_population.append([i,'BP',date_1_dosis,date_2_dosis,2]) # (6)\n",
    "\n",
    "#'MD' (Moderna)\n",
    "fake_dates = [fake.date_between_dates(date_start=datetime.strptime('2024-01-01', '%Y-%m-%d'), \n",
    "                                      date_end=datetime.strptime('2024-12-31', '%Y-%m-%d')) for x in range(0,populationMD)]\n",
    "md_population = []\n",
    "for i in range(0,len(fake_dates)):\n",
    "    date_2_dosis = fake_dates[i] \n",
    "    days_between = days_from_first_dosis_md[i]\n",
    "    date_1_dosis = date_2_dosis - timedelta(days=days_between)\n",
    "    real_period = period_md[i]\n",
    "    ref_date = vaccination_starting_date  + relativedelta(months=real_period)\n",
    "    num_months = (date_1_dosis.year - ref_date.year) * 12 + (date_1_dosis.month - ref_date.month)\n",
    "    date_1_dosis = date_1_dosis - relativedelta(months=num_months)\n",
    "    date_2_dosis = date_2_dosis - relativedelta(months=num_months)\n",
    "    md_population.append([i,'MD',date_1_dosis,date_2_dosis,2])\n",
    "\n",
    "#'AZ' (Astrazeneca)  \n",
    "fake_dates = [fake.date_between_dates(date_start=datetime.strptime('2024-01-01', '%Y-%m-%d'), \n",
    "                                      date_end=datetime.strptime('2024-12-31', '%Y-%m-%d')) for x in range(0,populationAZ)]\n",
    "az_population = []\n",
    "for i in range(0,len(fake_dates)):\n",
    "    date_2_dosis = fake_dates[i]\n",
    "    days_between = days_from_first_dosis_az[i]\n",
    "    date_1_dosis = date_2_dosis - timedelta(days=days_between)\n",
    "    real_period = period_az[i]\n",
    "    ref_date = vaccination_starting_date  + relativedelta(months=real_period)\n",
    "    num_months = (date_1_dosis.year - ref_date.year) * 12 + (date_1_dosis.month - ref_date.month)\n",
    "    date_1_dosis = date_1_dosis - relativedelta(months=num_months)\n",
    "    date_2_dosis = date_2_dosis - relativedelta(months=num_months)\n",
    "    az_population.append([i,'AZ',date_1_dosis,date_2_dosis,2]) \n",
    "\n",
    "\n",
    "# 'JJ' (Johnson-Johnson)\n",
    "# In this case it is only necessary to transfer the first dose because it does not have a second dose\n",
    "fake_dates = [fake.date_between_dates(date_start=datetime.strptime('2024-01-01', '%Y-%m-%d'),\n",
    "                                      date_end=datetime.strptime('2024-12-31', '%Y-%m-%d')) for x in range(0,populationJJ)]\n",
    "jj_population = []\n",
    "for i in range(0,len(fake_dates)):\n",
    "    date_2_dosis = fake_dates[i]\n",
    "    real_period = period_jj[i]\n",
    "    ref_date = vaccination_starting_date  + relativedelta(months=real_period)\n",
    "    num_months = (date_1_dosis.year - ref_date.year) * 12 + (date_1_dosis.month - ref_date.month)\n",
    "    date_1_dosis = date_1_dosis - relativedelta(months=num_months)\n",
    "    jj_population.append([i,'JJ',date_1_dosis,None,1]) \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2c63bd2-1119-4bdf-b2e2-c0bcf233c5d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Assign third dose\n",
    "# Probability of having third dose =0.6\n",
    "third_dosis_bl_bp = np.random.choice(\n",
    "    [False, True], populationBP, p=[0.40, 0.60])\n",
    "third_dosis_bl_md = np.random.choice(\n",
    "    [False, True], populationMD, p=[0.40, 0.60])\n",
    "third_dosis_bl_az = np.random.choice(\n",
    "    [False, True], populationAZ, p=[0.40, 0.60])\n",
    "third_dosis_bl_jj = np.random.choice(\n",
    "    [False, True], populationJJ, p=[0.40, 0.60])\n",
    "\n",
    "# Obtain records with a future third dose\n",
    "bp_population_with_3 = [bp_population[index] for index,\n",
    "                        value in enumerate(third_dosis_bl_bp) if value == True]\n",
    "md_population_with_3 = [md_population[index] for index,\n",
    "                        value in enumerate(third_dosis_bl_md) if value == True]\n",
    "az_population_with_3 = [az_population[index] for index,\n",
    "                        value in enumerate(third_dosis_bl_az) if value == True]\n",
    "jj_population_with_3 = [jj_population[index] for index,\n",
    "                        value in enumerate(third_dosis_bl_jj) if value == True]\n",
    "\n",
    "# Laboratories for third dose 'BP' (BioNTech-Pfizer) and 'MD' (Moderna) with probabilities 0.9 and 0.1 respectively (1)\n",
    "# Time between the second and third doses, is set at random between 150 and 210 days (2)\n",
    "# Number of doses, is set at random between 3 and 6 (3)\n",
    "# Generate Dataframe from the data calculated in the previous section (4)\n",
    "data = []\n",
    "for p in bp_population_with_3:\n",
    "    prob_lab = np.random.choice(np.arange(0, 2), 1, p=[0.90, 0.10])  # (1)\n",
    "    dose_3_brand_cd = 'BP' if prob_lab == 0 else 'MD'\n",
    "    data.append({'person_id': p[0],\n",
    "                 'dose_1_brand_cd': p[1], 'dose_1_dt': p[2], 'dose_2_brand_cd': p[1],\n",
    "                 'dose_2_dt': p[3], 'dose_3_brand_cd': dose_3_brand_cd, 'dose_3_dt': p[3]+timedelta(days=np.random.randint(150, 210)), # (2)\n",
    "                'fully_vaccinated_dt': p[3], 'fully_vaccinated_bl': True, 'number_doses': np.random.randint(3, 6)}) # (3) (4)\n",
    "\n",
    "for p in md_population_with_3:\n",
    "    prob_lab = np.random.choice(np.arange(0, 2), 1, p=[0.90, 0.10])\n",
    "    dose_3_brand_cd = 'BP' if prob_lab == 0 else 'MD'\n",
    "    data.append({'person_id': p[0], 'dose_1_brand_cd': p[1], 'dose_1_dt': p[2], 'dose_2_brand_cd': p[1], 'dose_2_dt': p[3], 'dose_3_brand_cd': dose_3_brand_cd, 'dose_3_dt': p[3]+timedelta(days=np.random.randint(150, 210)),\n",
    "                'fully_vaccinated_dt': p[3], 'fully_vaccinated_bl': True, 'number_doses': np.random.randint(3, 6)})\n",
    "\n",
    "for p in az_population_with_3:\n",
    "    prob_lab = np.random.choice(np.arange(0, 2), 1, p=[0.90, 0.10])\n",
    "    dose_3_brand_cd = 'BP' if prob_lab == 0 else 'MD'\n",
    "    data.append({'person_id': p[0], 'dose_1_brand_cd': p[1], 'dose_1_dt': p[2], 'dose_2_brand_cd': p[1],\n",
    "                 'dose_2_dt': p[3], 'dose_3_brand_cd': dose_3_brand_cd, 'dose_3_dt': p[3]+timedelta(days=np.random.randint(150, 210)),\n",
    "                'fully_vaccinated_dt': p[3], 'fully_vaccinated_bl': True, 'number_doses': np.random.randint(3, 6)})\n",
    "\n",
    "for p in jj_population_with_3:\n",
    "    prob_lab = np.random.choice(np.arange(0, 2), 1, p=[0.90, 0.10])\n",
    "    dose_3_brand_cd = 'BP' if prob_lab == 0 else 'MD'\n",
    "    data.append({'person_id': p[0], 'dose_1_brand_cd': p[1], 'dose_1_dt': p[2], 'dose_2_brand_cd': dose_3_brand_cd,\n",
    "                 'dose_2_dt': p[2]+timedelta(days=np.random.randint(150, 210)),\n",
    "                'fully_vaccinated_dt': p[2], 'fully_vaccinated_bl': True, 'number_doses': np.random.randint(2, 6)})\n",
    "\n",
    "bp_population_without_3 = [bp_population[index]\n",
    "                           for index, value in enumerate(third_dosis_bl_bp) if value == 0]\n",
    "md_population_without_3 = [md_population[index]\n",
    "                           for index, value in enumerate(third_dosis_bl_md) if value == 0]\n",
    "az_population_without_3 = [az_population[index]\n",
    "                           for index, value in enumerate(third_dosis_bl_az) if value == 0]\n",
    "jj_population_without_3 = [jj_population[index]\n",
    "                           for index, value in enumerate(third_dosis_bl_jj) if value == 0]\n",
    "\n",
    "for p in bp_population_without_3:\n",
    "    data.append({'person_id': p[0], 'dose_1_brand_cd': p[1], 'dose_1_dt': p[2], 'dose_2_brand_cd': p[1],\n",
    "                 'dose_2_dt': p[3],\n",
    "                'fully_vaccinated_dt': p[3], 'fully_vaccinated_bl': True, 'number_doses': p[4]})\n",
    "\n",
    "for p in md_population_without_3:\n",
    "    data.append({'person_id': p[0], 'dose_1_brand_cd': p[1], 'dose_1_dt': p[2], 'dose_2_brand_cd': p[1], 'dose_2_dt': p[3],\n",
    "                'fully_vaccinated_dt': p[3], 'fully_vaccinated_bl': True, 'number_doses': p[4]})\n",
    "\n",
    "for p in az_population_without_3:\n",
    "    data.append({'person_id': p[0], 'dose_1_brand_cd': p[1], 'dose_1_dt': p[2], 'dose_2_brand_cd': p[1],\n",
    "                 'dose_2_dt': p[3],\n",
    "                'fully_vaccinated_dt': p[3], 'fully_vaccinated_bl': True, 'number_doses': p[4]})\n",
    "\n",
    "for p in jj_population_without_3:\n",
    "    data.append({'person_id': p[0], 'dose_1_brand_cd': p[1], 'dose_1_dt': p[2],\n",
    "                'fully_vaccinated_dt': p[2], 'fully_vaccinated_bl': True, 'number_doses': p[4]})\n",
    "\n",
    "df_vacunas = pd.DataFrame(data)\n",
    "df_vacunas['person_id'] = np.arange(len(df_vacunas))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e98b74cf-2fba-412f-a473-0d4861780df1",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# - Create the first dataset with the basic information defined at the beginning\n",
    "# - Clean previously mentioned missmatches (e.g. essential worker with an age of 6 years) (2)\n",
    "# - Perform left join with the generated vaccine dataset  (3)\n",
    "# - This gives us vaccinated and non-vaccinated dataset\n",
    "\n",
    "import numpy as np\n",
    "data = {'age_nm': age_nm,\n",
    "        'sex_cd': sex_cd,\n",
    "        'socecon_lvl_cd': socecon_lvl_cd,\n",
    "        'residence_area_cd': residence_area_cd,\n",
    "        'country_cd': country_cd,\n",
    "        'foreign_bl': foreign_bl,\n",
    "        'essential_worker_bl': essential_worker_bl,\n",
    "        'institutionalized_bl': institutionalized_bl}\n",
    "\n",
    "df1 = pd.DataFrame(data, columns=['age_nm', 'sex_cd', 'socecon_lvl_cd', 'residence_area_cd',\n",
    "                   'country_cd', 'foreign_bl', 'essential_worker_bl', 'institutionalized_bl'])\n",
    "\n",
    "\n",
    "df1['essential_worker_bl'] = df1.apply(lambda x: True if ((x['age_nm'] >= 18 and x['essential_worker_bl'] == True )) else False,axis=1) # (2)\n",
    "df1['institutionalized_bl'] = df1.apply(lambda x: True if ((x['age_nm'] >= 18 and x['institutionalized_bl'] == True )) else False,axis=1) # (2)\n",
    "\n",
    "\n",
    "# shuffle dataframe \n",
    "df1 = df1.sample(frac=1).reset_index(drop=True)\n",
    "\n",
    "# Generate person_id to match with vaccinated\n",
    "df1['person_id'] = np.arange(len(df1))   \n",
    "df_total = pd.merge(df1, df_vacunas, on='person_id', how='left') # (3)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ebd9a0b-208f-45cf-b4a1-e659fd7f33c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Generate vaccination_schedule_cd, and concatenate: dose_1_brand_cd and dose_2_brand_cd, taking into account that 'JJ' (Johnson-Johnson) only has a single dose (1)\n",
    "# Set to null the invalid concatenations because they are not vaccinated (2)\n",
    "# Set fully_vaccinated_bl to False for all records with fully_vaccinated_bl with null value (3)\n",
    "# If the number of doses is null, set the value to 0 (4)\n",
    "\n",
    "df_total['vaccination_schedule_cd'] = df_total.apply(lambda x: str(x['dose_1_brand_cd'])+'-'+str(x['dose_2_brand_cd']) if x['dose_1_brand_cd'] !=\"JJ\" else x['dose_1_brand_cd'] ,axis=1) # (1)\n",
    "df_total['vaccination_schedule_cd'] = np.where((df_total['vaccination_schedule_cd']=='nan-nan'), None,df_total['vaccination_schedule_cd']) # (2)\n",
    "df_total['fully_vaccinated_bl'] = np.where((df_total['fully_vaccinated_bl'].isnull()), False,True) # (3)\n",
    "df_total['number_doses'] = np.where((df_total['number_doses'].isnull()), 0,df_total['number_doses']) # (4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9aa68688-7b81-46dd-927a-1382d89c16b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note that there may be records with only one dose administered (NOT Johnson-Johnson)\n",
    "# The population with only one dose is 5% of the unvaccinated (1)\n",
    "# Generate a random date in the first year of vaccination (2)\n",
    "# Assign to the record the date of first vaccination and only one dose received (3)\n",
    "\n",
    "index_list= list(np.where(df_total[\"fully_vaccinated_bl\"] == False)[0]) # (1)\n",
    "pop_with_1_dosis = int(len(index_list)*0.05) # (1)\n",
    "index_list = np.random.choice(index_list, pop_with_1_dosis)\n",
    "lab_distribution = np.random.choice(['BP','MD','AZ'],pop_with_1_dosis, p=[0.10, 0.4, 0.5]) # Laboratories available with their probabilities.\n",
    "i = 0\n",
    "for index in index_list:\n",
    "    date = fake.date_between_dates(date_start=datetime.strptime('2021-01-15', '%Y-%m-%d') , date_end=datetime.strptime('2021-10-15', '%Y-%m-%d'))  # (2)\n",
    "    df_total.at[index,'dose_1_brand_cd']=lab_distribution[i] # (3)\n",
    "    df_total.at[index,'dose_1_dt']=date # (3)\n",
    "    df_total.at[index,'number_doses']=1 # (3)\n",
    "    i+=1\n",
    "\n",
    "df_total['dose_1_dt'] = pd.to_datetime(df_total['dose_1_dt'])\n",
    "df_total['dose_2_dt'] = pd.to_datetime(df_total['dose_2_dt'])\n",
    "df_total['dose_3_dt'] = pd.to_datetime(df_total['dose_3_dt'])\n",
    "df_total['fully_vaccinated_dt'] = pd.to_datetime(df_total['fully_vaccinated_dt'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6462cb59-05e7-4a22-bf1d-21da1fd54003",
   "metadata": {},
   "outputs": [],
   "source": [
    "# variant_cd Options: 'alpha', 'beta', 'delta', 'gamma', 'omicron BA.1', 'omicron BA.2', 'lambda', 'mu'\n",
    "# variant_cd data at indiviual level is not available for Aragon, Spain\n",
    "\n",
    "df_total['variant_cd'] = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13344063-6060-401a-80f2-319844e1250c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this section we try to simulate the behavior of different epidemiological waves (in this case the behavior of Aragon (Spain)).\n",
    "# Reference: https://www.aragon.es/coronavirus/situacion-actual/evolucion-diaria\n",
    "\n",
    "# Methodology:\n",
    "# - Select dates and their respective values representing peaks and valleys of the distribution\n",
    "#   Change the dates by the number of days between the date and the date of onset of cases (1)\n",
    "# - Generate an interpolating polynomial, in this case 1-D interpolating spline of 5 degree (2)\n",
    "# - For all pandemic days generate the number of cases from the interpolating polynomial (3)\n",
    "# - Add noise to the resulting distribution (4)\n",
    "#!!Note that the values of confirmed cases per day to generate the polynomial do not fit the number of records we are going to generate. \n",
    "# - Apply a reduction factor to the above distribution taking into account the number of records to be generated (5)\n",
    "# - Generate as many dates as cases per day and day to be assigned as confirmed_case_dt \n",
    "#   (fill with null dates until the total population (not confirmed) is reached) (6)\n",
    "# -  Assign to the record its confirmation date plus the type of test used for confirmation (7)\n",
    "# - Add information on confirmed cases to the general dataset (8)\n",
    "\n",
    "date = ['2020-02-29', '2020-03-29', '2020-05-29', '2020-08-08', '2020-09-22', '2020-10-30', '2020-12-16', '2021-01-27',\n",
    "        '2021-03-24', '2021-04-30', '2021-06-19', '2021-07-21', '2021-10-21', '2022-01-11', '2022-04-10', '2022-05-10', '2022-05-20']  # (1)\n",
    "values = [2, 266, 17, 699, 317, 1264, 247, 869,  145, 379, 90, 1016, 58, 8224, 230, 692, 200] # (1)\n",
    "\n",
    "days = []\n",
    "for f in date:\n",
    "    day = datetime.strptime(f, '%Y-%m-%d') - covid_starting_date\n",
    "    day = day.days\n",
    "    days.append(day) # (1)\n",
    "\n",
    "p = InterpolatedUnivariateSpline(days, values, k=5) #(2)\n",
    "\n",
    "s = end_date - covid_starting_date # (3)\n",
    "s = np.arange(s.days) # (3)\n",
    "days = [x+1 for x in s] # (3)\n",
    "cases = p(days)  # (3)\n",
    "cases = [1 if x < 0 else x for x in cases] # (3)\n",
    "\n",
    "\n",
    "cases = [x + np.random.randint(-30, 120) if x >\n",
    "         40 else x + np.random.randint(5, 50) for x in cases] # (4)\n",
    "\n",
    "plt.plot(days, cases)\n",
    "df_distr_cases = pd.DataFrame({'day': days, 'cases': cases},\n",
    "                         columns=['day', 'cases'])\n",
    "\n",
    "\n",
    "reduction_factor = pop_with_covid / sum(df_distr_cases['cases']) # (5)\n",
    "df_distr_cases['cases'] = df_distr_cases.apply(\n",
    "    lambda x: int(reduction_factor*x['cases']), axis=1)\n",
    "df_distr_cases.plot.line(x='day', y='cases')  # (5)\n",
    "\n",
    "\n",
    "list_dates = np.ndarray(0)\n",
    "for index, row in df_distr_cases.iterrows():\n",
    "    date_confirmed = (covid_starting_date + timedelta(days=int(row['day'])))\n",
    "    if date_confirmed >= end_date:\n",
    "        date_confirmed = end_date\n",
    "        date_confirmed - timedelta(days=1)\n",
    "    partial_list = np.full(int(row['cases']), date_confirmed)\n",
    "    list_dates = np.concatenate((list_dates, partial_list), axis=None) # (6)\n",
    "\n",
    "\n",
    "list_dates = np.concatenate((list_dates, np.full(\n",
    "    pop_total - len(list_dates), pd.NaT)), axis=None) # (6)\n",
    "np.random.shuffle(list_dates) # (6)\n",
    "\n",
    "\n",
    "data = {'confirmed_case_dt': list_dates}\n",
    "\n",
    "\n",
    "df_confirmed = pd.DataFrame(data, columns=['confirmed_case_dt']) # (7)\n",
    "lista_test = np.random.choice(['PCR', 'AG', 'other', None], len(\n",
    "    df_confirmed), p=[0.60, 0.35, 0.04, 0.01])\n",
    "np.random.shuffle(lista_test)\n",
    "\n",
    "df_confirmed['person_id'] = np.arange(len(df_confirmed)) # (7)\n",
    "df_confirmed['confirmed_case_bl'] = np.where(\n",
    "    (df_confirmed['confirmed_case_dt'].isnull()), False, True) # (7)\n",
    "\n",
    "df_confirmed['test_type_cd'] = None # (7)\n",
    "\n",
    "index_list = list(np.where(df_confirmed[\"confirmed_case_bl\"] == True)[0]) # (7)\n",
    "i = 0\n",
    "for index in index_list:\n",
    "    df_confirmed.at[index, 'test_type_cd'] = lista_test[i] # (7)\n",
    "    i += 1\n",
    "\n",
    "\n",
    "df_total = pd.merge(df_total, df_confirmed, on='person_id', how='left') # (8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b42cb1fd-1c8c-4371-aaba-6dbe88c33015",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate previous infections\n",
    "# Assign a date between 90 and 280 days prior to the date of case confirmation\n",
    "\n",
    "df_total['previous_infection_dt'] = pd.NaT\n",
    "df_total['previous_infection_bl'] = False\n",
    "\n",
    "index_list= list(np.where(df_total[\"confirmed_case_dt\"] >= '2021-11-11')[0])\n",
    "index_list = np.random.choice(index_list, pop_with_previous)\n",
    "\n",
    "for index in index_list:\n",
    "    last_confirmed = df_total.iloc[index]['confirmed_case_dt']\n",
    "    days_ = np.random.randint(90,280)\n",
    "    previous_infection_dt = last_confirmed - timedelta(days=days_)\n",
    "    df_total.at[index,'previous_infection_dt']=previous_infection_dt\n",
    "    df_total.at[index,'previous_infection_bl']=True\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a060f709-a7bd-460c-9c48-96322784ae10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate comorbidities with a 10% flat prevalence in the total population (1)\n",
    "# Please, note that nonsensical combinations may occur. Profile comorbidities in a meaningful way (2)\n",
    "# Please, note that changes in comorbidities profiling might change the expected distributions generated thus changing the overall percentages\n",
    "\n",
    "# Modified:\n",
    "# Diabetes - Probability distribution: prob(TRUE) = 0.05, prob(FALSE) = 0.95\n",
    "# Heart failure - Probability distribution: prob(TRUE) = 0.02, prob(FALSE) = 0.98\n",
    "# COPD - Probability distribution: prob(TRUE) = 0.05, prob(FALSE) = 0.95\n",
    "# Solid tumor without metastasis - Probability distribution: prob(TRUE) = 0.01, prob(FALSE) = 0.99\n",
    "# Hypertension - Probability distribution: prob(TRUE) = 0.14, prob(FALSE) = 0.86\n",
    "# Blood cancer - Probability distribution: prob(TRUE) = 0.01, prob(FALSE) = 0.99\n",
    "# Immunosuppression - Probability distribution: prob(TRUE) = 0.01,prob(FALSE) = 0.99\n",
    "\n",
    "data = {'diabetes_bl': np.random.choice([True,False],pop_total, p=[0.05, 0.95]),\n",
    "       'obesity_bl':np.random.choice([True,False],pop_total, p=[0.10, 0.90]),\n",
    "       'heart_failure_bl':np.random.choice([True,False],pop_total, p=[0.02, 0.98]),\n",
    "       'copd_bl':np.random.choice([True,False],pop_total, p=[0.05, 0.95]),\n",
    "       'solid_tumor_without_metastasis_bl':np.random.choice([True,False],pop_total, p=[0.01, 0.99]),\n",
    "       'chronic_kidney_disease_bl':np.random.choice([True,False],pop_total, p=[0.10, 0.90]),\n",
    "       'sickle_cell_disease_bl':np.random.choice([True,False],pop_total, p=[0.10, 0.90]),\n",
    "       'hypertension_bl':np.random.choice([True,False],pop_total, p=[0.14, 0.86]),\n",
    "       'chronic_liver_disease_bl':np.random.choice([True,False],pop_total, p=[0.10, 0.90]),\n",
    "       'blood_cancer_bl':np.random.choice([True,False],pop_total, p=[0.01, 0.99]),\n",
    "       'transplanted_bl':np.random.choice([True,False],pop_total, p=[0.10, 0.90]),\n",
    "       'hiv_infection_bl':np.random.choice([True,False],pop_total, p=[0.10, 0.90]),\n",
    "       'immunosuppression_bl':np.random.choice([True,False],pop_total, p=[0.01, 0.99]),\n",
    "       'primary_immunodeficiency_bl':np.random.choice([True,False],pop_total, p=[0.10, 0.90]),\n",
    "       'pregnancy_bl':np.random.choice([True,False],pop_total, p=[0.10, 0.90])}\n",
    "\n",
    "\n",
    "\n",
    "df1=pd.DataFrame(data, columns=['diabetes_bl','obesity_bl','heart_failure_bl','copd_bl','solid_tumor_without_metastasis_bl',\n",
    "                                'chronic_kidney_disease_bl','sickle_cell_disease_bl','hypertension_bl','chronic_liver_disease_bl','blood_cancer_bl','transplanted_bl',\n",
    "                                'hiv_infection_bl','immunosuppression_bl','primary_immunodeficiency_bl','pregnancy_bl'])\n",
    "df1['person_id'] = np.arange(len(df1))\n",
    "\n",
    "df_total = pd.merge(df_total,df1,on='person_id',how='left')\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcf8dae9-2544-44b5-9106-32675601ffca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Profiling comorbidities\n",
    "# - Assign false to those records that pregnancy_bl is true but do not meet age and gender criteria (1)\n",
    "# - Assign false to records where they do not satisfy the age criteria (2)\n",
    "# Please, note that comorbidities profiles could be expanded to fit more complex scenarios\n",
    "\n",
    "df_total['pregnancy_bl'] = df_total.apply(lambda x: True if ((x['age_nm'] >= 15 and x['age_nm']<=45 \n",
    "                                                              and x['sex_cd']==2 and x['pregnancy_bl'] == True )) else False,axis=1) # (1)\n",
    "\n",
    "df_total['heart_failure_bl'] = df_total.apply(lambda x: True if (x['age_nm'] > 40 and x['heart_failure_bl'] == True ) else False,axis=1) # (2)\n",
    "df_total['copd_bl'] = df_total.apply(lambda x: True if (x['age_nm'] > 40 and x['copd_bl'] == True ) else False,axis=1) # (2)\n",
    "df_total['chronic_kidney_disease_bl'] = df_total.apply(lambda x: True if (x['age_nm'] > 40 and x['chronic_kidney_disease_bl'] == True ) else False,axis=1) # (2)\n",
    "df_total['chronic_liver_disease_bl'] = df_total.apply(lambda x: True if (x['age_nm'] > 40 and x['chronic_liver_disease_bl'] == True ) else False,axis=1) # (2)\n",
    "df_total['hypertension_bl'] = df_total.apply(lambda x: True if (x['age_nm'] > 40 and x['hypertension_bl'] == True ) else False,axis=1) # (2)\n",
    "df_total['hiv_infection_bl'] = df_total.apply(lambda x: True if (x['age_nm'] > 40 and x['hiv_infection_bl'] == True ) else False,axis=1) # (2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e4faa6e-64cb-418e-8627-a884ecac00b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# In this section we try to simulate the observed daily mortality behavior for our study period (in this case the behavior of Aragon (Spain))\n",
    "# Reference: https://momo.isciii.es/panel_momo/ (Aragon)\n",
    "\n",
    "# Methodology:\n",
    "# - Select dates and their respective values representing peaks and valleys of the distribution\n",
    "#   Change the dates by the number of days between the date and the date of onset of cases (1)\n",
    "# - Generate an interpolating polynomial, in this case 1-D interpolating spline of 5 degree (2)\n",
    "# - For all pandemic days generate the number of deaths from the interpolating polynomial (3)\n",
    "# - Add noise to the resulting distribution (4)\n",
    "#!!Note that the values of exitus cases per day to generate the polynomial do not fit the number of records we are going to generate.\n",
    "#   Apply a reduction factor to the above distribution taking into account the number of records to be generated (5)\n",
    "# - Generate as many dates as cases per day and day to be assigned as exitus_dt (6)\n",
    "# - Assign to the record its exitus date plus exitus_bl (7)\n",
    "# - Delete dates of the first 10 days, to accumulate population susceptible to death (8) \n",
    "# - Delete dates after date of death (assigned previously) (9)\n",
    "# - Add the information of the exitus cases to the general dataset (10)\n",
    "\n",
    "date = ['2020-02-28','2020-03-01','2020-04-01','2020-05-23','2020-07-02','2020-08-07','2020-09-13', '2020-11-13','2020-12-22','2021-01-22','2021-03-30',\n",
    "        '2021-06-09','2021-10-16','2022-01-09','2022-02-15','2022-03-27','2022-04-22','2022-05-12']  # (1)\n",
    "values = [41,41,87,27,37,65,44,80,43,61,43,39,44,52,50,32,32,40] # (1)\n",
    "\n",
    "days = []\n",
    "for f in date:\n",
    "    day = datetime.strptime(f, '%Y-%m-%d') - covid_starting_date\n",
    "    day = day.days\n",
    "    days.append(day) # (1)\n",
    "\n",
    "p = InterpolatedUnivariateSpline(days, values, k=5) # (2)\n",
    "\n",
    "s = end_date - covid_starting_date # (3)\n",
    "s = np.arange(s.days) # (3)\n",
    "days = [x+1 for x in s] # (3)\n",
    "existus_cases = p(days)  # (3)\n",
    "existus_cases = [np.random.randint(20,23) if x < 10 else int(x) for x in existus_cases] # (3)\n",
    "\n",
    "existus_cases = [x + np.random.randint(-3, 3) for x in existus_cases] # (4)\n",
    "\n",
    "plt.plot(days, existus_cases)\n",
    "df_distr_cases = pd.DataFrame({'day': days, 'existus_cases': existus_cases},\n",
    "                         columns=['day', 'existus_cases'])\n",
    "\n",
    "\n",
    "reduction_factor = pop_exitus / sum(df_distr_cases['existus_cases']) # (5)\n",
    "df_distr_cases['existus_cases'] = df_distr_cases.apply(\n",
    "    lambda x: int(reduction_factor*x['existus_cases']), axis=1)\n",
    "\n",
    "\n",
    "\n",
    "df_distr_cases.plot.line(x='day', y='existus_cases')  # (5)\n",
    "\n",
    "list_dates = []\n",
    "for index, row in df_distr_cases.iterrows():\n",
    "    date_exitus = (covid_starting_date + timedelta(days=int(row['day'])))\n",
    "    if date_exitus >= end_date:\n",
    "        date_exitus = end_date\n",
    "        date_exitus - timedelta(days=1)\n",
    "    list_dates.append({'exitus_dt':date_exitus,'counts':row['existus_cases']}) # (6)\n",
    "\n",
    "df_group_exitus = pd.DataFrame(list_dates, columns=['exitus_dt','counts']) # (6) \n",
    "\n",
    "\n",
    "\n",
    "df_total['exitus_bl'] = False  # (7)\n",
    "df_total['exitus_dt'] = pd.NaT # (7)\n",
    "\n",
    "\n",
    "n = 10 \n",
    "df_group_exitus = df_group_exitus.iloc[10:,:] # (8)\n",
    "\n",
    "\n",
    "for index,row in df_group_exitus.iterrows():    \n",
    "    # Select susceptible population, not deceased, confirmed with date prior to death and at least one previous vaccination    \n",
    "    exitus_dt_i = row.exitus_dt    \n",
    "    pop1 = np.where((df_total['exitus_bl']==False) & ((df_total['confirmed_case_dt'] <= exitus_dt_i) | (df_total['dose_1_dt'].isnull())))[0]\n",
    "    pop2 = np.where((df_total['exitus_bl']==False) & ((df_total['dose_1_dt'] <= exitus_dt_i) & (df_total['confirmed_case_dt'] <= exitus_dt_i)))[0]\n",
    "    pop3 = np.where((df_total['exitus_bl']==False) & ((df_total['dose_1_dt'] <= exitus_dt_i) & (df_total['confirmed_case_dt'].isnull())))[0]\n",
    "    pop = list(set(pop1).union(set(pop2),set(pop3)))\n",
    "                    \n",
    "    _sample = np.random.choice(pop,size=row.counts)\n",
    "  \n",
    "    for i in _sample:\n",
    "        dose_2_dt = df_total.iloc[i]['dose_2_dt']\n",
    "        dose_3_dt = df_total.iloc[i]['dose_3_dt']\n",
    "        number_doses = df_total.iloc[i]['number_doses']\n",
    "        if i in pop1:\n",
    "            pass\n",
    "        else:\n",
    "            if number_doses == 2:\n",
    "                if dose_2_dt > exitus_dt_i:                    \n",
    "                    df_total.at[i,'dose_2_dt']=pd.NaT              # (9)\n",
    "                    df_total.at[i,'dose_3_dt']=pd.NaT              # (9)\n",
    "                    df_total.at[i,'fully_vaccinated_bl']=False     # (9)\n",
    "                    df_total.at[i,'fully_vaccinated_dt']=pd.NaT    # (9)\n",
    "                    df_total.at[i,'vaccination_schedule_cd']=None  # (9)          \n",
    "                    df_total.at[i,'number_doses']=1                 \n",
    "            elif number_doses >= 3:  \n",
    "                if dose_2_dt > exitus_dt_i:\n",
    "                    df_total.at[i,'dose_2_dt']=pd.NaT              # (9)\n",
    "                    df_total.at[i,'dose_3_dt']=pd.NaT              # (9)\n",
    "                    df_total.at[i,'fully_vaccinated_bl']=False     # (9)\n",
    "                    df_total.at[i,'fully_vaccinated_dt']=pd.NaT    # (9)\n",
    "                    df_total.at[i,'vaccination_schedule_cd']=None  # (9)           \n",
    "                    df_total.at[i,'number_doses']=1                    # (9)\n",
    "                else:\n",
    "                    if dose_3_dt > exitus_dt_i:\n",
    "                        df_total.at[i,'dose_3_dt']=pd.NaT          # (9)\n",
    "                        df_total.at[i,'number_doses']=2                # (9)\n",
    "        df_total.at[i,'exitus_bl'] = True                          #(10)\n",
    "        df_total.at[i,'exitus_dt'] = exitus_dt_i                   #(10)\n",
    "\n",
    "# !!!Note: May generate an unexpected behavior (e.g. some people vaccinated after death). \n",
    "# This behavior is not fixed because it is used to simulate registration errors.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dd1c406-a742-415a-b1e9-eaae9dd394dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check variable data types\n",
    "print(len(df_total.columns))\n",
    "print(df_total.dtypes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1527ece6-e783-4885-959d-b01e7e140197",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Assign a more realistic person_id (1)\n",
    "# Save dataset with the same order of variables as the data model (2)\n",
    "\n",
    "df_total['person_id'] = df_total.apply(lambda x: fake.pystr(max_chars=10) ,axis=1) # (1)\n",
    "df_total = df_total[['person_id','age_nm','sex_cd','socecon_lvl_cd','residence_area_cd','country_cd','foreign_bl','essential_worker_bl','institutionalized_bl',\n",
    "'dose_1_brand_cd','dose_1_dt','dose_2_brand_cd','dose_2_dt','dose_3_brand_cd','dose_3_dt','number_doses','fully_vaccinated_dt','fully_vaccinated_bl','vaccination_schedule_cd',\n",
    "'confirmed_case_dt','confirmed_case_bl','previous_infection_dt','previous_infection_bl','test_type_cd','variant_cd','diabetes_bl','obesity_bl','heart_failure_bl',\n",
    "'copd_bl','solid_tumor_without_metastasis_bl','chronic_kidney_disease_bl','sickle_cell_disease_bl','hypertension_bl','chronic_liver_disease_bl','blood_cancer_bl','transplanted_bl','hiv_infection_bl',\n",
    "'primary_immunodeficiency_bl','immunosuppression_bl','pregnancy_bl','exitus_dt','exitus_bl']] # (2)\n",
    "df_total.to_csv('vaccine_effectiveness_synthetic_pop_10k.csv',index=False,sep='|') # (2)\n",
    "\n",
    "#!!!Note: The synthetic dataset is saved as a CSV file, UTF-8 format, pipe-separated (\"|\")\n"
   ]
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
