import pyarrow as pa

from .base import OMOPSchemaBase


class OMOPSchemaV5(OMOPSchemaBase):
    """
    A class to define and manage the schema for OMOP CDM tables for version 5.0.
    """

    def _load_schema(self):
        return {
            "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_unit_concept_id": pa.int64(),
                "valid_start_date": pa.date64(),
                "valid_end_date": pa.date64(),
                "invalid_reason": pa.string(),
            },
            "attribute_definition": {
                "attribute_definition_id": pa.int64(),
                "attribute_name": pa.string(),
                "attribute_description": pa.large_string(),
                "attribute_type_concept_id": pa.int64(),
                "attribute_syntax": pa.large_string(),
            },
            "cdm_source": {
                "cdm_source_name": pa.string(),
                "cdm_source_abbreviation": pa.string(),
                "cdm_holder": pa.string(),
                "source_description": pa.large_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(),
            },
            "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_source_value": pa.string(),
                "plan_source_value": pa.string(),
                "family_source_value": pa.string(),
            },
            "visit_cost": {
                "visit_cost_id": pa.int64(),
                "visit_occurrence_id": pa.int64(),
                "currency_concept_id": pa.int64(),
                "paid_copay": pa.float64(),
                "paid_coinsurance": pa.float64(),
                "paid_toward_deductible": pa.float64(),
                "paid_by_payer": pa.float64(),
                "paid_by_coordination_benefits": pa.float64(),
                "total_out_of_pocket": pa.float64(),
                "total_paid": pa.float64(),
                "payer_plan_period_id": pa.int64(),
            },
            "procedure_cost": {
                "procedure_cost_id": pa.int64(),
                "procedure_occurrence_id": pa.int64(),
                "currency_concept_id": pa.int64(),
                "paid_copay": pa.float64(),
                "paid_coinsurance": pa.float64(),
                "paid_toward_deductible": pa.float64(),
                "paid_by_payer": pa.float64(),
                "paid_by_coordination_benefits": pa.float64(),
                "total_out_of_pocket": pa.float64(),
                "total_paid": pa.float64(),
                "revenue_code_concept_id": pa.int64(),
                "payer_plan_period_id": pa.int64(),
                "revenue_code_source_value": pa.string(),
            },
            "drug_cost": {
                "drug_cost_id": pa.int64(),
                "drug_exposure_id": pa.int64(),
                "currency_concept_id": pa.int64(),
                "paid_copay": pa.float64(),
                "paid_coinsurance": pa.float64(),
                "paid_toward_deductible": pa.float64(),
                "paid_by_payer": pa.float64(),
                "paid_by_coordination_benefits": pa.float64(),
                "total_out_of_pocket": pa.float64(),
                "total_paid": pa.float64(),
                "ingredient_cost": pa.float64(),
                "dispensing_fee": pa.float64(),
                "average_wholesale_price": pa.float64(),
                "payer_plan_period_id": pa.int64(),
            },
            "device_cost": {
                "device_cost_id": pa.int64(),
                "device_exposure_id": pa.int64(),
                "currency_concept_id": pa.int64(),
                "paid_copay": pa.float64(),
                "paid_coinsurance": pa.float64(),
                "paid_toward_deductible": pa.float64(),
                "paid_by_payer": pa.float64(),
                "paid_by_coordination_benefits": pa.float64(),
                "total_out_of_pocket": pa.float64(),
                "total_paid": pa.float64(),
                "payer_plan_period_id": pa.int64(),
            },
            "cohort": {
                "cohort_definition_id": pa.int64(),
                "subject_id": pa.int64(),
                "cohort_start_date": pa.date64(),
                "cohort_end_date": pa.date64(),
            },
            "cohort_attribute": {
                "cohort_definition_id": pa.int64(),
                "cohort_start_date": pa.date64(),
                "cohort_end_date": pa.date64(),
                "subject_id": pa.int64(),
                "attribute_definition_id": pa.int64(),
                "value_as_number": pa.float64(),
                "value_as_concept_id": pa.int64(),
            },
            "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(),
            },
        }
