import pandas as pd
import os
from tqdm import tqdm
import json
import numpy as np
from scipy import stats
import multiprocessing as mp
from scipy.stats import ks_2samp

def remove_outliers_iqr(df, column, multiplier=5):
    """
    Remove outliers from a specified column in a pandas DataFrame using the IQR method.
    """
    value = np.log1p(df[column].astype(float))
    Q1 = value.quantile(0.25)
    Q3 = value.quantile(0.75)
    IQR = Q3 - Q1
    lower_bound = Q1 - multiplier * IQR
    upper_bound = Q3 + multiplier * IQR
    return np.where(value.between(lower_bound, upper_bound), df[column], pd.NA)

def remove_outliers(df, col="median_value", multiplier=5):
    df =  df.copy()
    df = df[df[col] >= 0]
    df[col] = remove_outliers_iqr(df, col, multiplier)
    return df.dropna(subset=[col])

def clean_measurements_helper(g):
    from rwe.parsers.aou.config import PLAUSIBLE, UNIT_DROPS, UNIT_CONVERSIONS
    m = 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(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 rwe.parsers.aou.config import BUCKET, CDR, GOOGLE_PROJECT, PLAUSIBLE
    numerical_measurements_df = pd.read_parquet(f"{BUCKET}/data/rwe_info/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).reset_index(drop=True)
    cleaned_nm_df.to_parquet(f"{BUCKET}/data/rwe_info/processed/selected_numerical_measurements.parquet")
    return

def clean_surveys():
    from rwe.parsers.aou.config import BUCKET, CDR, GOOGLE_PROJECT
    survey_df = pd.read_parquet(f"{BUCKET}/data/rwe_info/raw/survey_health.parquet")

    questions = [
    'Overall Health: General Health',
    'Overall Health: General Mental Health',
    'Overall Health: General Physical Health',
    'Overall Health: General Quality',
    ]

    selected_survey_df = survey_df.loc[survey_df.question.isin(questions)].copy()
    selected_survey_df.to_parquet(f"{BUCKET}/data/rwe_info/processed/selected_surveys.parquet")
    return
