app.py

# app.py — Gradio app for "Phone Before Bed" classifier (tabular, AutoGluon)
# Strategy: FORCE a non-fastai child model (LightGBM/XGB/Cat/Sklearn) to avoid fastai unpickling.

import os, pathlib, shutil, zipfile, traceback
import pandas as pd
import gradio as gr
from huggingface_hub import hf_hub_download
from autogluon.tabular import TabularPredictor

# --- MODEL ARTIFACT ON HUB ---
MODEL_REPO_ID = "jennifee/classical_automl_model"
ZIP_FILENAME  = "autogluon_predictor_dir.zip"
CACHE_DIR     = pathlib.Path("hf_assets")
EXTRACT_DIR   = CACHE_DIR / "predictor_dir"

# predictor expects these 5 features
REQUIRED_FEATURES = ["phone_hours", "computer_hours", "device_count", "sleep_time", "sleep_hours"]

# extra (ignored by model but kept in UI)
SLEEP_QUALITY_KEY     = "sleep_quality"
SLEEP_QUALITY_CHOICES = ["Poor", "Fair", "Good", "Excellent"]

LABEL_MAP = {
    0: "Does NOT use phone before bed",
    1: "Uses phone before bed",
    "0": "Does NOT use phone before bed",
    "1": "Uses phone before bed",
}

def _extract_dir(local_zip: str) -> str:
    if EXTRACT_DIR.exists():
        shutil.rmtree(EXTRACT_DIR)
    EXTRACT_DIR.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(local_zip, "r") as zf:
        zf.extractall(str(EXTRACT_DIR))
    items = list(EXTRACT_DIR.iterdir())
    return str(items[0] if len(items)==1 and items[0].is_dir() else EXTRACT_DIR)

def _load_predictor():
    local_zip = hf_hub_download(
        repo_id=MODEL_REPO_ID,
        filename=ZIP_FILENAME,
        repo_type="model",
        local_dir=str(CACHE_DIR),
        local_dir_use_symlinks=False,
    )
    predictor_dir = _extract_dir(local_zip)
    # relax version checks to be robust in Spaces
    return TabularPredictor.load(
        predictor_dir,
        require_version_match=False,
        require_py_version_match=False,
    )

PREDICTOR = _load_predictor()

# choose a safe (non-fastai) child model
SAFE_PATTERNS = ["LIGHTGBM", "XGBOOST", "CATBOOST", "LR", "RF", "XT", "KNN", "SKLEARN", "GBM", "XGB"]
AVOID_PATTERNS = ["FASTAI", "NN", "NEURAL"]

def _choose_safe_model(pred):
    try:
        names = pred.get_model_names()
    except Exception:
        return None, "ℹ️ Could not list child models; using default."
    def is_safe(name: str) -> bool:
        u = (name or "").upper()
        if any(a in u for a in AVOID_PATTERNS):
            return False
        if any(s in u for s in SAFE_PATTERNS):
            return True
        # allow plain sklearn models (often have short names)
        return True
    # prefer known good
    preferred = [m for m in names if any(s in m.upper() for s in SAFE_PATTERNS)]
    if preferred:
        return preferred[0], f"✅ Using model: {preferred[0]}"
    # else first non-fastai
    fallback = [m for m in names if is_safe(m)]
    if fallback:
        return fallback[0], f"✅ Using model: {fallback[0]}"
    return None, "⚠️ No non-fastai child models found; using default."

def predict_once(phone_hours, computer_hours, device_count, sleep_time, sleep_hours, sleep_quality):
    try:
        row = {
            "phone_hours": float(phone_hours),
            "computer_hours": float(computer_hours),
            "device_count": int(device_count),
            "sleep_time": float(sleep_time),
            "sleep_hours": float(sleep_hours),
            SLEEP_QUALITY_KEY: str(sleep_quality),
        }
        X = pd.DataFrame([row], columns=REQUIRED_FEATURES + [SLEEP_QUALITY_KEY])

        model_name, note = _choose_safe_model(PREDICTOR)

        # prediction
        if model_name:
            raw_label = PREDICTOR.predict(X, model=model_name).iloc[0]
        else:
            raw_label = PREDICTOR.predict(X).iloc[0]
        label = LABEL_MAP.get(raw_label, str(raw_label))

        # proba (best-effort)
        proba_text = ""
        try:
            if model_name:
                proba_df = PREDICTOR.predict_proba(X, model=model_name)
            else:
                proba_df = PREDICTOR.predict_proba(X)
            if proba_df is not None:
                row0 = proba_df.iloc[0]
                lines = [f"{LABEL_MAP.get(k, str(k))}: {int(round(float(v)*100))}%" for k, v in row0.items()]
                proba_text = "\n" + "\n".join(lines)
        except Exception as e:
            proba_text = f"\nℹ️ Could not compute probabilities: {e}"

        out = f"Prediction: {label}{proba_text}\n{note}"
        return out

    except Exception as e:
        # show error in the UI, not a crash
        return "❌ Prediction error: " + str(e) + "\n" + traceback.format_exc()

# ----------- Gradio UI -----------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📱 Use-Before-Bed Predictor")
    gr.Markdown("Predict whether a student uses their phone before bed based on device use and sleep habits.")

    with gr.Row():
        phone_hours    = gr.Slider(0, 24, step=0.1, value=2.5, label="Phone Hours / day")
        computer_hours = gr.Slider(0, 24, step=0.1, value=4.0, label="Computer Hours / day")
        device_count   = gr.Number(value=3, precision=0, label="Device Count")

    with gr.Row():
        sleep_time  = gr.Slider(0, 24, step=0.5, value=23.0, label="Weekday Sleep Start (0–24h)")
        sleep_hours = gr.Slider(0, 12, step=0.5, value=7.0, label="Sleep Duration (hours)")

    sleep_quality = gr.Radio(choices=SLEEP_QUALITY_CHOICES, value="Good", label="Sleep Quality")

    out = gr.Textbox(lines=8, label="Prediction & Confidence")

    inputs = [phone_hours, computer_hours, device_count, sleep_time, sleep_hours, sleep_quality]
    for c in inputs:
        c.change(fn=predict_once, inputs=inputs, outputs=out)

    gr.Examples(
        examples=[
            [2.5, 4, 3, 23.0, 7.0, "Good"],
            [1.0, 8, 5,  1.0, 5.0, "Poor"],
            [5.0, 2, 2, 22.5, 8.5, "Excellent"],
            [0.5,10, 4,  0.0, 6.0, "Fair"],
        ],
        inputs=inputs,
        label="Representative examples",
        examples_per_page=5,
        cache_examples=False,
    )

if __name__ == "__main__":
    demo.launch()