import pandas as pd
import os
from tqdm import tqdm

# person info
def add_sample_info(df, auxiliary_path = "gs://fc-aou-datasets-controlled/v8/wgs/short_read/snpindel/aux"):
    # read AoU provided central qc and ancestry files
    related_samples_path = f"{auxiliary_path}/relatedness/relatedness.tsv"
    flagged_samples_path = f"{auxiliary_path}/qc/flagged_samples.tsv"
    ancestry_path = f"{auxiliary_path}/ancestry/ancestry_preds.tsv"
    rel_df = pd.read_csv(related_samples_path, sep="\t", storage_options={"project": GOOGLE_PROJECT, "requester_pays": True})
    qc_sample_df = pd.read_csv(flagged_samples_path, sep="\t", storage_options={"project": GOOGLE_PROJECT, "requester_pays": True})
    ancestry_df = pd.read_csv(ancestry_path, sep="\t", storage_options={"project": GOOGLE_PROJECT, "requester_pays": True}).rename(columns={"research_id": "person_id"})
    # flag duplicates in pheno file
    df["duplicates"] = df.person_id.isin(rel_df.loc[rel_df.kin>0.354, "i.s"].values)
    # flag samples that failed qc
    df["qc_failed"] = df.person_id.isin(qc_sample_df.s)
    # add ancestry predictions and genetic pcs to pheno file
    ancestry_df = ancestry_df.loc[:, ["person_id", "ancestry_pred_other"]]
    df = df.merge(ancestry_df, on="person_id", how="left")
    return df

def save_person_info(cdr, bucket):
    person_table_query = f"""
    SELECT 
        person_id,
        gender,
        sex_at_birth,
        race,
        ethnicity,
        dob,
        age_at_consent,
        age_at_cdr,
        has_ehr_data
    FROM {cdr}.cb_search_person
    WHERE has_whole_genome_variant = 1
    """
    person_df = pd.read_gbq(
        person_table_query, dialect="standard",
        use_bqstorage_api=("BIGQUERY_STORAGE_API_ENABLED" in os.environ),
        progress_bar_type="tqdm_notebook"
    )    
    person_df = add_sample_info(person_df)
    person_df.to_csv(f"{bucket}/data/rwe_info/raw/person.csv.gz", index=False)
    return

# save numerical measurements info
def save_numerical_measurements_info(cdr, bucket, min_indv=4000):
    numerical_measurements_raw_query = f"""
    WITH eligible_measurements AS (
    SELECT
        m.measurement_concept_id
    FROM `{cdr}.measurement` m
    LEFT JOIN `{cdr}.cb_search_person` p
        ON p.person_id = m.person_id
    WHERE p.has_whole_genome_variant = 1
        AND m.value_as_number IS NOT NULL
    GROUP BY m.measurement_concept_id
    HAVING COUNT(DISTINCT m.person_id) >= {min_indv}
    ),

    numeric_raw_table AS (
    SELECT
    m.person_id,
    m.measurement_concept_id,
    c.concept_name AS measurement,
    m.measurement_datetime,
    m.value_as_number,
    uc.concept_name AS unit,
    m.unit_concept_id,
    m.unit_source_value
    FROM `{cdr}.measurement` m
    JOIN eligible_measurements em
    ON m.measurement_concept_id = em.measurement_concept_id
    JOIN `{cdr}.concept` c
    ON m.measurement_concept_id = c.concept_id
    LEFT JOIN `{cdr}.concept` uc
    ON m.unit_concept_id = uc.concept_id
    LEFT JOIN `{cdr}.cb_search_person` p
    ON p.person_id = m.person_id
    WHERE p.has_whole_genome_variant = 1
    AND m.value_as_number IS NOT NULL
    )

    SELECT
    person_id,
    measurement_concept_id,
    ANY_VALUE(measurement) AS measurement,
    unit_concept_id,
    ANY_VALUE(unit) AS unit,
    ANY_VALUE(unit_source_value) AS unit_source_value,
    APPROX_QUANTILES(value_as_number, 2)[OFFSET(1)] AS median_value
    FROM numeric_raw_table
    GROUP BY
    person_id,
    measurement_concept_id,
    unit_concept_id
    """

    numerical_measurements_df = pd.read_gbq(
        numerical_measurements_raw_query, dialect="standard",
        use_bqstorage_api=("BIGQUERY_STORAGE_API_ENABLED" in os.environ),
        progress_bar_type="tqdm_notebook"
    )
    numerical_measurements_df.to_parquet(f"{bucket}/data/rwe_info/raw/numerical_measurements.parquet", index=False)
    return

