import pyarrow as pa

from .base import OMOPSchemaBase


class OMOPSchemaV53(OMOPSchemaBase):
    """
    A class to define and manage the schema for OMOP CDM tables.
    The base schema is currently defined for version 5.3.
    """

    def _load_schema(self):
        return {
            "person": {
                "person_id": pa.int64(),
                "gender_concept_id": pa.int64(),
                "year_of_birth": pa.int64(),
                "month_of_birth": pa.int64(),
                "day_of_birth": pa.int64(),
                "birth_datetime": pa.timestamp("us"),
                "race_concept_id": pa.int64(),
                "ethnicity_concept_id": pa.int64(),
                "location_id": pa.int64(),
                "provider_id": pa.int64(),
                "care_site_id": pa.int64(),
                "person_source_value": pa.string(),
                "gender_source_value": pa.string(),
                "gender_source_concept_id": pa.int64(),
                "race_source_value": pa.string(),
                "race_source_concept_id": pa.int64(),
                "ethnicity_source_value": pa.string(),
                "ethnicity_source_concept_id": pa.int64(),
            },
            "observation_period": {
                "observation_period_id": pa.int64(),
                "person_id": pa.int64(),
                "observation_period_start_date": pa.date64(),
                "observation_period_end_date": pa.date64(),
                "period_type_concept_id": pa.int64(),
            },
            "visit_occurrence": {
                "visit_occurrence_id": pa.int64(),
                "person_id": pa.int64(),
                "visit_concept_id": pa.int64(),
                "visit_start_date": pa.date64(),
                "visit_start_datetime": pa.timestamp("us"),
                "visit_end_date": pa.date64(),
                "visit_end_datetime": pa.timestamp("us"),
                "visit_type_concept_id": pa.int64(),
                "provider_id": pa.int64(),
                "care_site_id": pa.int64(),
                "visit_source_value": pa.string(),
                "visit_source_concept_id": pa.int64(),
                "admitting_source_concept_id": pa.int64(),
                "admitting_source_value": pa.string(),
                "discharge_to_concept_id": pa.int64(),
                "discharge_to_source_value": pa.string(),
                "preceding_visit_occurrence_id": pa.int64(),
            },
            "visit_detail": {
                "visit_detail_id": pa.int64(),
                "person_id": pa.int64(),
                "visit_detail_concept_id": pa.int64(),
                "visit_detail_start_date": pa.date64(),
                "visit_detail_start_datetime": pa.timestamp("us"),
                "visit_detail_end_date": pa.date64(),
                "visit_detail_end_datetime": pa.timestamp("us"),
                "visit_detail_type_concept_id": pa.int64(),
                "provider_id": pa.int64(),
                "care_site_id": pa.int64(),
                "visit_detail_source_value": pa.string(),
                "visit_detail_source_concept_id": pa.int64(),
                "admitting_source_value": pa.string(),
                "admitting_source_concept_id": pa.int64(),
                "discharge_to_source_value": pa.string(),
                "discharge_to_concept_id": pa.int64(),
                "preceding_visit_detail_id": pa.int64(),
                "visit_detail_parent_id": pa.int64(),
                "visit_occurrence_id": pa.int64(),
            },
            "condition_occurrence": {
                "condition_occurrence_id": pa.int64(),
                "person_id": pa.int64(),
                "condition_concept_id": pa.int64(),
                "condition_start_date": pa.date64(),
                "condition_start_datetime": pa.timestamp("us"),
                "condition_end_date": pa.date64(),
                "condition_end_datetime": pa.timestamp("us"),
                "condition_type_concept_id": pa.int64(),
                "condition_status_concept_id": pa.int64(),
                "stop_reason": pa.string(),
                "provider_id": pa.int64(),
                "visit_occurrence_id": pa.int64(),
                "visit_detail_id": pa.int64(),
                "condition_source_value": pa.string(),
                "condition_source_concept_id": pa.int64(),
                "condition_status_source_value": pa.string(),
            },
            "drug_exposure": {
                "drug_exposure_id": pa.int64(),
                "person_id": pa.int64(),
                "drug_concept_id": pa.int64(),
                "drug_exposure_start_date": pa.date64(),
                "drug_exposure_start_datetime": pa.timestamp("us"),
                "drug_exposure_end_date": pa.date64(),
                "drug_exposure_end_datetime": pa.timestamp("us"),
                "verbatim_end_date": pa.date64(),
                "drug_type_concept_id": pa.int64(),
                "stop_reason": pa.string(),
                "refills": pa.int64(),
                "quantity": pa.float64(),
                "days_supply": pa.int64(),
                "sig": pa.string(),
                "route_concept_id": pa.int64(),
                "lot_number": pa.string(),
                "provider_id": pa.int64(),
                "visit_occurrence_id": pa.int64(),
                "visit_detail_id": pa.int64(),
                "drug_source_value": pa.string(),
                "drug_source_concept_id": pa.int64(),
                "route_source_value": pa.string(),
                "dose_unit_source_value": pa.string(),
            },
            "procedure_occurrence": {
                "procedure_occurrence_id": pa.int64(),
                "person_id": pa.int64(),
                "procedure_concept_id": pa.int64(),
                "procedure_date": pa.date64(),
                "procedure_datetime": pa.timestamp("us"),
                "procedure_type_concept_id": pa.int64(),
                "modifier_concept_id": pa.int64(),
                "quantity": pa.int64(),
                "provider_id": pa.int64(),
                "visit_occurrence_id": pa.int64(),
                "visit_detail_id": pa.int64(),
                "procedure_source_value": pa.string(),
                "procedure_source_concept_id": pa.int64(),
                "modifier_source_value": pa.string(),
            },
            "device_exposure": {
                "device_exposure_id": pa.int64(),
                "person_id": pa.int64(),
                "device_concept_id": pa.int64(),
                "device_exposure_start_date": pa.date64(),
                "device_exposure_start_datetime": pa.timestamp("us"),
                "device_exposure_end_date": pa.date64(),
                "device_exposure_end_datetime": pa.timestamp("us"),
                "device_type_concept_id": pa.int64(),
                "unique_device_id": pa.string(),
                "quantity": pa.int64(),
                "provider_id": pa.int64(),
                "visit_occurrence_id": pa.int64(),
                "visit_detail_id": pa.int64(),
                "device_source_value": pa.string(),
                "device_source_concept_id": pa.int64(),
            },
            "measurement": {
                "measurement_id": pa.int64(),
                "person_id": pa.int64(),
                "measurement_concept_id": pa.int64(),
                "measurement_date": pa.date64(),
                "measurement_datetime": pa.timestamp("us"),
                "measurement_time": pa.string(),
                "measurement_type_concept_id": pa.int64(),
                "operator_concept_id": pa.int64(),
                "value_as_number": pa.float64(),
                "value_as_concept_id": pa.int64(),
                "unit_concept_id": pa.int64(),
                "range_low": pa.float64(),
                "range_high": pa.float64(),
                "provider_id": pa.int64(),
                "visit_occurrence_id": pa.int64(),
                "visit_detail_id": pa.int64(),
                "measurement_source_value": pa.string(),
                "measurement_source_concept_id": pa.int64(),
                "unit_source_value": pa.string(),
                "value_source_value": pa.string(),
            },
            "observation": {
                "observation_id": pa.int64(),
                "person_id": pa.int64(),
                "observation_concept_id": pa.int64(),
                "observation_date": pa.date64(),
                "observation_datetime": pa.timestamp("us"),
                "observation_type_concept_id": pa.int64(),
                "value_as_number": pa.float64(),
                "value_as_string": pa.string(),
                "value_as_concept_id": pa.int64(),
                "qualifier_concept_id": pa.int64(),
                "unit_concept_id": pa.int64(),
                "provider_id": pa.int64(),
                "visit_occurrence_id": pa.int64(),
                "visit_detail_id": pa.int64(),
                "observation_source_value": pa.string(),
                "observation_source_concept_id": pa.int64(),
                "unit_source_value": pa.string(),
                "qualifier_source_value": pa.string(),
            },
            "death": {
                "person_id": pa.int64(),
                "death_date": pa.date64(),
                "death_datetime": pa.timestamp("us"),
                "death_type_concept_id": pa.int64(),
                "cause_concept_id": pa.int64(),
                "cause_source_value": pa.string(),
                "cause_source_concept_id": pa.int64(),
            },
            "note": {
                "note_id": pa.int64(),
                "person_id": pa.int64(),
                "note_date": pa.date64(),
                "note_datetime": pa.timestamp("us"),
                "note_type_concept_id": pa.int64(),
                "note_class_concept_id": pa.int64(),
                "note_title": pa.string(),
                "note_text": pa.string(),
                "encoding_concept_id": pa.int64(),
                "language_concept_id": pa.int64(),
                "provider_id": pa.int64(),
                "visit_occurrence_id": pa.int64(),
                "visit_detail_id": pa.int64(),
                "note_source_value": pa.string(),
            },
            "note_nlp": {
                "note_nlp_id": pa.int64(),
                "note_id": pa.int64(),
                "section_concept_id": pa.int64(),
                "snippet": pa.string(),
                "offset": pa.string(),
                "lexical_variant": pa.string(),
                "note_nlp_concept_id": pa.int64(),
                "note_nlp_source_concept_id": pa.int64(),
                "nlp_system": pa.string(),
                "nlp_date": pa.date64(),
                "nlp_datetime": pa.timestamp("us"),
                "term_exists": pa.string(),
                "term_temporal": pa.string(),
                "term_modifiers": pa.string(),
            },
            "specimen": {
                "specimen_id": pa.int64(),
                "person_id": pa.int64(),
                "specimen_concept_id": pa.int64(),
                "specimen_type_concept_id": pa.int64(),
                "specimen_date": pa.date64(),
                "specimen_datetime": pa.timestamp("us"),
                "quantity": pa.float64(),
                "unit_concept_id": pa.int64(),
                "anatomic_site_concept_id": pa.int64(),
                "disease_status_concept_id": pa.int64(),
                "specimen_source_id": pa.string(),
                "specimen_source_value": pa.string(),
                "unit_source_value": pa.string(),
                "anatomic_site_source_value": pa.string(),
                "disease_status_source_value": pa.string(),
            },
            "fact_relationship": {
                "domain_concept_id_1": pa.int64(),
                "fact_id_1": pa.int64(),
                "domain_concept_id_2": pa.int64(),
                "fact_id_2": pa.int64(),
                "relationship_concept_id": pa.int64(),
            },
            "location": {
                "location_id": pa.int64(),
                "address_1": pa.string(),
                "address_2": pa.string(),
                "city": pa.string(),
                "state": pa.string(),
                "zip": pa.string(),
                "county": pa.string(),
                "location_source_value": pa.string(),
            },
            "care_site": {
                "care_site_id": pa.int64(),
                "care_site_name": pa.string(),
                "place_of_service_concept_id": pa.int64(),
                "location_id": pa.int64(),
                "care_site_source_value": pa.string(),
                "place_of_service_source_value": pa.string(),
            },
            "provider": {
                "provider_id": pa.int64(),
                "provider_name": pa.string(),
                "npi": pa.string(),
                "dea": pa.string(),
                "specialty_concept_id": pa.int64(),
                "care_site_id": pa.int64(),
                "year_of_birth": pa.int64(),
                "gender_concept_id": pa.int64(),
                "provider_source_value": pa.string(),
                "specialty_source_value": pa.string(),
                "specialty_source_concept_id": pa.int64(),
                "gender_source_value": pa.string(),
                "gender_source_concept_id": pa.int64(),
            },
            "payer_plan_period": {
                "payer_plan_period_id": pa.int64(),
                "person_id": pa.int64(),
                "payer_plan_period_start_date": pa.date64(),
                "payer_plan_period_end_date": pa.date64(),
                "payer_concept_id": pa.int64(),
                "payer_source_value": pa.string(),
                "payer_source_concept_id": pa.int64(),
                "plan_concept_id": pa.int64(),
                "plan_source_value": pa.string(),
                "plan_source_concept_id": pa.int64(),
                "sponsor_concept_id": pa.int64(),
                "sponsor_source_value": pa.string(),
                "sponsor_source_concept_id": pa.int64(),
                "family_source_value": pa.string(),
                "stop_reason_concept_id": pa.int64(),
                "stop_reason_source_value": pa.string(),
                "stop_reason_source_concept_id": pa.int64(),
            },
            "cost": {
                "cost_id": pa.int64(),
                "cost_event_id": pa.int64(),
                "cost_domain_id": pa.string(),
                "cost_type_concept_id": pa.int64(),
                "currency_concept_id": pa.int64(),
                "total_charge": pa.float64(),
                "total_cost": pa.float64(),
                "total_paid": pa.float64(),
                "paid_by_payer": pa.float64(),
                "paid_by_patient": pa.float64(),
                "paid_patient_copay": pa.float64(),
                "paid_patient_coinsurance": pa.float64(),
                "paid_patient_deductible": pa.float64(),
                "paid_by_primary": pa.float64(),
                "paid_ingredient_cost": pa.float64(),
                "paid_dispensing_fee": pa.float64(),
                "payer_plan_period_id": pa.int64(),
                "amount_allowed": pa.float64(),
                "revenue_code_concept_id": pa.int64(),
                "revenue_code_source_value": pa.string(),
                "drg_concept_id": pa.int64(),
                "drg_source_value": pa.string(),
            },
            "drug_era": {
                "drug_era_id": pa.int64(),
                "person_id": pa.int64(),
                "drug_concept_id": pa.int64(),
                "drug_era_start_date": pa.date64(),
                "drug_era_end_date": pa.date64(),
                "drug_exposure_count": pa.int64(),
                "gap_days": pa.int64(),
            },
            "dose_era": {
                "dose_era_id": pa.int64(),
                "person_id": pa.int64(),
                "drug_concept_id": pa.int64(),
                "unit_concept_id": pa.int64(),
                "dose_value": pa.float64(),
                "dose_era_start_date": pa.date64(),
                "dose_era_end_date": pa.date64(),
            },
            "condition_era": {
                "condition_era_id": pa.int64(),
                "person_id": pa.int64(),
                "condition_concept_id": pa.int64(),
                "condition_era_start_date": pa.date64(),
                "condition_era_end_date": pa.date64(),
                "condition_occurrence_count": pa.int64(),
            },
            "metadata": {
                "metadata_concept_id": pa.int64(),
                "metadata_type_concept_id": pa.int64(),
                "name": pa.string(),
                "value_as_string": pa.string(),
                "value_as_concept_id": pa.int64(),
                "metadata_date": pa.date64(),
                "metadata_datetime": pa.timestamp("us"),
            },
            "cdm_source": {
                "cdm_source_name": pa.string(),
                "cdm_source_abbreviation": pa.string(),
                "cdm_holder": pa.string(),
                "source_description": pa.string(),
                "source_documentation_reference": pa.string(),
                "cdm_etl_reference": pa.string(),
                "source_release_date": pa.date64(),
                "cdm_release_date": pa.date64(),
                "cdm_version": pa.string(),
                "vocabulary_version": pa.string(),
            },
            "concept": {
                "concept_id": pa.int64(),
                "concept_name": pa.string(),
                "domain_id": pa.string(),
                "vocabulary_id": pa.string(),
                "concept_class_id": pa.string(),
                "standard_concept": pa.string(),
                "concept_code": pa.string(),
                "valid_start_date": pa.date64(),
                "valid_end_date": pa.date64(),
                "invalid_reason": pa.string(),
            },
            "vocabulary": {
                "vocabulary_id": pa.string(),
                "vocabulary_name": pa.string(),
                "vocabulary_reference": pa.string(),
                "vocabulary_version": pa.string(),
                "vocabulary_concept_id": pa.int64(),
            },
            "domain": {
                "domain_id": pa.string(),
                "domain_name": pa.string(),
                "domain_concept_id": pa.int64(),
            },
            "concept_class": {
                "concept_class_id": pa.string(),
                "concept_class_name": pa.string(),
                "concept_class_concept_id": pa.int64(),
            },
            "concept_relationship": {
                "concept_id_1": pa.int64(),
                "concept_id_2": pa.int64(),
                "relationship_id": pa.string(),
                "valid_start_date": pa.date64(),
                "valid_end_date": pa.date64(),
                "invalid_reason": pa.string(),
            },
            "relationship": {
                "relationship_id": pa.string(),
                "relationship_name": pa.string(),
                "is_hierarchical": pa.string(),
                "defines_ancestry": pa.string(),
                "reverse_relationship_id": pa.string(),
                "relationship_concept_id": pa.int64(),
            },
            "concept_synonym": {
                "concept_id": pa.int64(),
                "concept_synonym_name": pa.string(),
                "language_concept_id": pa.int64(),
            },
            "concept_ancestor": {
                "ancestor_concept_id": pa.int64(),
                "descendant_concept_id": pa.int64(),
                "min_levels_of_separation": pa.int64(),
                "max_levels_of_separation": pa.int64(),
            },
            "source_to_concept_map": {
                "source_code": pa.string(),
                "source_concept_id": pa.int64(),
                "source_vocabulary_id": pa.string(),
                "source_code_description": pa.string(),
                "target_concept_id": pa.int64(),
                "target_vocabulary_id": pa.string(),
                "valid_start_date": pa.date64(),
                "valid_end_date": pa.date64(),
                "invalid_reason": pa.string(),
            },
            "drug_strength": {
                "drug_concept_id": pa.int64(),
                "ingredient_concept_id": pa.int64(),
                "amount_value": pa.float64(),
                "amount_unit_concept_id": pa.int64(),
                "numerator_value": pa.float64(),
                "numerator_unit_concept_id": pa.int64(),
                "denominator_value": pa.float64(),
                "denominator_unit_concept_id": pa.int64(),
                "box_size": pa.int64(),
                "valid_start_date": pa.date64(),
                "valid_end_date": pa.date64(),
                "invalid_reason": pa.string(),
            },
            "cohort_definition": {
                "cohort_definition_id": pa.int64(),
                "cohort_definition_name": pa.string(),
                "cohort_definition_description": pa.string(),
                "definition_type_concept_id": pa.int64(),
                "cohort_definition_syntax": pa.string(),
                "subject_concept_id": pa.int64(),
                "cohort_initiation_date": pa.date64(),
            },
            "attribute_definition": {
                "attribute_definition_id": pa.int64(),
                "attribute_name": pa.string(),
                "attribute_description": pa.string(),
                "attribute_type_concept_id": pa.int64(),
                "attribute_syntax": pa.string(),
            },
        }
