import pandas as pd
import os
import pandas_gbq

from scipy.stats import ks_2samp
from biobanking.preprocess.helpers import remove_outliers_iqr
from biobanking.preprocess.aou.mrules import PLAUSIBLE, UNIT_CONVERSIONS, UNIT_DROPS, MEASUREMENT_NAMES_MAP


#--------------- collect numerical measurements info ---------------#
def collect_numerical_measurements_info(min_indv=4000):
    from biobanking.utils.config import AllofUs
    aou = AllofUs()
    numerical_measurements_raw_query = f"""
    WITH eligible_measurements AS (
        SELECT
            m.measurement_concept_id
        FROM `{aou.cdr}.measurement` m
        LEFT JOIN `{aou.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 `{aou.cdr}.measurement` m
    JOIN eligible_measurements em
        ON m.measurement_concept_id = em.measurement_concept_id
    JOIN `{aou.cdr}.concept` c
        ON m.measurement_concept_id = c.concept_id
    LEFT JOIN `{aou.cdr}.concept` uc
        ON m.unit_concept_id = uc.concept_id
    LEFT JOIN `{aou.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 = pandas_gbq.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"{aou.bucket}/data/phenotypes/raw/numerical_measurements.parquet", index=False)
    return

#--------------- Process numerical measurements ---------------#
def clean_measurements_helper(g):
    m = g.name if "measurement" not in g.columns else g["measurement"].iat[0]
    g["median_value"] = pd.to_numeric(g["median_value"], errors="coerce")
    g = g.dropna(subset=["median_value"])
    # 1) drop units (manual)
    drops = UNIT_DROPS.get(m, set())
    if drops:
        g = g[~g["unit"].isin(drops)].copy()
    if g.empty:
        return g
    # 2) convert units (manual)
    conv = UNIT_CONVERSIONS.get(m, {})
    for u, fn in conv.items():
        mask = g["unit"].eq(u)
        if mask.any():
            g.loc[mask, "median_value"] = fn(g.loc[mask, "median_value"].astype(float))
    # 3) plausible range filter
    lo, hi = PLAUSIBLE[m]
    g = g[g["median_value"].between(lo, hi)].copy()
    # 4) IQR outlier removal (pooled; before KS)
    g = remove_outliers_iqr(g, col="median_value", multiplier=5)
    if g.empty:
        return g
    # 5) remove units present in < MIN_UNIT_N samples
    unit_counts = g["unit"].value_counts()
    keep_units = unit_counts[unit_counts >= 5].index
    g = g[g["unit"].isin(keep_units)].copy()
    if g.empty or g["unit"].nunique() == 1:
        return g
    # 6) dissimilar distributions via KS test vs most common unit
    ref_unit = g["unit"].value_counts().idxmax()
    ref = g.loc[g["unit"].eq(ref_unit), "median_value"].astype(float).dropna()

    keep = {ref_unit}
    for u, sub in g.groupby("unit"):
        if u == ref_unit:
            continue
        x = sub["median_value"].astype(float).dropna()
        res = ks_2samp(ref, x, alternative="two-sided", mode="auto")  # res.statistic, res.pvalue
        if (res.pvalue>=0.001) or (res.statistic<=0.25):
            keep.add(u)
    return g[g["unit"].isin(keep)].copy()

def clean_measurements():
    from biobanking.utils.config import AllofUs
    aou = AllofUs()
    numerical_measurements_df = pd.read_parquet(f"{aou.bucket}/data/phenotypes/raw/numerical_measurements.parquet")
    selected_nm_df = numerical_measurements_df.loc[numerical_measurements_df.measurement.isin(PLAUSIBLE.keys())]
    cleaned_nm_df = selected_nm_df.groupby("measurement").apply(clean_measurements_helper, include_groups=False).reset_index(drop=True)
    cleaned_nm_df.to_parquet(f"{aou.bucket}/data/phenotypes/processed/selected_numerical_measurements.parquet")
    return

def save_measurements_in_wide_format():
    from biobanking.utils.config import AllofUs
    aou = AllofUs()
    lab_df = pd.read_parquet(f"{aou.bucket}/data/phenotypes/processed/selected_numerical_measurements.parquet")
    lab_df = lab_df.groupby(["person_id", "measurement"]).agg({"median_value": "median"}).reset_index()
    lab_df = lab_df.pivot_table(index="person_id", columns="measurement", values="median_value")
    lab_df = lab_df.rename(columns=MEASUREMENT_NAMES_MAP)
    lab_df.reset_index().to_csv(f"{aou.bucket}/data/phenotypes/measurements.csv.gz", index=False)
    return