# save categorical measurements info
def save_categorical_measurements_info(cdr, bucket, min_indv=4000):
    categorical_measurements_query = f"""
    WITH eligible_measurements AS (
    SELECT
        m.measurement_concept_id
    FROM `{cdr}.measurement` m
    LEFT JOIN `{cdr}.cb_search_person` p
        ON p.person_id = m.person_id
    WHERE p.has_whole_genome_variant = 1
        AND m.value_as_number IS NULL
        AND m.value_as_concept_id IS NOT NULL
    GROUP BY m.measurement_concept_id
    HAVING COUNT(DISTINCT m.person_id) >= {min_indv}
    ),

    categorical_raw_table AS (
    SELECT
        m.person_id,
        m.measurement_concept_id,
        c.concept_name AS measurement,
        m.measurement_datetime,
        m.value_as_concept_id,
        vc.concept_name AS value_category
    FROM `{cdr}.measurement` m
    JOIN eligible_measurements em
        ON m.measurement_concept_id = em.measurement_concept_id
    JOIN `{cdr}.concept` c
        ON m.measurement_concept_id = c.concept_id
    LEFT JOIN `{cdr}.concept` vc
        ON m.value_as_concept_id = vc.concept_id
    LEFT JOIN `{cdr}.cb_search_person` p
        ON p.person_id = m.person_id
    WHERE p.has_whole_genome_variant = 1
        AND m.value_as_number IS NULL
        AND m.value_as_concept_id IS NOT NULL
        AND vc.concept_name NOT IN ('No matching concept')
    ),

    consistency AS (
    -- compute counts and the (unique) categorical value per person/measurement
    SELECT
        person_id,
        measurement_concept_id,
        COUNT(*) AS n_obs,
        COUNT(DISTINCT value_as_concept_id) AS n_distinct_values,
        MIN(value_as_concept_id) AS value_as_concept_id
    FROM categorical_raw_table
    GROUP BY person_id, measurement_concept_id
    )

    SELECT
    cs.person_id,
    cs.measurement_concept_id,
    mc.concept_name AS measurement,
    cs.value_as_concept_id,
    vc.concept_name AS value_category,
    cs.n_obs
    FROM consistency cs
    JOIN `{cdr}.concept` mc
    ON cs.measurement_concept_id = mc.concept_id
    LEFT JOIN `{cdr}.concept` vc
    ON cs.value_as_concept_id = vc.concept_id
    WHERE cs.n_distinct_values = 1
    AND cs.n_obs >= 1  -- change to ">= 1" if you want to keep single-read cases
    """

    categorical_measurements_df = pd.read_gbq(
        categorical_measurements_query, dialect="standard",
        use_bqstorage_api=("BIGQUERY_STORAGE_API_ENABLED" in os.environ),
        progress_bar_type="tqdm_notebook"
    )
    categorical_measurements_df.to_parquet(f"{bucket}/data/rwe_info/raw/categorical_measurements.parquet", index=False)
    return    

