import os, pathlib, zipfile, tempfile
import pandas as pd
import gradio as gr
from huggingface_hub import hf_hub_download
import autogluon.tabular as ag

MODEL_REPO_ID = "jennifee/classical_automl_model"
ZIP_FILENAME  = "autogluon_predictor_dir.zip"

def _prepare_predictor_dir() -> str:
    local_zip = hf_hub_download(
        repo_id=MODEL_REPO_ID,
        filename=ZIP_FILENAME,
        repo_type="model",
    )
    workdir = tempfile.mkdtemp(prefix="ag_predictor_")
    with zipfile.ZipFile(local_zip, "r") as zf:
        zf.extractall(workdir)
    entries = list(pathlib.Path(workdir).iterdir())
    if len(entries) == 1 and entries[0].is_dir():
        return str(entries[0])
    return workdir

PREDICTOR_DIR = _prepare_predictor_dir()
PREDICTOR = ag.TabularPredictor.load(PREDICTOR_DIR, require_py_version_match=False, require_version_match=True)

FEATURE_COLS = ["phone_hours","computer_hours","device_count","sleep_quality","sleep_time","sleep_hours"]
LABEL_MAP = {0: "No (does not use phone before bed)", 1: "Yes (uses phone before bed)"}

def _non_fastai_models():
    # Filter out potential FastAI/NN models by name to avoid importing fastai if not installed
    try:
        names = PREDICTOR.get_model_names()
        names_upper = {n: n.upper() for n in names}
        keep = [n for n in names if ("FASTAI" not in names_upper[n]) and ("NN" not in names_upper[n])]
        return keep if keep else names
    except Exception:
        return None

def _predict_probs_with_fallback(X: pd.DataFrame):

    try:
        proba = PREDICTOR.predict_proba(X)
        if isinstance(proba, pd.Series):
            proba = proba.to_frame().T
        return proba
    except ModuleNotFoundError as e:
        # If fastai is missing, try excluding those models
        if "fastai" in str(e).lower():
            model_list = _non_fastai_models()
            if model_list:
                proba = PREDICTOR.predict_proba(X, model=model_list)
                if isinstance(proba, pd.Series):
                    proba = proba.to_frame().T
                return proba
        raise  # bubble up any other error

def do_predict(phone_hours, computer_hours, device_count, sleep_quality, sleep_time, sleep_hours):
    row = {
        "phone_hours":    int(phone_hours),
        "computer_hours": int(computer_hours),
        "device_count":   int(device_count),
        "sleep_quality":  str(sleep_quality),
        "sleep_time":     int(sleep_time),
        "sleep_hours":    int(sleep_hours),
    }
    X = pd.DataFrame([row], columns=FEATURE_COLS)

    try:
        proba = _predict_probs_with_fallback(X)
        row0 = proba.iloc[0]
        out = {LABEL_MAP.get(cls, str(cls)): float(val) for cls, val in row0.items()}
        out = dict(sorted(out.items(), key=lambda kv: kv[1], reverse=True))
    except Exception:
        # Ultimate fallback: hard label via best available model(s)
        pred = PREDICTOR.predict(X).iloc[0]
        out = {LABEL_MAP.get(pred, str(pred)): 1.0}
    return out

with gr.Blocks() as demo:
    gr.Markdown("# Sleep Habits → Phone-Before-Bed Predictor")
    with gr.Row():
        phone_hours    = gr.Slider(0, 24, step=1, value=3,  label="phone_hours (hrs/day)")
        computer_hours = gr.Slider(0, 24, step=1, value=5,  label="computer_hours (hrs/day)")
        device_count   = gr.Slider(0, 10, step=1, value=3,  label="device_count (# devices)")
    with gr.Row():
        sleep_quality  = gr.Dropdown(choices=["bad","medium","good"], value="good", label="sleep_quality")
        sleep_time     = gr.Slider(0, 23, step=1, value=23, label="sleep_time (hour 0–23)")
        sleep_hours    = gr.Slider(0, 14, step=1, value=7,  label="sleep_hours (hrs/night)")

    out = gr.Label(num_top_classes=2, label="Class probabilities")

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

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