"""
Suppress COVID EHR vaccine concepts.

Original Issues: DC-1692
"""

# Python imports
import logging

# Project imports
from cdr_cleaner.cleaning_rules.deid.concept_suppression import AbstractBqLookupTableConceptSuppression
from constants.cdr_cleaner import clean_cdr as cdr_consts
from common import JINJA_ENV, CDM_TABLES
from utils import pipeline_logging

# Third party imports
from google.cloud.exceptions import GoogleCloudError

LOGGER = logging.getLogger(__name__)

SUPPRESSION_RULE_CONCEPT_TABLE = 'covid_vaccine_concepts'

COVID_VACCINE_CONCEPT_QUERY = JINJA_ENV.from_string("""
CREATE OR REPLACE TABLE `{{project_id}}.{{sandbox_id}}.{{concept_suppression_lookup_table}}` AS
with covid_vacc as (
    SELECT *
    FROM `{{project_id}}.{{dataset_id}}.concept` 
    WHERE (
        -- done by name and vocab --
        REGEXP_CONTAINS(concept_name, r'(?i)(COVID)') AND
        REGEXP_CONTAINS(concept_name, r'(?i)(VAC)') AND 
        vocabulary_id not in ('PPI')
    ) OR (
        -- done by code  and vocab --
        REGEXP_CONTAINS(concept_code, r'(207)|(208)|(210)|(211)|(212)')
        and vocabulary_id = 'CVX'
    ) OR (
        -- done by code and vocab --
        REGEXP_CONTAINS(concept_code, r'(91300)|(91301)|(91302)|(91303)|(91304)')
        and vocabulary_id = 'CPT4'
    )
),
concepts_via_cr as (
select distinct c.*
from `{{project_id}}.{{dataset_id}}.concept`as c
left join `{{project_id}}.{{dataset_id}}.concept_relationship`
on c.concept_id = concept_id_1
where concept_id_2 in (select concept_id from covid_vacc)
# and concept_id_1 not in (select concept_id from covid_vacc)
and (
        relationship_id not in ('Subsumes', 'RxNorm dose form of', 'Dose form group of', 'RxNorm - SPL') OR 
            (relationship_id = 'RxNorm - SPL' and REGEXP_CONTAINS(concept_name, r'(?i)(COVID)'))
    )
),
concepts_via_ca as (
  select c.*
    from `{{project_id}}.{{dataset_id}}.concept`as c
    left join `{{project_id}}.{{dataset_id}}.concept_ancestor` as ca
    on c.concept_id = ca.descendant_concept_id
    where ca.ancestor_concept_id in (select concept_id from covid_vacc)  
)
select distinct * from covid_vacc 
union distinct
select distinct * from concepts_via_ca 
union distinct 
select distinct * from concepts_via_cr
""")


class CovidEHRVaccineConceptSuppression(AbstractBqLookupTableConceptSuppression
                                       ):

    def __init__(self,
                 project_id,
                 dataset_id,
                 sandbox_dataset_id,
                 table_namer=None):
        """
        Initialize the class with proper information.

        Set the issue numbers, description and affected datasets. As other tickets may affect
        this SQL, append them to the list of Jira Issues.
        DO NOT REMOVE ORIGINAL JIRA ISSUE NUMBERS!
        """
        desc = "Suppress COVID EHR vaccine concepts."
        super().__init__(
            issue_numbers=['DC1692'],
            description=desc,
            affected_datasets=[cdr_consts.REGISTERED_TIER_DEID],
            affected_tables=CDM_TABLES,
            project_id=project_id,
            dataset_id=dataset_id,
            sandbox_dataset_id=sandbox_dataset_id,
            concept_suppression_lookup_table=SUPPRESSION_RULE_CONCEPT_TABLE,
            table_namer=table_namer)

    def create_suppression_lookup_table(self, client):
        concept_suppression_lookup_query = COVID_VACCINE_CONCEPT_QUERY.render(
            project_id=self.project_id,
            dataset_id=self.dataset_id,
            sandbox_id=self.sandbox_dataset_id,
            concept_suppression_lookup_table=self.
            concept_suppression_lookup_table)
        query_job = client.query(concept_suppression_lookup_query)
        result = query_job.result()

        if hasattr(result, 'errors') and result.errors:
            LOGGER.error(f"Error running job {result.job_id}: {result.errors}")
            raise GoogleCloudError(
                f"Error running job {result.job_id}: {result.errors}")

    def validate_rule(self, client, *args, **keyword_args):
        """
        Validates the cleaning rule which deletes or updates the data from the tables

        Method to run validation on cleaning rules that will be updating the values.
        For example:
        if your class updates all the datetime fields you should be implementing the
        validation that checks if the date time values that needs to be updated no
        longer exists in the table.

        if your class deletes a subset of rows in the tables you should be implementing
        the validation that checks if the count of final final row counts + deleted rows
        should equals to initial row counts of the affected tables.

        Raises RunTimeError if the validation fails.
        """

        raise NotImplementedError("Please fix me.")

    def setup_validation(self, client, *args, **keyword_args):
        """
        Run required steps for validation setup

        Method to run to setup validation on cleaning rules that will be updating or deleting the values.
        For example:
        if your class updates all the datetime fields you should be implementing the
        logic to get the initial list of values which adhere to a condition we are looking for.

        if your class deletes a subset of rows in the tables you should be implementing
        the logic to get the row counts of the tables prior to applying cleaning rule

        """
        raise NotImplementedError("Please fix me.")


