# coding=utf-8
# Copyright 2023 The TensorFlow Datasets Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Largest Covid-19 epidemiological database with a powerful set of expansive covariates.

It includes open, publicly sourced, licensed data relating to demographics,
economy, epidemiology, geography, health, hospitalizations, mobility, government
response, weather, and more.
"""

import numpy as np
from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
import tensorflow_datasets.public_api as tfds

_DESCRIPTION = """
This repository attempts to assemble the largest Covid-19 epidemiological database
in addition to a powerful set of expansive covariates. It includes open, publicly sourced,
licensed data relating to demographics, economy, epidemiology, geography, health, hospitalizations,
mobility, government response, weather, and more.

This particular dataset corresponds to a join of all the different tables that are part of the repository.
Therefore, expect the resulting samples to be highly sparse.
"""

_CITATION = """
@article{Wahltinez2020,
  author = "O. Wahltinez and others",
  year = 2020,
  title = "COVID-19 Open-Data: curating a fine-grained, global-scale data repository for SARS-CoV-2",
  note = "Work in progress",
  url = {https://goo.gle/covid-19-open-data},
}
"""

_N_RECORDS = 11266423  # It should be an upper bound
_BATCH_SIZE = 10000


class Covid19(tfds.core.GeneratorBasedBuilder):
  """DatasetBuilder for covid19 dataset."""

  VERSION = tfds.core.Version('1.0.0')
  RELEASE_NOTES = {
      '1.0.0': 'Initial release.',
  }

  def _info(self) -> tfds.core.DatasetInfo:
    """Returns the dataset metadata."""
    return tfds.core.DatasetInfo(
        builder=self,
        description=_DESCRIPTION,
        features=tfds.features.FeaturesDict({
            'location_key': np.str_,
            'date': np.str_,
            'place_id': np.str_,
            'wikidata_id': np.str_,
            'datacommons_id': np.str_,
            'country_code': np.str_,
            'country_name': np.str_,
            'iso_3166_1_alpha_2': np.str_,
            'iso_3166_1_alpha_3': np.str_,
            'aggregation_level': np.float64,
            'new_confirmed': np.float64,
            'new_deceased': np.float64,
            'cumulative_confirmed': np.float64,
            'cumulative_deceased': np.float64,
            'cumulative_tested': np.float64,
            'new_persons_vaccinated': np.float64,
            'cumulative_persons_vaccinated': np.float64,
            'new_persons_fully_vaccinated': np.float64,
            'cumulative_persons_fully_vaccinated': np.float64,
            'new_vaccine_doses_administered': np.float64,
            'cumulative_vaccine_doses_administered': np.float64,
            'population': np.float64,
            'population_male': np.float64,
            'population_female': np.float64,
            'population_rural': np.float64,
            'population_urban': np.float64,
            'population_density': np.float64,
            'human_development_index': np.float64,
            'population_age_00_09': np.float64,
            'population_age_10_19': np.float64,
            'population_age_20_29': np.float64,
            'population_age_30_39': np.float64,
            'population_age_40_49': np.float64,
            'population_age_50_59': np.float64,
            'population_age_60_69': np.float64,
            'population_age_70_79': np.float64,
            'population_age_80_and_older': np.float64,
            'gdp_usd': np.float64,
            'gdp_per_capita_usd': np.float64,
            'openstreetmap_id': np.str_,
            'latitude': np.float64,
            'longitude': np.float64,
            'area_sq_km': np.float64,
            'smoking_prevalence': np.float64,
            'diabetes_prevalence': np.float64,
            'infant_mortality_rate': np.float64,
            'nurses_per_1000': np.float64,
            'physicians_per_1000': np.float64,
            'health_expenditure_usd': np.float64,
            'out_of_pocket_health_expenditure_usd': np.float64,
            'school_closing': np.float64,
            'workplace_closing': np.float64,
            'cancel_public_events': np.float64,
            'restrictions_on_gatherings': np.float64,
            'public_transport_closing': np.float64,
            'stay_at_home_requirements': np.float64,
            'restrictions_on_internal_movement': np.float64,
            'international_travel_controls': np.float64,
            'income_support': np.float64,
            'debt_relief': np.float64,
            'fiscal_measures': np.float64,
            'international_support': np.float64,
            'public_information_campaigns': np.float64,
            'testing_policy': np.float64,
            'contact_tracing': np.float64,
            'emergency_investment_in_healthcare': np.float64,
            'investment_in_vaccines': np.float64,
            'facial_coverings': np.float64,
            'vaccination_policy': np.float64,
            'stringency_index': np.float64,
            'average_temperature_celsius': np.float64,
            'minimum_temperature_celsius': np.float64,
            'maximum_temperature_celsius': np.float64,
            'rainfall_mm': np.float64,
            'snowfall_mm': np.float64,
            'dew_point': np.float64,
            'relative_humidity': np.float64,
            'new_tested': np.float64,
            'population_largest_city': np.float64,
            'population_clustered': np.float64,
            'human_capital_index': np.float64,
            'area_rural_sq_km': np.float64,
            'area_urban_sq_km': np.float64,
            'life_expectancy': np.float64,
            'adult_male_mortality_rate': np.float64,
            'adult_female_mortality_rate': np.float64,
            'pollution_mortality_rate': np.float64,
            'comorbidity_mortality_rate': np.float64,
            'mobility_retail_and_recreation': np.float64,
            'mobility_grocery_and_pharmacy': np.float64,
            'mobility_parks': np.float64,
            'mobility_transit_stations': np.float64,
            'mobility_workplaces': np.float64,
            'mobility_residential': np.float64,
            'hospital_beds_per_1000': np.float64,
            'subregion1_code': np.str_,
            'subregion1_name': np.str_,
            'new_recovered': np.float64,
            'cumulative_recovered': np.float64,
            'elevation_m': np.float64,
            'new_hospitalized_patients': np.float64,
            'cumulative_hospitalized_patients': np.float64,
            'new_intensive_care_patients': np.float64,
            'cumulative_intensive_care_patients': np.float64,
            'new_confirmed_age_0': np.float64,
            'new_confirmed_age_1': np.float64,
            'new_confirmed_age_2': np.float64,
            'new_confirmed_age_3': np.float64,
            'new_confirmed_age_4': np.float64,
            'new_confirmed_age_5': np.float64,
            'new_confirmed_age_6': np.float64,
            'new_confirmed_age_7': np.float64,
            'new_confirmed_age_8': np.float64,
            'new_confirmed_age_9': np.float64,
            'cumulative_confirmed_age_0': np.float64,
            'cumulative_confirmed_age_1': np.float64,
            'cumulative_confirmed_age_2': np.float64,
            'cumulative_confirmed_age_3': np.float64,
            'cumulative_confirmed_age_4': np.float64,
            'cumulative_confirmed_age_5': np.float64,
            'cumulative_confirmed_age_6': np.float64,
            'cumulative_confirmed_age_7': np.float64,
            'cumulative_confirmed_age_8': np.float64,
            'cumulative_confirmed_age_9': np.float64,
            'new_deceased_age_0': np.float64,
            'new_deceased_age_1': np.float64,
            'new_deceased_age_2': np.float64,
            'new_deceased_age_3': np.float64,
            'new_deceased_age_4': np.float64,
            'new_deceased_age_5': np.float64,
            'new_deceased_age_6': np.float64,
            'new_deceased_age_7': np.float64,
            'new_deceased_age_8': np.float64,
            'new_deceased_age_9': np.float64,
            'cumulative_deceased_age_0': np.float64,
            'cumulative_deceased_age_1': np.float64,
            'cumulative_deceased_age_2': np.float64,
            'cumulative_deceased_age_3': np.float64,
            'cumulative_deceased_age_4': np.float64,
            'cumulative_deceased_age_5': np.float64,
            'cumulative_deceased_age_6': np.float64,
            'cumulative_deceased_age_7': np.float64,
            'cumulative_deceased_age_8': np.float64,
            'cumulative_deceased_age_9': np.float64,
            'new_tested_age_0': np.float64,
            'new_tested_age_1': np.float64,
            'new_tested_age_2': np.float64,
            'new_tested_age_3': np.float64,
            'new_tested_age_4': np.float64,
            'new_tested_age_5': np.float64,
            'new_tested_age_6': np.float64,
            'new_tested_age_7': np.float64,
            'new_tested_age_8': np.float64,
            'new_tested_age_9': np.float64,
            'cumulative_tested_age_0': np.float64,
            'cumulative_tested_age_1': np.float64,
            'cumulative_tested_age_2': np.float64,
            'cumulative_tested_age_3': np.float64,
            'cumulative_tested_age_4': np.float64,
            'cumulative_tested_age_5': np.float64,
            'cumulative_tested_age_6': np.float64,
            'cumulative_tested_age_7': np.float64,
            'cumulative_tested_age_8': np.float64,
            'cumulative_tested_age_9': np.float64,
            'new_hospitalized_patients_age_0': np.float64,
            'new_hospitalized_patients_age_1': np.float64,
            'new_hospitalized_patients_age_2': np.float64,
            'new_hospitalized_patients_age_3': np.float64,
            'new_hospitalized_patients_age_4': np.float64,
            'new_hospitalized_patients_age_5': np.float64,
            'new_hospitalized_patients_age_6': np.float64,
            'new_hospitalized_patients_age_7': np.float64,
            'new_hospitalized_patients_age_8': np.float64,
            'new_hospitalized_patients_age_9': np.float64,
            'cumulative_hospitalized_patients_age_0': np.float64,
            'cumulative_hospitalized_patients_age_1': np.float64,
            'cumulative_hospitalized_patients_age_2': np.float64,
            'cumulative_hospitalized_patients_age_3': np.float64,
            'cumulative_hospitalized_patients_age_4': np.float64,
            'cumulative_hospitalized_patients_age_5': np.float64,
            'cumulative_hospitalized_patients_age_6': np.float64,
            'cumulative_hospitalized_patients_age_7': np.float64,
            'cumulative_hospitalized_patients_age_8': np.float64,
            'cumulative_hospitalized_patients_age_9': np.float64,
            'new_intensive_care_patients_age_0': np.float64,
            'new_intensive_care_patients_age_1': np.float64,
            'new_intensive_care_patients_age_2': np.float64,
            'new_intensive_care_patients_age_3': np.float64,
            'new_intensive_care_patients_age_4': np.float64,
            'new_intensive_care_patients_age_5': np.float64,
            'new_intensive_care_patients_age_6': np.float64,
            'new_intensive_care_patients_age_7': np.float64,
            'new_intensive_care_patients_age_8': np.float64,
            'new_intensive_care_patients_age_9': np.float64,
            'cumulative_intensive_care_patients_age_0': np.float64,
            'cumulative_intensive_care_patients_age_1': np.float64,
            'cumulative_intensive_care_patients_age_2': np.float64,
            'cumulative_intensive_care_patients_age_3': np.float64,
            'cumulative_intensive_care_patients_age_4': np.float64,
            'cumulative_intensive_care_patients_age_5': np.float64,
            'cumulative_intensive_care_patients_age_6': np.float64,
            'cumulative_intensive_care_patients_age_7': np.float64,
            'cumulative_intensive_care_patients_age_8': np.float64,
            'cumulative_intensive_care_patients_age_9': np.float64,
            # Stored as string because values look like "30-39"
            'age_bin_0': np.str_,
            'age_bin_1': np.str_,
            'age_bin_2': np.str_,
            'age_bin_3': np.str_,
            'age_bin_4': np.str_,
            'age_bin_5': np.str_,
            'age_bin_6': np.str_,
            'age_bin_7': np.str_,
            'age_bin_8': np.str_,
            'age_bin_9': np.str_,
            'new_confirmed_male': np.float64,
            'new_confirmed_female': np.float64,
            'cumulative_confirmed_male': np.float64,
            'cumulative_confirmed_female': np.float64,
            'new_deceased_male': np.float64,
            'new_deceased_female': np.float64,
            'cumulative_deceased_male': np.float64,
            'cumulative_deceased_female': np.float64,
            'new_tested_male': np.float64,
            'new_tested_female': np.float64,
            'cumulative_tested_male': np.float64,
            'cumulative_tested_female': np.float64,
            'new_hospitalized_patients_male': np.float64,
            'new_hospitalized_patients_female': np.float64,
            'cumulative_hospitalized_patients_male': np.float64,
            'cumulative_hospitalized_patients_female': np.float64,
            'new_intensive_care_patients_male': np.float64,
            'new_intensive_care_patients_female': np.float64,
            'cumulative_intensive_care_patients_male': np.float64,
            'cumulative_intensive_care_patients_female': np.float64,
            'subregion2_code': np.str_,
            'subregion2_name': np.str_,
            'current_hospitalized_patients': np.float64,
            'current_intensive_care_patients': np.float64,
            'current_ventilator_patients': np.float64,
            'search_trends_abdominal_obesity': np.float64,
            'search_trends_abdominal_pain': np.float64,
            'search_trends_acne': np.float64,
            'search_trends_actinic_keratosis': np.float64,
            'search_trends_acute_bronchitis': np.float64,
            'search_trends_adrenal_crisis': np.float64,
            'search_trends_ageusia': np.float64,
            'search_trends_alcoholism': np.float64,
            'search_trends_allergic_conjunctivitis': np.float64,
            'search_trends_allergy': np.float64,
            'search_trends_amblyopia': np.float64,
            'search_trends_amenorrhea': np.float64,
            'search_trends_amnesia': np.float64,
            'search_trends_anal_fissure': np.float64,
            'search_trends_anaphylaxis': np.float64,
            'search_trends_anemia': np.float64,
            'search_trends_angina_pectoris': np.float64,
            'search_trends_angioedema': np.float64,
            'search_trends_angular_cheilitis': np.float64,
            'search_trends_anosmia': np.float64,
            'search_trends_anxiety': np.float64,
            'search_trends_aphasia': np.float64,
            'search_trends_aphonia': np.float64,
            'search_trends_apnea': np.float64,
            'search_trends_arthralgia': np.float64,
            'search_trends_arthritis': np.float64,
            'search_trends_ascites': np.float64,
            'search_trends_asperger_syndrome': np.float64,
            'search_trends_asphyxia': np.float64,
            'search_trends_asthma': np.float64,
            'search_trends_astigmatism': np.float64,
            'search_trends_ataxia': np.float64,
            'search_trends_atheroma': np.float64,
            'search_trends_attention_deficit_hyperactivity_disorder': (
                np.float64
            ),
            'search_trends_auditory_hallucination': np.float64,
            'search_trends_autoimmune_disease': np.float64,
            'search_trends_avoidant_personality_disorder': np.float64,
            'search_trends_back_pain': np.float64,
            'search_trends_bacterial_vaginosis': np.float64,
            'search_trends_balance_disorder': np.float64,
            'search_trends_beaus_lines': np.float64,
            'search_trends_bells_palsy': np.float64,
            'search_trends_biliary_colic': np.float64,
            'search_trends_binge_eating': np.float64,
            'search_trends_bleeding': np.float64,
            'search_trends_bleeding_on_probing': np.float64,
            'search_trends_blepharospasm': np.float64,
            'search_trends_bloating': np.float64,
            'search_trends_blood_in_stool': np.float64,
            'search_trends_blurred_vision': np.float64,
            'search_trends_blushing': np.float64,
            'search_trends_boil': np.float64,
            'search_trends_bone_fracture': np.float64,
            'search_trends_bone_tumor': np.float64,
            'search_trends_bowel_obstruction': np.float64,
            'search_trends_bradycardia': np.float64,
            'search_trends_braxton_hicks_contractions': np.float64,
            'search_trends_breakthrough_bleeding': np.float64,
            'search_trends_breast_pain': np.float64,
            'search_trends_bronchitis': np.float64,
            'search_trends_bruise': np.float64,
            'search_trends_bruxism': np.float64,
            'search_trends_bunion': np.float64,
            'search_trends_burn': np.float64,
            'search_trends_burning_chest_pain': np.float64,
            'search_trends_burning_mouth_syndrome': np.float64,
            'search_trends_candidiasis': np.float64,
            'search_trends_canker_sore': np.float64,
            'search_trends_cardiac_arrest': np.float64,
            'search_trends_carpal_tunnel_syndrome': np.float64,
            'search_trends_cataplexy': np.float64,
            'search_trends_cataract': np.float64,
            'search_trends_chancre': np.float64,
            'search_trends_cheilitis': np.float64,
            'search_trends_chest_pain': np.float64,
            'search_trends_chills': np.float64,
            'search_trends_chorea': np.float64,
            'search_trends_chronic_pain': np.float64,
            'search_trends_cirrhosis': np.float64,
            'search_trends_cleft_lip_and_cleft_palate': np.float64,
            'search_trends_clouding_of_consciousness': np.float64,
            'search_trends_cluster_headache': np.float64,
            'search_trends_colitis': np.float64,
            'search_trends_coma': np.float64,
            'search_trends_common_cold': np.float64,
            'search_trends_compulsive_behavior': np.float64,
            'search_trends_compulsive_hoarding': np.float64,
            'search_trends_confusion': np.float64,
            'search_trends_congenital_heart_defect': np.float64,
            'search_trends_conjunctivitis': np.float64,
            'search_trends_constipation': np.float64,
            'search_trends_convulsion': np.float64,
            'search_trends_cough': np.float64,
            'search_trends_crackles': np.float64,
            'search_trends_cramp': np.float64,
            'search_trends_crepitus': np.float64,
            'search_trends_croup': np.float64,
            'search_trends_cyanosis': np.float64,
            'search_trends_dandruff': np.float64,
            'search_trends_delayed_onset_muscle_soreness': np.float64,
            'search_trends_dementia': np.float64,
            'search_trends_dentin_hypersensitivity': np.float64,
            'search_trends_depersonalization': np.float64,
            'search_trends_depression': np.float64,
            'search_trends_dermatitis': np.float64,
            'search_trends_desquamation': np.float64,
            'search_trends_developmental_disability': np.float64,
            'search_trends_diabetes': np.float64,
            'search_trends_diabetic_ketoacidosis': np.float64,
            'search_trends_diarrhea': np.float64,
            'search_trends_dizziness': np.float64,
            'search_trends_dry_eye_syndrome': np.float64,
            'search_trends_dysautonomia': np.float64,
            'search_trends_dysgeusia': np.float64,
            'search_trends_dysmenorrhea': np.float64,
            'search_trends_dyspareunia': np.float64,
            'search_trends_dysphagia': np.float64,
            'search_trends_dysphoria': np.float64,
            'search_trends_dystonia': np.float64,
            'search_trends_dysuria': np.float64,
            'search_trends_ear_pain': np.float64,
            'search_trends_eczema': np.float64,
            'search_trends_edema': np.float64,
            'search_trends_encephalitis': np.float64,
            'search_trends_encephalopathy': np.float64,
            'search_trends_epidermoid_cyst': np.float64,
            'search_trends_epilepsy': np.float64,
            'search_trends_epiphora': np.float64,
            'search_trends_erectile_dysfunction': np.float64,
            'search_trends_erythema': np.float64,
            'search_trends_erythema_chronicum_migrans': np.float64,
            'search_trends_esophagitis': np.float64,
            'search_trends_excessive_daytime_sleepiness': np.float64,
            'search_trends_eye_pain': np.float64,
            'search_trends_eye_strain': np.float64,
            'search_trends_facial_nerve_paralysis': np.float64,
            'search_trends_facial_swelling': np.float64,
            'search_trends_fasciculation': np.float64,
            'search_trends_fatigue': np.float64,
            'search_trends_fatty_liver_disease': np.float64,
            'search_trends_fecal_incontinence': np.float64,
            'search_trends_fever': np.float64,
            'search_trends_fibrillation': np.float64,
            'search_trends_fibrocystic_breast_changes': np.float64,
            'search_trends_fibromyalgia': np.float64,
            'search_trends_flatulence': np.float64,
            'search_trends_floater': np.float64,
            'search_trends_focal_seizure': np.float64,
            'search_trends_folate_deficiency': np.float64,
            'search_trends_food_craving': np.float64,
            'search_trends_food_intolerance': np.float64,
            'search_trends_frequent_urination': np.float64,
            'search_trends_gastroesophageal_reflux_disease': np.float64,
            'search_trends_gastroparesis': np.float64,
            'search_trends_generalized_anxiety_disorder': np.float64,
            'search_trends_genital_wart': np.float64,
            'search_trends_gingival_recession': np.float64,
            'search_trends_gingivitis': np.float64,
            'search_trends_globus_pharyngis': np.float64,
            'search_trends_goitre': np.float64,
            'search_trends_gout': np.float64,
            'search_trends_grandiosity': np.float64,
            'search_trends_granuloma': np.float64,
            'search_trends_guilt': np.float64,
            'search_trends_hair_loss': np.float64,
            'search_trends_halitosis': np.float64,
            'search_trends_hay_fever': np.float64,
            'search_trends_headache': np.float64,
            'search_trends_heart_arrhythmia': np.float64,
            'search_trends_heart_murmur': np.float64,
            'search_trends_heartburn': np.float64,
            'search_trends_hematochezia': np.float64,
            'search_trends_hematoma': np.float64,
            'search_trends_hematuria': np.float64,
            'search_trends_hemolysis': np.float64,
            'search_trends_hemoptysis': np.float64,
            'search_trends_hemorrhoids': np.float64,
            'search_trends_hepatic_encephalopathy': np.float64,
            'search_trends_hepatitis': np.float64,
            'search_trends_hepatotoxicity': np.float64,
            'search_trends_hiccup': np.float64,
            'search_trends_hip_pain': np.float64,
            'search_trends_hives': np.float64,
            'search_trends_hot_flash': np.float64,
            'search_trends_hydrocephalus': np.float64,
            'search_trends_hypercalcaemia': np.float64,
            'search_trends_hypercapnia': np.float64,
            'search_trends_hypercholesterolemia': np.float64,
            'search_trends_hyperemesis_gravidarum': np.float64,
            'search_trends_hyperglycemia': np.float64,
            'search_trends_hyperhidrosis': np.float64,
            'search_trends_hyperkalemia': np.float64,
            'search_trends_hyperlipidemia': np.float64,
            'search_trends_hypermobility': np.float64,
            'search_trends_hyperpigmentation': np.float64,
            'search_trends_hypersomnia': np.float64,
            'search_trends_hypertension': np.float64,
            'search_trends_hyperthermia': np.float64,
            'search_trends_hyperthyroidism': np.float64,
            'search_trends_hypertriglyceridemia': np.float64,
            'search_trends_hypertrophy': np.float64,
            'search_trends_hyperventilation': np.float64,
            'search_trends_hypocalcaemia': np.float64,
            'search_trends_hypochondriasis': np.float64,
            'search_trends_hypoglycemia': np.float64,
            'search_trends_hypogonadism': np.float64,
            'search_trends_hypokalemia': np.float64,
            'search_trends_hypomania': np.float64,
            'search_trends_hyponatremia': np.float64,
            'search_trends_hypotension': np.float64,
            'search_trends_hypothyroidism': np.float64,
            'search_trends_hypoxemia': np.float64,
            'search_trends_hypoxia': np.float64,
            'search_trends_impetigo': np.float64,
            'search_trends_implantation_bleeding': np.float64,
            'search_trends_impulsivity': np.float64,
            'search_trends_indigestion': np.float64,
            'search_trends_infection': np.float64,
            'search_trends_inflammation': np.float64,
            'search_trends_inflammatory_bowel_disease': np.float64,
            'search_trends_ingrown_hair': np.float64,
            'search_trends_insomnia': np.float64,
            'search_trends_insulin_resistance': np.float64,
            'search_trends_intermenstrual_bleeding': np.float64,
            'search_trends_intracranial_pressure': np.float64,
            'search_trends_iron_deficiency': np.float64,
            'search_trends_irregular_menstruation': np.float64,
            'search_trends_itch': np.float64,
            'search_trends_jaundice': np.float64,
            'search_trends_kidney_failure': np.float64,
            'search_trends_kidney_stone': np.float64,
            'search_trends_knee_pain': np.float64,
            'search_trends_kyphosis': np.float64,
            'search_trends_lactose_intolerance': np.float64,
            'search_trends_laryngitis': np.float64,
            'search_trends_leg_cramps': np.float64,
            'search_trends_lesion': np.float64,
            'search_trends_leukorrhea': np.float64,
            'search_trends_lightheadedness': np.float64,
            'search_trends_low_back_pain': np.float64,
            'search_trends_low_grade_fever': np.float64,
            'search_trends_lymphedema': np.float64,
            'search_trends_major_depressive_disorder': np.float64,
            'search_trends_malabsorption': np.float64,
            'search_trends_male_infertility': np.float64,
            'search_trends_manic_disorder': np.float64,
            'search_trends_melasma': np.float64,
            'search_trends_melena': np.float64,
            'search_trends_meningitis': np.float64,
            'search_trends_menorrhagia': np.float64,
            'search_trends_middle_back_pain': np.float64,
            'search_trends_migraine': np.float64,
            'search_trends_milium': np.float64,
            'search_trends_mitral_insufficiency': np.float64,
            'search_trends_mood_disorder': np.float64,
            'search_trends_mood_swing': np.float64,
            'search_trends_morning_sickness': np.float64,
            'search_trends_motion_sickness': np.float64,
            'search_trends_mouth_ulcer': np.float64,
            'search_trends_muscle_atrophy': np.float64,
            'search_trends_muscle_weakness': np.float64,
            'search_trends_myalgia': np.float64,
            'search_trends_mydriasis': np.float64,
            'search_trends_myocardial_infarction': np.float64,
            'search_trends_myoclonus': np.float64,
            'search_trends_nasal_congestion': np.float64,
            'search_trends_nasal_polyp': np.float64,
            'search_trends_nausea': np.float64,
            'search_trends_neck_mass': np.float64,
            'search_trends_neck_pain': np.float64,
            'search_trends_neonatal_jaundice': np.float64,
            'search_trends_nerve_injury': np.float64,
            'search_trends_neuralgia': np.float64,
            'search_trends_neutropenia': np.float64,
            'search_trends_night_sweats': np.float64,
            'search_trends_night_terror': np.float64,
            'search_trends_nocturnal_enuresis': np.float64,
            'search_trends_nodule': np.float64,
            'search_trends_nosebleed': np.float64,
            'search_trends_nystagmus': np.float64,
            'search_trends_obesity': np.float64,
            'search_trends_onychorrhexis': np.float64,
            'search_trends_oral_candidiasis': np.float64,
            'search_trends_orthostatic_hypotension': np.float64,
            'search_trends_osteopenia': np.float64,
            'search_trends_osteophyte': np.float64,
            'search_trends_osteoporosis': np.float64,
            'search_trends_otitis': np.float64,
            'search_trends_otitis_externa': np.float64,
            'search_trends_otitis_media': np.float64,
            'search_trends_pain': np.float64,
            'search_trends_palpitations': np.float64,
            'search_trends_pancreatitis': np.float64,
            'search_trends_panic_attack': np.float64,
            'search_trends_papule': np.float64,
            'search_trends_paranoia': np.float64,
            'search_trends_paresthesia': np.float64,
            'search_trends_pelvic_inflammatory_disease': np.float64,
            'search_trends_pericarditis': np.float64,
            'search_trends_periodontal_disease': np.float64,
            'search_trends_periorbital_puffiness': np.float64,
            'search_trends_peripheral_neuropathy': np.float64,
            'search_trends_perspiration': np.float64,
            'search_trends_petechia': np.float64,
            'search_trends_phlegm': np.float64,
            'search_trends_photodermatitis': np.float64,
            'search_trends_photophobia': np.float64,
            'search_trends_photopsia': np.float64,
            'search_trends_pleural_effusion': np.float64,
            'search_trends_pleurisy': np.float64,
            'search_trends_pneumonia': np.float64,
            'search_trends_podalgia': np.float64,
            'search_trends_polycythemia': np.float64,
            'search_trends_polydipsia': np.float64,
            'search_trends_polyneuropathy': np.float64,
            'search_trends_polyuria': np.float64,
            'search_trends_poor_posture': np.float64,
            'search_trends_post_nasal_drip': np.float64,
            'search_trends_postural_orthostatic_tachycardia_syndrome': (
                np.float64
            ),
            'search_trends_prediabetes': np.float64,
            'search_trends_proteinuria': np.float64,
            'search_trends_pruritus_ani': np.float64,
            'search_trends_psychosis': np.float64,
            'search_trends_ptosis': np.float64,
            'search_trends_pulmonary_edema': np.float64,
            'search_trends_pulmonary_hypertension': np.float64,
            'search_trends_purpura': np.float64,
            'search_trends_pus': np.float64,
            'search_trends_pyelonephritis': np.float64,
            'search_trends_radiculopathy': np.float64,
            'search_trends_rectal_pain': np.float64,
            'search_trends_rectal_prolapse': np.float64,
            'search_trends_red_eye': np.float64,
            'search_trends_renal_colic': np.float64,
            'search_trends_restless_legs_syndrome': np.float64,
            'search_trends_rheum': np.float64,
            'search_trends_rhinitis': np.float64,
            'search_trends_rhinorrhea': np.float64,
            'search_trends_rosacea': np.float64,
            'search_trends_round_ligament_pain': np.float64,
            'search_trends_rumination': np.float64,
            'search_trends_scar': np.float64,
            'search_trends_sciatica': np.float64,
            'search_trends_scoliosis': np.float64,
            'search_trends_seborrheic_dermatitis': np.float64,
            'search_trends_self_harm': np.float64,
            'search_trends_sensitivity_to_sound': np.float64,
            'search_trends_sexual_dysfunction': np.float64,
            'search_trends_shallow_breathing': np.float64,
            'search_trends_sharp_pain': np.float64,
            'search_trends_shivering': np.float64,
            'search_trends_shortness_of_breath': np.float64,
            'search_trends_shyness': np.float64,
            'search_trends_sinusitis': np.float64,
            'search_trends_skin_condition': np.float64,
            'search_trends_skin_rash': np.float64,
            'search_trends_skin_tag': np.float64,
            'search_trends_skin_ulcer': np.float64,
            'search_trends_sleep_apnea': np.float64,
            'search_trends_sleep_deprivation': np.float64,
            'search_trends_sleep_disorder': np.float64,
            'search_trends_snoring': np.float64,
            'search_trends_sore_throat': np.float64,
            'search_trends_spasticity': np.float64,
            'search_trends_splenomegaly': np.float64,
            'search_trends_sputum': np.float64,
            'search_trends_stomach_rumble': np.float64,
            'search_trends_strabismus': np.float64,
            'search_trends_stretch_marks': np.float64,
            'search_trends_stridor': np.float64,
            'search_trends_stroke': np.float64,
            'search_trends_stuttering': np.float64,
            'search_trends_subdural_hematoma': np.float64,
            'search_trends_suicidal_ideation': np.float64,
            'search_trends_swelling': np.float64,
            'search_trends_swollen_feet': np.float64,
            'search_trends_swollen_lymph_nodes': np.float64,
            'search_trends_syncope': np.float64,
            'search_trends_tachycardia': np.float64,
            'search_trends_tachypnea': np.float64,
            'search_trends_telangiectasia': np.float64,
            'search_trends_tenderness': np.float64,
            'search_trends_testicular_pain': np.float64,
            'search_trends_throat_irritation': np.float64,
            'search_trends_thrombocytopenia': np.float64,
            'search_trends_thyroid_nodule': np.float64,
            'search_trends_tic': np.float64,
            'search_trends_tinnitus': np.float64,
            'search_trends_tonsillitis': np.float64,
            'search_trends_toothache': np.float64,
            'search_trends_tremor': np.float64,
            'search_trends_trichoptilosis': np.float64,
            'search_trends_tumor': np.float64,
            'search_trends_type_2_diabetes': np.float64,
            'search_trends_unconsciousness': np.float64,
            'search_trends_underweight': np.float64,
            'search_trends_upper_respiratory_tract_infection': np.float64,
            'search_trends_urethritis': np.float64,
            'search_trends_urinary_incontinence': np.float64,
            'search_trends_urinary_tract_infection': np.float64,
            'search_trends_urinary_urgency': np.float64,
            'search_trends_uterine_contraction': np.float64,
            'search_trends_vaginal_bleeding': np.float64,
            'search_trends_vaginal_discharge': np.float64,
            'search_trends_vaginitis': np.float64,
            'search_trends_varicose_veins': np.float64,
            'search_trends_vasculitis': np.float64,
            'search_trends_ventricular_fibrillation': np.float64,
            'search_trends_ventricular_tachycardia': np.float64,
            'search_trends_vertigo': np.float64,
            'search_trends_viral_pneumonia': np.float64,
            'search_trends_visual_acuity': np.float64,
            'search_trends_vomiting': np.float64,
            'search_trends_wart': np.float64,
            'search_trends_water_retention': np.float64,
            'search_trends_weakness': np.float64,
            'search_trends_weight_gain': np.float64,
            'search_trends_wheeze': np.float64,
            'search_trends_xeroderma': np.float64,
            'search_trends_xerostomia': np.float64,
            'search_trends_yawn': np.float64,
            'new_recovered_age_0': np.float64,
            'new_recovered_age_1': np.float64,
            'new_recovered_age_2': np.float64,
            'new_recovered_age_3': np.float64,
            'new_recovered_age_4': np.float64,
            'new_recovered_age_5': np.float64,
            'new_recovered_age_6': np.float64,
            'new_recovered_age_7': np.float64,
            'new_recovered_age_8': np.float64,
            'new_recovered_age_9': np.float64,
            'cumulative_recovered_age_0': np.float64,
            'cumulative_recovered_age_1': np.float64,
            'cumulative_recovered_age_2': np.float64,
            'cumulative_recovered_age_3': np.float64,
            'cumulative_recovered_age_4': np.float64,
            'cumulative_recovered_age_5': np.float64,
            'cumulative_recovered_age_6': np.float64,
            'cumulative_recovered_age_7': np.float64,
            'cumulative_recovered_age_8': np.float64,
            'cumulative_recovered_age_9': np.float64,
            'new_recovered_male': np.float64,
            'new_recovered_female': np.float64,
            'cumulative_recovered_male': np.float64,
            'cumulative_recovered_female': np.float64,
            'locality_code': np.str_,
            'locality_name': np.str_,
            'new_ventilator_patients': np.float64,
            'cumulative_ventilator_patients': np.float64,
            'new_persons_fully_vaccinated_pfizer': np.float64,
            'cumulative_persons_fully_vaccinated_pfizer': np.float64,
            'new_vaccine_doses_administered_pfizer': np.float64,
            'cumulative_vaccine_doses_administered_pfizer': np.float64,
            'new_persons_fully_vaccinated_moderna': np.float64,
            'cumulative_persons_fully_vaccinated_moderna': np.float64,
            'new_vaccine_doses_administered_moderna': np.float64,
            'cumulative_vaccine_doses_administered_moderna': np.float64,
            'new_persons_fully_vaccinated_janssen': np.float64,
            'cumulative_persons_fully_vaccinated_janssen': np.float64,
            'new_vaccine_doses_administered_janssen': np.float64,
            'cumulative_vaccine_doses_administered_janssen': np.float64,
        }),
        supervised_keys=None,
        homepage='https://github.com/GoogleCloudPlatform/covid-19-open-data',
        citation=_CITATION,
    )

  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
    """Returns SplitGenerators."""
    archive_path = dl_manager.download(
        'https://storage.googleapis.com/covid19-open-data/v3/aggregated.csv.gz?generation=1620814656792419'
    )
    return {tfds.Split.TRAIN: self._generate_examples(dl_manager, archive_path)}

  def _generate_examples(
      self, dl_manager: tfds.download.DownloadManager, archive_path
  ):
    """Yields examples."""
    pd = tfds.core.lazy_imports.pandas
    beam = tfds.core.lazy_imports.apache_beam

    def cast_according_to_column(feature_type, v):
      if feature_type == tf.string and isinstance(v, (float, int)):
        return str(v)
      return v

    file_handles = dl_manager.iter_archive(archive_path)
    _, file = next(file_handles)

    columns = pd.read_csv(file, nrows=1).columns

    def load_shard(index: int):
      # There is only one file so by using the for we guarantee that the file
      # will be closed.
      for _, file in dl_manager.iter_archive(archive_path):
        df = pd.read_csv(file, skiprows=index, nrows=_BATCH_SIZE)
        features = self.info.features
        result = []
        for i, row in df.iterrows():
          example = {
              k: cast_according_to_column(features[k].dtype, v)
              for k, v in zip(columns, row.values)
          }
          result.append((index + i, example))
        return result

    return beam.Create(list(range(0, _N_RECORDS, _BATCH_SIZE))) | beam.FlatMap(
        load_shard
    )
