# Streamlit UI that downloads a scikit-learn pipeline from HF
import os, sys, logging, joblib, numpy as np, pandas as pd
from huggingface_hub import hf_hub_download
from sklearn.exceptions import InconsistentVersionWarning
import warnings

warnings.filterwarnings("ignore", category=InconsistentVersionWarning)

# Silence noisy logs when not run via `streamlit run`
if "streamlit" not in " ".join(sys.argv).lower():
    for name in ("streamlit.runtime.scriptrunner.script_run_context",
                 "streamlit.runtime.scriptrunner","streamlit"):
        lg = logging.getLogger(name); lg.setLevel(logging.CRITICAL); lg.propagate=False; lg.disabled=True

HF_MODEL_REPO = os.getenv("HF_MODEL_REPO", "dhani10/engine-condition-model")
MODEL_FILE    = os.getenv("MODEL_FILE", "model/best_engine_model.joblib")
HF_TOKEN      = os.getenv("HF_TOKEN")  # add as Space Secret if model repo is private

HF_CACHE_ROOT = os.getenv("HF_HOME", "/tmp/huggingface")
os.environ["HF_HOME"] = HF_CACHE_ROOT
os.environ["HF_HUB_CACHE"] = os.path.join(HF_CACHE_ROOT, "hub")
os.makedirs(os.environ["HF_HUB_CACHE"], exist_ok=True)

def _load_model_impl():
    path = hf_hub_download(
        repo_id=HF_MODEL_REPO,
        filename=MODEL_FILE,
        repo_type="model",
        token=HF_TOKEN,                    # None if public
        cache_dir=os.environ["HF_HUB_CACHE"],
    )
    return joblib.load(path)

def get_expected_input_columns(clf):
    pre = getattr(getattr(clf, "named_steps", {}), "get", lambda *_: None)("preprocessor")
    if pre is not None:
        transformers = getattr(pre, "transformers_", getattr(pre, "transformers", []))
        cols = []
        for _, __, selected in transformers:
            if selected in (None, "drop"): continue
            if isinstance(selected, list): cols.extend(selected)
            elif hasattr(selected, "__iter__"): cols.extend(list(selected))
        cols = list(dict.fromkeys(cols))
        if cols: return cols
    fni = getattr(clf, "feature_names_in_", None)
    return list(fni) if fni is not None else [
        "engine_rpm","lub_oil_pressure","fuel_pressure",
        "coolant_pressure","lub_oil_temp","coolant_temp"
    ]

def coerce_numeric_df(df: pd.DataFrame) -> pd.DataFrame:
    out = df.copy()
    for c in out.columns: out[c] = pd.to_numeric(out[c], errors="ignore")
    return out

def predict_with_pipeline(model, X: pd.DataFrame):
    y = model.predict(X); p = None
    if hasattr(model, "predict_proba"):
        try:
            P = model.predict_proba(X); p = P[:,1] if (P.ndim==2 and P.shape[1]>=2) else P.ravel()
        except Exception: pass
    return y, p

def main():
    import streamlit as st
    st.set_page_config(page_title="Engine Condition Predictor", layout="centered")
    st.title("Predictive Maintenance — Engine Condition")
    st.caption(f"Model: {HF_MODEL_REPO} → {MODEL_FILE}")

    @st.cache_resource(show_spinner=True)
    def load_model(): return _load_model_impl()

    model = load_model()
    EXPECTED_COLS = get_expected_input_columns(model)

    with st.form("predict_form"):
        col1, col2 = st.columns(2)
        with col1:
            engine_rpm       = st.number_input("Engine RPM", min_value=0, max_value=5000, value=1200, step=10)
            lub_oil_pressure = st.number_input("Lubricating Oil Pressure (bar)", value=3.0, step=0.1)
            fuel_pressure    = st.number_input("Fuel Pressure (bar)", value=5.0, step=0.1)
        with col2:
            coolant_pressure = st.number_input("Coolant Pressure (bar)", value=2.0, step=0.1)
            lub_oil_temp     = st.number_input("Lubricating Oil Temperature (°C)", value=80.0, step=0.1)
            coolant_temp     = st.number_input("Coolant Temperature (°C)", value=75.0, step=0.1)
        submitted = st.form_submit_button("Predict")

    if submitted:
        row = pd.DataFrame({c:[np.nan] for c in EXPECTED_COLS})
        for k,v in {
            "engine_rpm":engine_rpm,"lub_oil_pressure":lub_oil_pressure,"fuel_pressure":fuel_pressure,
            "coolant_pressure":coolant_pressure,"lub_oil_temp":lub_oil_temp,"coolant_temp":coolant_temp
        }.items():
            if k in row.columns: row.at[0,k]=v
        try:
            X = coerce_numeric_df(row)
            y, p = predict_with_pipeline(model, X)
            pred = int(y[0])
            if pred==1:
                msg = "⚠️ Faulty Engine Detected"
                if p is not None: msg += f" (Confidence: {float(p[0]):.2f})"
                import streamlit as st; st.error(msg)
            else:
                msg = "✅ Engine is Healthy"
                if p is not None: msg += f" (Confidence: {1 - float(p[0]):.2f})"
                import streamlit as st; st.success(msg)
            with st.expander("Inputs sent to the model"):
                st.dataframe(X)
        except Exception as e:
            import streamlit as st
            st.error(f"Prediction failed: {e}")
            st.write("Expected columns:", EXPECTED_COLS)

if __name__ == "__main__":
    if "streamlit" in " ".join(sys.argv).lower(): main()
    else: print("Tip: run this app with: streamlit run streamlit_app.py")