import polars as pl

def check_perfect_separation(
    df: pl.DataFrame,
    outcome_col: str,
    predictor_cols: list[str]
) -> str | None:
    """
    Checks which predictor causes perfect separation.
    Returns the name of the predictor or None.
    """
    for col in predictor_cols:
        if col not in df.columns:
            continue
            
        # Check if continuous (Float)
        # We assume Float types are continuous and ignore them as per instructions
        dtype = df.schema[col]
        if dtype in [pl.Float32, pl.Float64]:
            continue
            
        # Group by predictor and check distinct outcomes
        # If any value of the predictor has only 1 unique outcome, it causes separation
        counts = df.group_by(col).agg(pl.col(outcome_col).n_unique().alias("n_outcomes"))
        
        # We look for n_outcomes == 1
        if counts.select((pl.col("n_outcomes") == 1).any()).item():
            return col
            
    return None

def test_separation():
    # Case 1: Binary predictor causes separation
    # Predictor 'A': 0 -> Outcome 0, 1 -> Outcome 1 (Perfect)
    df1 = pl.DataFrame({
        "outcome": [0, 0, 1, 1],
        "A": [0, 0, 1, 1], # Separation
        "B": [0, 1, 0, 1]  # No separation
    })
    
    res1 = check_perfect_separation(df1, "outcome", ["A", "B"])
    print(f"Test 1 (Expect A): {res1}")
    assert res1 == "A"

    # Case 2: Predictor 'A' has one level that separates, another that doesn't
    # A=0 -> Outcome 0
    # A=1 -> Outcome 0 or 1
    # This is "quasi-complete" separation, but usually enough to cause issues/warnings
    df2 = pl.DataFrame({
        "outcome": [0, 0, 0, 1],
        "A":       [0, 0, 1, 1], 
        # A=0 -> [0, 0] (Outcome=0 only) -> n_unique=1 -> Separation
        # A=1 -> [0, 1] (Outcome mixed) -> n_unique=2
    })
    res2 = check_perfect_separation(df2, "outcome", ["A"])
    print(f"Test 2 (Expect A): {res2}")
    assert res2 == "A"

    # Case 3: Continuous variable (Float) - Should be ignored
    df3 = pl.DataFrame({
        "outcome": [0, 0, 1, 1],
        "C": [1.1, 1.2, 2.1, 2.2], # Effectively separates, but is float
    }).with_columns(pl.col("C").cast(pl.Float64))
    
    res3 = check_perfect_separation(df3, "outcome", ["C"])
    print(f"Test 3 (Expect None): {res3}")
    assert res3 is None

    # Case 4: No separation
    df4 = pl.DataFrame({
        "outcome": [0, 0, 1, 1],
        "A":       [0, 1, 0, 1]
    })
    res4 = check_perfect_separation(df4, "outcome", ["A"])
    print(f"Test 4 (Expect None): {res4}")
    assert res4 is None

if __name__ == "__main__":
    test_separation()