if __name__ == '__main__':
    import cdr_cleaner.args_parser as parser
    import cdr_cleaner.clean_cdr_engine as clean_engine

    ARGS = parser.parse_args()
    pipeline_logging.configure(level=logging.DEBUG, add_console_handler=True)

    if ARGS.list_queries:
        clean_engine.add_console_logging()
        query_list = clean_engine.get_query_list(
            ARGS.project_id, ARGS.dataset_id, ARGS.sandbox_dataset_id,
            [(CovidEHRVaccineConceptSuppression,)])

        for query in query_list:
            LOGGER.info(query)
    else:
        clean_engine.add_console_logging(ARGS.console_log)
        clean_engine.clean_dataset(ARGS.project_id, ARGS.dataset_id,
                                   ARGS.sandbox_dataset_id,
                                   [(CovidEHRVaccineConceptSuppression,)])

[((19, 9, 19, 36), 'logging.getLogger', 'logging.getLogger', ({(19, 27, 19, 35): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((23, 30, 67, 4), 'common.JINJA_ENV.from_string', 'JINJA_ENV.from_string', ({(23, 52, 67, 3): '"""\nCREATE OR REPLACE TABLE `{{project_id}}.{{sandbox_id}}.{{concept_suppression_lookup_table}}` AS\nwith covid_vacc as (\n    SELECT *\n    FROM `{{project_id}}.{{dataset_id}}.concept` \n    WHERE (\n        -- done by name and vocab --\n        REGEXP_CONTAINS(concept_name, r\'(?i)(COVID)\') AND\n        REGEXP_CONTAINS(concept_name, r\'(?i)(VAC)\') AND \n        vocabulary_id not in (\'PPI\')\n    ) OR (\n        -- done by code  and vocab --\n        REGEXP_CONTAINS(concept_code, r\'(207)|(208)|(210)|(211)|(212)\')\n        and vocabulary_id = \'CVX\'\n    ) OR (\n        -- done by code and vocab --\n        REGEXP_CONTAINS(concept_code, r\'(91300)|(91301)|(91302)|(91303)|(91304)\')\n        and vocabulary_id = \'CPT4\'\n    )\n),\nconcepts_via_cr as (\nselect distinct c.*\nfrom `{{project_id}}.{{dataset_id}}.concept`as c\nleft join `{{project_id}}.{{dataset_id}}.concept_relationship`\non c.concept_id = concept_id_1\nwhere concept_id_2 in (select concept_id from covid_vacc)\n# and concept_id_1 not in (select concept_id from covid_vacc)\nand (\n        relationship_id not in (\'Subsumes\', \'RxNorm dose form of\', \'Dose form group of\', \'RxNorm - SPL\') OR \n            (relationship_id = \'RxNorm - SPL\' and REGEXP_CONTAINS(concept_name, r\'(?i)(COVID)\'))\n    )\n),\nconcepts_via_ca as (\n  select c.*\n    from `{{project_id}}.{{dataset_id}}.concept`as c\n    left join `{{project_id}}.{{dataset_id}}.concept_ancestor` as ca\n    on c.concept_id = ca.descendant_concept_id\n    where ca.ancestor_concept_id in (select concept_id from covid_vacc)  \n)\nselect distinct * from covid_vacc \nunion distinct\nselect distinct * from concepts_via_ca \nunion distinct \nselect distinct * from concepts_via_cr\n"""'}, {}), '(\n    """\nCREATE OR REPLACE TABLE `{{project_id}}.{{sandbox_id}}.{{concept_suppression_lookup_table}}` AS\nwith covid_vacc as (\n    SELECT *\n    FROM `{{project_id}}.{{dataset_id}}.concept` \n    WHERE (\n        -- done by name and vocab --\n        REGEXP_CONTAINS(concept_name, r\'(?i)(COVID)\') AND\n        REGEXP_CONTAINS(concept_name, r\'(?i)(VAC)\') AND \n        vocabulary_id not in (\'PPI\')\n    ) OR (\n        -- done by code  and vocab --\n        REGEXP_CONTAINS(concept_code, r\'(207)|(208)|(210)|(211)|(212)\')\n        and vocabulary_id = \'CVX\'\n    ) OR (\n        -- done by code and vocab --\n        REGEXP_CONTAINS(concept_code, r\'(91300)|(91301)|(91302)|(91303)|(91304)\')\n        and vocabulary_id = \'CPT4\'\n    )\n),\nconcepts_via_cr as (\nselect distinct c.*\nfrom `{{project_id}}.{{dataset_id}}.concept`as c\nleft join `{{project_id}}.{{dataset_id}}.concept_relationship`\non c.concept_id = concept_id_1\nwhere concept_id_2 in (select concept_id from covid_vacc)\n# and concept_id_1 not in (select concept_id from covid_vacc)\nand (\n        relationship_id not in (\'Subsumes\', \'RxNorm dose form of\', \'Dose form group of\', \'RxNorm - SPL\') OR \n            (relationship_id = \'RxNorm - SPL\' and REGEXP_CONTAINS(concept_name, r\'(?i)(COVID)\'))\n    )\n),\nconcepts_via_ca as (\n  select c.*\n    from `{{project_id}}.{{dataset_id}}.concept`as c\n    left join `{{project_id}}.{{dataset_id}}.concept_ancestor` as ca\n    on c.concept_id = ca.descendant_concept_id\n    where ca.ancestor_concept_id in (select concept_id from covid_vacc)  \n)\nselect distinct * from covid_vacc \nunion distinct\nselect distinct * from concepts_via_ca \nunion distinct \nselect distinct * from concepts_via_cr\n"""\n    )', False, 'from common import JINJA_ENV, CDM_TABLES\n'), ((151, 11, 151, 30), 'cdr_cleaner.args_parser.parse_args', 'parser.parse_args', ({}, {}), '()', True, 'import cdr_cleaner.args_parser as parser\n'), ((152, 4, 152, 77), 'utils.pipeline_logging.configure', 'pipeline_logging.configure', (), '', False, 'from utils import pipeline_logging\n'), ((155, 8, 155, 42), 'cdr_cleaner.clean_cdr_engine.add_console_logging', 'clean_engine.add_console_logging', ({}, {}), '()', True, 'import cdr_cleaner.clean_cdr_engine as clean_engine\n'), ((156, 21, 158, 51), 'cdr_cleaner.clean_cdr_engine.get_query_list', 'clean_engine.get_query_list', ({(157, 12, 157, 27): 'ARGS.project_id', (157, 29, 157, 44): 'ARGS.dataset_id', (157, 46, 157, 69): 'ARGS.sandbox_dataset_id', (158, 12, 158, 50): '[(CovidEHRVaccineConceptSuppression,)]'}, {}), '(ARGS.project_id, ARGS.dataset_id, ARGS.\n    sandbox_dataset_id, [(CovidEHRVaccineConceptSuppression,)])', True, 'import cdr_cleaner.clean_cdr_engine as clean_engine\n'), ((163, 8, 163, 58), 'cdr_cleaner.clean_cdr_engine.add_console_logging', 'clean_engine.add_console_logging', ({(163, 41, 163, 57): 'ARGS.console_log'}, {}), '(ARGS.console_log)', True, 'import cdr_cleaner.clean_cdr_engine as clean_engine\n'), ((164, 8, 166, 74), 'cdr_cleaner.clean_cdr_engine.clean_dataset', 'clean_engine.clean_dataset', ({(164, 35, 164, 50): 'ARGS.project_id', (164, 52, 164, 67): 'ARGS.dataset_id', (165, 35, 165, 58): 'ARGS.sandbox_dataset_id', (166, 35, 166, 73): '[(CovidEHRVaccineConceptSuppression,)]'}, {}), '(ARGS.project_id, ARGS.dataset_id, ARGS.\n    sandbox_dataset_id, [(CovidEHRVaccineConceptSuppression,)])', True, 'import cdr_cleaner.clean_cdr_engine as clean_engine\n'), ((109, 18, 110, 70), 'google.cloud.exceptions.GoogleCloudError', 'GoogleCloudError', ({(110, 16, 110, 69): 'f"""Error running job {result.job_id}: {result.errors}"""'}, {}), "(f'Error running job {result.job_id}: {result.errors}')", False, 'from google.cloud.exceptions import GoogleCloudError\n')]