import os
import pandas as pd
import pandas_gbq

# person info
def add_sample_info(df, google_project, 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
    genetic_pca_df = ancestry_df.pca_features.str.strip("[]").str.split(", ", expand=True).astype(float)
    ancestry_df[[f"genetic_pca{i}" for i in range(1, genetic_pca_df.shape[1]+1)]] = genetic_pca_df
    ancestry_df = ancestry_df.loc[:, ["person_id", "ancestry_pred_other"]+[f"genetic_pca{i}" for i in range(1, genetic_pca_df.shape[1]+1)]]
    df = df.merge(ancestry_df, on="person_id", how="left")
    return df

def collect_person_info():
    from biobanking.utils.config import AllofUs
    aou = AllofUs()
    person_table_query = f"""
    SELECT 
        person_id,
        gender,
        sex_at_birth,
        race,
        ethnicity,
        dob,
        age_at_consent,
        age_at_cdr,
        has_ehr_data
    FROM `{aou.cdr}.cb_search_person`
    WHERE has_whole_genome_variant = 1
    """
    person_df = pandas_gbq.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, aou.google_project)

    person_df.to_csv(f"{aou.bucket}/data/phenotypes/raw/individuals.csv.gz", index=False)
    return

def clean_and_save_person_info():
    from biobanking.utils.config import AllofUs
    aou = AllofUs()
    person_df = pd.read_csv(f"{aou.bucket}/data/phenotypes/raw/individuals.csv.gz")
    person_df = person_df.rename(columns={"age_at_cdr": "age"})
    person_df["sex_at_birth"] = person_df.sex_at_birth.apply(lambda x: 1 if x=="Female" else (0 if x=="Male" else pd.NA))
    person_df.to_csv(f"{aou.bucket}/data/phenotypes/individuals.csv.gz", index=False)
    return