# save survey info
def save_survey_info(cdr, bucket):
    ds_survey_query = f"""
    SELECT
    s.answer_concept_id,
    s.person_id,
    s.survey_datetime,
    s.survey,
    s.question,
    s.answer,
    ac.concept_name AS answer_category,
    s.survey_version_name
    
    FROM `{cdr}.ds_survey` s
    LEFT JOIN `{cdr}.concept` ac
    ON s.answer_concept_id = ac.concept_id
    LEFT JOIN `{cdr}.cb_search_person` p
    ON p.person_id = s.person_id
    WHERE p.has_whole_genome_variant = 1
        AND s.survey IN ('Overall Health', 'Personal and Family Health History')
    """
    survey_df = pd.read_gbq(
        ds_survey_query, dialect="standard",
        use_bqstorage_api=("BIGQUERY_STORAGE_API_ENABLED" in os.environ),
        progress_bar_type="tqdm_notebook"
    )
    survey_df.to_parquet(f"{bucket}/data/rwe_info/raw/survey_health.parquet", index=False)
    return

# save observation info
def save_observation_info(cdr, bucket):
    observation_table_query = f"""
    SELECT
    o.person_id,
    c.concept_name AS observation,
    o.observation_datetime,
    o.value_as_number,
    vc.concept_name AS value_category,
    uc.concept_name AS unit,
    o.value_as_string,
    o.value_source_value,
    o.observation_source_value,
    tc.concept_name AS observation_type
    FROM `{cdr}.observation` o
    JOIN `{cdr}.concept` c
    ON o.observation_concept_id = c.concept_id
    LEFT JOIN `{cdr}.concept` vc
    ON o.value_as_concept_id = vc.concept_id
    LEFT JOIN `{cdr}.concept` uc
    ON o.unit_concept_id = uc.concept_id
    LEFT JOIN `{cdr}.concept` tc
    ON o.observation_type_concept_id = tc.concept_id
    LEFT JOIN `{cdr}.cb_search_person` p
    ON p.person_id = o.person_id
    WHERE p.has_whole_genome_variant = 1
    """

    observation_df = pd.read_gbq(
        observation_table_query, dialect="standard",
        use_bqstorage_api=("BIGQUERY_STORAGE_API_ENABLED" in os.environ),
        progress_bar_type="tqdm_notebook"
    )
    observation_df.to_csv(f"{bucket}/data/rwe_info/raw/observation.csv.gz")
    return

# save condition info
def save_condition_info(cdr, bucket):
    condition_table_query = f"""
    SELECT
    co.person_id,
    c.concept_name AS condition,
    co.condition_start_datetime,
    co.condition_end_datetime,
    tc.concept_name AS condition_type,
    sc.concept_name AS condition_status,
    co.condition_source_value,
    co.condition_status_source_value,
    co.visit_occurrence_id
    FROM `{cdr}.condition_occurrence` co
    JOIN `{cdr}.concept` c
    ON co.condition_concept_id = c.concept_id
    LEFT JOIN `{cdr}.concept` tc
    ON co.condition_type_concept_id = tc.concept_id
    LEFT JOIN `{cdr}.concept` sc
    ON co.condition_status_concept_id = sc.concept_id
    LEFT JOIN `{cdr}.cb_search_person` p
    ON p.person_id = co.person_id
    WHERE p.has_whole_genome_variant = 1
    """

    condition_df = pd.read_gbq(
        condition_table_query, dialect="standard",
        use_bqstorage_api=("BIGQUERY_STORAGE_API_ENABLED" in os.environ),
        progress_bar_type="tqdm_notebook"
    )
    condition_df.to_csv(f"{bucket}/data/rwe_info/raw/condition.csv.gz")
    return

# save procedure info
def save_procedure_info(cdr, bucket):
    procedure_table_query = f"""
    SELECT
    po.person_id,
    c.concept_name AS procedure,
    po.procedure_datetime,
    tc.concept_name AS procedure_type,
    po.procedure_source_value,
    po.visit_occurrence_id
    FROM `{cdr}.procedure_occurrence` po
    JOIN `{cdr}.concept` c
    ON po.procedure_concept_id = c.concept_id
    LEFT JOIN `{cdr}.concept` tc
    ON po.procedure_type_concept_id = tc.concept_id
    LEFT JOIN `{cdr}.cb_search_person` p
    ON p.person_id = po.person_id
    WHERE p.has_whole_genome_variant = 1
    """

    procedure_df = pd.read_gbq(
        procedure_table_query, dialect="standard",
        use_bqstorage_api=("BIGQUERY_STORAGE_API_ENABLED" in os.environ),
        progress_bar_type="tqdm_notebook"
    )
    procedure_df.to_csv(f"{bucket}/data/rwe_info/raw/procedure.csv.gz")
    return

