import polars as pl
import statsmodels.formula.api as smf
import warnings
from statsmodels.tools.sm_exceptions import PerfectSeparationWarning, ConvergenceWarning

# Replicate filters from run_PheWAS.py
warnings.filterwarnings("error", category=RuntimeWarning)
warnings.filterwarnings("ignore", category=ConvergenceWarning)
warnings.filterwarnings("error", category=PerfectSeparationWarning)

# Load data
counts = pl.read_csv("tests/inputs/phecodes_file_10000_samples_binary_predictor_30_phecodes_perfect_sep.txt")
covariates = pl.read_csv("tests/inputs/covariates_file_10000_samples_binary_predictor_30_phecodes_perfect_sep.txt")

# Prepare data for phecode_10
phecode_name = "phecode_10"
cases = counts.filter(pl.col("phecode") == phecode_name).select("id").to_series().to_list()

df = covariates.with_columns(
    pl.when(pl.col("id").is_in(cases))
    .then(1)
    .otherwise(0)
    .alias("phecode_status")
)

model_eq = "phecode_status ~ predictor + age + sex"
df_pandas = df.to_pandas()

print(f"Testing {phecode_name}...")
try:
    result = smf.logit(model_eq, data=df_pandas).fit(disp=0, maxiter=200)
    print(f"Converged: {result.converged}")
    print(f"Params:\n{result.params}")
except Exception as e:
    print(f"Caught exception: {type(e).__name__}: {e}")

# Now test identify_separation_source
def identify_separation_source(df, analysis_str):
    parts = analysis_str.split("~")
    outcome_col = parts[0].strip()
    predictor_cols = [p.strip() for p in parts[1].split("+")]

    for col in predictor_cols:
        if col not in df.columns:
            continue
        dtype = df.schema[col]
        if dtype in [pl.Float32, pl.Float64]:
            continue
        counts = df.group_by(col).agg(pl.col(outcome_col).n_unique().alias("n_outcomes"))
        print(f"Checking {col}: {counts.to_dicts()}")
        if counts.select((pl.col("n_outcomes") == 1).any()).item():
            return col
    return None

sep_col = identify_separation_source(df, model_eq)
print(f"Identified separation source: {sep_col}")