# save visit occurrence info
def save_visit_occurrence_info(cdr, bucket):
    visit_table_query = f"""
    SELECT
    vo.person_id,
    c.concept_name AS visit_type,
    vo.visit_start_datetime,
    vo.visit_end_datetime,
    ac.concept_name AS admitting_source,
    dc.concept_name AS discharge_to,
    vo.visit_source_value
    FROM `{cdr}.visit_occurrence` vo
    LEFT JOIN `{cdr}.concept` c
    ON vo.visit_concept_id = c.concept_id
    LEFT JOIN `{cdr}.concept` ac
    ON vo.admitting_source_concept_id = ac.concept_id
    LEFT JOIN `{cdr}.concept` dc
    ON vo.discharge_to_concept_id = dc.concept_id
    LEFT JOIN `{cdr}.cb_search_person` p
    ON p.person_id = vo.person_id
    WHERE p.has_whole_genome_variant = 1
    """

    visit_df = pd.read_gbq(
        visit_table_query, dialect="standard",
        use_bqstorage_api=("BIGQUERY_STORAGE_API_ENABLED" in os.environ),
        progress_bar_type="tqdm_notebook"
    )
    visit_df.to_csv(f"{bucket}/data/rwe_info/raw/visit.csv.gz")
    return

# save drug exposure info
def save_drug_exposure_info(cdr, bucket):
    drug_table_query = f"""
    SELECT
    de.person_id,
    c.concept_name AS drug,
    de.drug_exposure_start_datetime,
    de.drug_exposure_end_datetime,
    rc.concept_name AS route,
    de.quantity,
    de.days_supply,
    tc.concept_name AS drug_type,
    de.drug_source_value,
    de.visit_occurrence_id
    FROM `{cdr}.drug_exposure` de
    JOIN `{cdr}.concept` c
    ON de.drug_concept_id = c.concept_id
    LEFT JOIN `{cdr}.concept` rc
    ON de.route_concept_id = rc.concept_id
    LEFT JOIN `{cdr}.concept` tc
    ON de.drug_type_concept_id = tc.concept_id
    LEFT JOIN `{cdr}.cb_search_person` p
    ON p.person_id = de.person_id
    WHERE p.has_whole_genome_variant = 1
    """

    drug_df = pd.read_gbq(
        drug_table_query, dialect="standard",
        use_bqstorage_api=("BIGQUERY_STORAGE_API_ENABLED" in os.environ),
        progress_bar_type="tqdm_notebook"
    )
    drug_df.to_csv(f"{bucket}/data/rwe_info/raw/drug.csv.gz")
    return

def save_rwe_info(
    person=True, numerical_measurements=True, categorical_measurements=True,
    survey=True, observation=True, condition=True,
    procedure=True, visit_occurrence=True, drug_exposure=True, 
    min_indv=4000
    ):
    from rwe.parsers.aou.config import BUCKET, CDR, GOOGLE_PROJECT
    if person:
        save_person_info(cdr=CDR, bucket=BUCKET)
    if numerical_measurements:
        save_numerical_measurements_info(cdr=CDR, bucket=BUCKET, min_indv=min_indv)
    if categorical_measurements:
        save_categorical_measurements_info(cdr=CDR, bucket=BUCKET, min_indv=min_indv)
    if survey:
        save_survey_info(cdr=CDR, bucket=BUCKET)
    if observation:
        save_observation_info(cdr=CDR, bucket=BUCKET)
    if condition:
        save_condition_info(cdr=CDR, bucket=BUCKET)
    if procedure:
        save_procedure_info(cdr=CDR, bucket=BUCKET)
    if visit_occurrence:
        save_visit_occurrence_info(cdr=CDR, bucket=BUCKET)
    if drug_exposure:
        save_drug_exposure_info(cdr=CDR, bucket=BUCKET)
    return
