Update app.py

app.py

@@ -1,114 +1,117 @@
-import io, json, os, base64
+# app.py
+import io, os, json, base64
 from pathlib import Path
-import streamlit as st
-import pandas as pd
+
 import numpy as np
+import pandas as pd
+import streamlit as st
 import joblib
-import matplotlib
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error
 
-# Plotly (for interactivity)
-HAVE_PLOTLY = True
-try:
-    import plotly.graph_objects as go
-    from plotly.subplots import make_subplots
-except Exception:
-    HAVE_PLOTLY = False
-
-# ---------------- Defaults ----------------
+# =========================
+# Constants / defaults
+# =========================
 FEATURES = ["Q, gpm", "SPP(psi)", "T (kft.lbf)", "WOB (klbf)", "ROP (ft/h)"]
 TARGET = "UCS"
+
 MODELS_DIR = Path("models")
 DEFAULT_MODEL = MODELS_DIR / "ucs_rf.joblib"
 MODEL_FALLBACKS = [MODELS_DIR / "model.joblib", MODELS_DIR / "model.pkl"]
 
-COLORS = {"pred": "#1f77b4", "actual": "#f2b702", "ref": "#5a5a5a"}
+COLORS = {
+    "pred": "#1f77b4",   # blue
+    "actual": "#f2c94c", # yellow
+    "ref": "#444444",    # 1:1 line
+}
 
-# ---------------- Page / Theme ----------------
+# =========================
+# Page config + CSS
+# =========================
 st.set_page_config(page_title="ST_GeoMech_UCS", page_icon="logo.png", layout="wide")
-st.markdown("<style>header, footer{visibility:hidden !important;}</style>", unsafe_allow_html=True)
-st.markdown(
-    """
-    <style>
-    .stApp { background: #FFFFFF; }
-    section[data-testid="stSidebar"] { background: #F6F9FC; }
-    .block-container { padding-top: .5rem; padding-bottom: .5rem; }
-
-    /* Default Streamlit button style (Run, Predict remain blue) */
-    .stButton>button{
-      background:#0d6efd; color:#fff; font-weight:bold; border-radius:8px; border:none; padding:10px 24px;
-    }
-    .stButton>button:hover{ filter: brightness(0.92); }
-
-    /* Hero */
-    .st-hero { display:flex; align-items:center; gap:16px; padding-top: 4px; }
-    .st-hero .brand { width:110px; height:110px; object-fit:contain; }
-    .st-hero h1 { margin:0; line-height:1.05; }
-    .st-hero .tagline { margin:2px 0 0 2px; color:#6b7280; font-size:1.05rem; font-style:italic; }
-    [data-testid="stBlock"]{ margin-top:0 !important; }
-
-    /* Color the sidebar buttons by order inside our wrappers */
-    .dev-actions > div.stButton:nth-child(1) button { background:#f59e0b !important; } /* Preview (orange) */
-    .dev-actions > div.stButton:nth-child(2) button { background:#0d6efd !important; } /* Run (blue) */
-    .dev-actions > div.stButton:nth-child(3) button { background:#198754 !important; } /* Proceed (green) */
-
-    .val-actions > div.stButton:nth-child(1) button { background:#f59e0b !important; } /* Preview (orange) */
-    .val-actions > div.stButton:nth-child(2) button { background:#0d6efd !important; } /* Predict (blue) */
-
-    .dev-actions .stButton button:disabled,
-    .val-actions .stButton button:disabled{ filter: grayscale(40%); opacity:.6; }
-    </style>
-    """,
-    unsafe_allow_html=True
-)
 
-# ---------------- Helpers ----------------
-try:
-    dialog = st.dialog
-except AttributeError:
-    def dialog(title):
-        def deco(fn):
-            def wrapper(*args, **kwargs):
-                with st.expander(title, expanded=True):
-                    return fn(*args, **kwargs)
-            return wrapper
-        return deco
-
-def _get_model_url(): return (os.environ.get("MODEL_URL", "") or "").strip()
-def rmse(y_true, y_pred): return float(np.sqrt(mean_squared_error(y_true, y_pred)))
-
-def ensure_cols(df, cols):
-    miss = [c for c in cols if c not in df.columns]
-    if miss: st.error(f"Missing columns: {miss}\nFound: {list(df.columns)}"); return False
-    return True
+st.markdown("""
+<style>
+/* Hide default header/footer chrome */
+header, footer {visibility: hidden !important;}
+.stApp { background: #ffffff; }
+
+/* Sidebar look */
+section[data-testid="stSidebar"] { background: #F6F9FC; }
+
+/* Hero */
+.st-hero { display:flex; align-items:center; gap:14px; padding: 4px 0 2px 0; }
+.st-hero .brand { width:90px; height:90px; object-fit:contain; }
+.st-hero h1 { margin:0; line-height:1.05; }
+.st-hero .tagline { margin:2px 0 0 2px; color:#6b7280; font-size:1.05rem; font-style:italic; }
+
+/* Keep hero snug to the top */
+[data-testid="stBlock"] { margin-top:0 !important; }
+
+/* Global primary button style (Run Model stays blue) */
+.stButton > button {
+  background:#2563eb; color:#fff; font-weight:600; border:none; border-radius:8px;
+  padding:9px 18px;
+}
+
+/* Orange preview button (scoped by wrapper) */
+#preview-btn button {
+  background:#f59e0b !important; color:#fff !important;
+}
+
+/* Green proceed button (scoped by wrapper) */
+#proceed-btn button {
+  background:#16a34a !important; color:#fff !important;
+}
+
+/* Info helper chip */
+.helper-note {
+  background:#e7f0ff; border-radius:10px; padding:14px 16px; border:1px solid #d4e3ff;
+  color:#0f172a;
+}
+
+/* Make tab content tighter */
+[data-baseweb="tab-border"] { margin-top: 0.2rem; }
+
+/* Plotly charts use white backgrounds via functions below */
+</style>
+""", unsafe_allow_html=True)
+
+# =========================
+# Utils
+# =========================
+def inline_logo(path="logo.png") -> str:
+    try:
+        p = Path(path)
+        if not p.exists(): return ""
+        return f"data:image/png;base64,{base64.b64encode(p.read_bytes()).decode('ascii')}"
+    except Exception:
+        return ""
 
-@st.cache_resource(show_spinner=False)
-def load_model(model_path: str): return joblib.load(model_path)
+def _get_model_url():
+    # Safe access (prevents the "No secrets files" banner)
+    try:
+        return (st.secrets.get("MODEL_URL", "") or os.environ.get("MODEL_URL", "") or "").strip()
+    except Exception:
+        return (os.environ.get("MODEL_URL", "") or "").strip()
 
 @st.cache_data(show_spinner=False)
-def parse_excel(data_bytes: bytes):
+def parse_excel_bytes(data_bytes: bytes):
     bio = io.BytesIO(data_bytes)
     xl = pd.ExcelFile(bio)
     return {sh: xl.parse(sh) for sh in xl.sheet_names}
 
-def read_book_bytes(data_bytes: bytes):
-    if not data_bytes: return {}
-    try: return parse_excel(data_bytes)
-    except Exception as e: st.error(f"Failed to read Excel: {e}"); return {}
-
-def find_sheet(book, names):
-    low2orig = {k.lower(): k for k in book.keys()}
-    for nm in names:
-        if nm.lower() in low2orig: return low2orig[nm.lower()]
-    return None
+def ensure_required_columns(df: pd.DataFrame, cols) -> bool:
+    miss = [c for c in cols if c not in df.columns]
+    if miss:
+        st.error(f"Missing columns: {miss}\nFound: {list(df.columns)}")
+        return False
+    return True
 
-def toast(msg):
-    try: st.toast(msg)
-    except Exception: st.info(msg)
+@st.cache_resource(show_spinner=False)
+def load_model(model_path: str):
+    return joblib.load(model_path)
 
 def infer_features_from_model(m):
+    # Try scikit-learn feature names if present
     try:
         if hasattr(m, "feature_names_in_") and len(getattr(m, "feature_names_in_")):
             return [str(x) for x in m.feature_names_in_]
@@ -121,56 +124,115 @@ def infer_features_from_model(m):
     except Exception: pass
     return None
 
-def inline_logo(path="logo.png") -> str:
-    try:
-        p = Path(path)
-        if not p.exists(): return ""
-        return f"data:image/png;base64,{base64.b64encode(p.read_bytes()).decode('ascii')}"
-    except Exception:
-        return ""
+def rmse(y_true, y_pred):  # convenience
+    from sklearn.metrics import mean_squared_error
+    return float(np.sqrt(mean_squared_error(y_true, y_pred)))
+
+# =========================
+# Model availability
+# =========================
+MODEL_URL = _get_model_url()
+
+def ensure_model_present() -> Path | None:
+    for p in [DEFAULT_MODEL, *MODEL_FALLBACKS]:
+        if p.exists(): return p
+    if MODEL_URL:
+        try:
+            import requests
+        except Exception:
+            st.error("Downloading the model requires 'requests'. Please add it to requirements.txt.")
+            return None
+        try:
+            DEFAULT_MODEL.parent.mkdir(parents=True, exist_ok=True)
+            with requests.get(MODEL_URL, stream=True) as r:
+                r.raise_for_status()
+                with open(DEFAULT_MODEL, "wb") as f:
+                    for chunk in r.iter_content(chunk_size=1<<20):
+                        f.write(chunk)
+            return DEFAULT_MODEL
+        except Exception as e:
+            st.error(f"Failed to download model from MODEL_URL. {e}")
+            return None
+    return None
+
+model_path = ensure_model_present()
+if not model_path:
+    st.error("Model not found. Upload models/ucs_rf.joblib (or set MODEL_URL in Settings → Variables).")
+    st.stop()
+
+try:
+    model = load_model(str(model_path))
+except Exception as e:
+    st.error(f"Failed to load model: {model_path}\n{e}")
+    st.stop()
 
-def export_workbook(sheets_dict: dict[str, pd.DataFrame], summary_df: pd.DataFrame|None):
+# Optional meta overrides
+meta_path = MODELS_DIR / "meta.json"
+if meta_path.exists():
     try:
-        import openpyxl  # noqa
+        meta = json.loads(meta_path.read_text(encoding="utf-8"))
+        FEATURES = meta.get("features", FEATURES)
+        TARGET = meta.get("target", TARGET)
     except Exception:
-        raise RuntimeError("Export requires openpyxl. Please add it to requirements.txt")
-    buf = io.BytesIO()
-    with pd.ExcelWriter(buf, engine="openpyxl") as xw:
-        for name, frame in sheets_dict.items():
-            frame.to_excel(xw, sheet_name=name[:31], index=False)
-        if summary_df is not None:
-            summary_df.to_excel(xw, sheet_name="Summary", index=False)
-    return buf.getvalue()
-
-# ---------- Plotting (Plotly first, MPL fallback) ----------
-def cross_plotly(actual, pred, title):
+        pass
+else:
+    _inf = infer_features_from_model(model)
+    if _inf: FEATURES = _inf
+
+# =========================
+# Plotly helpers (no titles, white background, safe margins)
+# =========================
+def _apply_plotly_base_layout(fig, *, top=40, left=60):
+    fig.update_layout(
+        margin=dict(l=left, r=10, t=top, b=40),
+        paper_bgcolor="#ffffff",
+        plot_bgcolor="#ffffff",
+        font=dict(size=12),
+    )
+    fig.update_xaxes(automargin=True, title_font=dict(size=12), tickfont=dict(size=11))
+    fig.update_yaxes(automargin=True, title_font=dict(size=12), tickfont=dict(size=11))
+    return fig
+
+def cross_plotly(actual, pred):
+    import plotly.graph_objects as go
     lo = float(np.nanmin([actual.min(), pred.min()]))
     hi = float(np.nanmax([actual.max(), pred.max()]))
     pad = 0.03 * (hi - lo if hi > lo else 1.0)
+
     fig = go.Figure()
     fig.add_trace(go.Scatter(
         x=actual, y=pred, mode="markers",
         marker=dict(size=6, color=COLORS["pred"]),
         hovertemplate="Actual: %{x:.2f}<br>Pred: %{y:.2f}<extra></extra>",
-        name="Points"
+        showlegend=False, name="Points",
     ))
     fig.add_trace(go.Scatter(
         x=[lo - pad, hi + pad], y=[lo - pad, hi + pad],
         mode="lines", line=dict(dash="dash", width=1.5, color=COLORS["ref"]),
-        hoverinfo="skip", showlegend=False
+        hoverinfo="skip", showlegend=False,
     ))
-    fig.update_layout(title=title, margin=dict(l=10, r=10, t=40, b=10), height=350,
-                      legend=dict(orientation="h", yanchor="bottom", y=1.02, x=0))
-    fig.update_xaxes(title_text="Actual UCS", scaleanchor="y", scaleratio=1)
-    fig.update_yaxes(title_text="Predicted UCS")
+
+    _apply_plotly_base_layout(fig, top=10, left=60)
+    fig.update_xaxes(
+        title_text="Actual UCS", title_standoff=10,
+        showgrid=True, gridcolor="rgba(0,0,0,0.12)",
+        zeroline=False, scaleanchor="y", scaleratio=1
+    )
+    fig.update_yaxes(
+        title_text="Predicted UCS", title_standoff=10,
+        showgrid=True, gridcolor="rgba(0,0,0,0.12)",
+        zeroline=False
+    )
     return fig
 
 def track_plotly(df, include_actual=True):
+    import plotly.graph_objects as go
     depth_col = next((c for c in df.columns if 'depth' in str(c).lower()), None)
     if depth_col is not None:
         y = df[depth_col]; y_label = depth_col
     else:
         y = np.arange(1, len(df) + 1); y_label = "Point Index"
+
     fig = go.Figure()
     fig.add_trace(go.Scatter(
         x=df["UCS_Pred"], y=y, mode="lines",
@@ -185,170 +247,88 @@ def track_plotly(df, include_actual=True):
             name="UCS (actual)",
             hovertemplate="UCS (actual): %{x:.2f}<br>"+y_label+": %{y}<extra></extra>"
         ))
-    fig.update_yaxes(autorange="reversed", title_text=y_label)
-    fig.update_xaxes(title_text="UCS", side="top")
-    fig.update_layout(margin=dict(l=10, r=10, t=40, b=10), height=650,
-                      legend=dict(orientation="h", yanchor="bottom", y=1.02, x=0))
+
+    _apply_plotly_base_layout(fig, top=60, left=70)
+    fig.update_layout(
+        legend=dict(orientation="h", yanchor="bottom", y=1.02, x=0),
+        height=650
+    )
+    fig.update_xaxes(
+        title_text="UCS", side="top", title_standoff=12,
+        showgrid=True, gridcolor="rgba(0,0,0,0.12)"
+    )
+    fig.update_yaxes(
+        title_text=y_label, autorange="reversed", title_standoff=10,
+        showgrid=True, gridcolor="rgba(0,0,0,0.12)"
+    )
     return fig
 
 def make_index_tracks_plotly(df: pd.DataFrame, cols: list[str]):
+    from plotly.subplots import make_subplots
+    import plotly.graph_objects as go
+
     cols = [c for c in cols if c in df.columns]
     if not cols:
         fig = go.Figure()
         fig.add_annotation(text="No selected columns in sheet", showarrow=False, x=0.5, y=0.5)
         fig.update_xaxes(visible=False); fig.update_yaxes(visible=False)
-        fig.update_layout(height=200, margin=dict(l=10,r=10,t=10,b=10))
+        fig.update_layout(height=200, margin=dict(l=10,r=10,t=10,b=10),
+                          paper_bgcolor="#ffffff", plot_bgcolor="#ffffff")
         return fig
+
     n = len(cols)
-    fig = make_subplots(rows=1, cols=n, shared_yaxes=True, horizontal_spacing=0.05)  # <-- FIX
+    # IMPORTANT: shared_yaxes (not shared_y)
+    fig = make_subplots(rows=1, cols=n, shared_yaxes=True, horizontal_spacing=0.05)
     idx = np.arange(1, len(df) + 1)
+
     for i, col in enumerate(cols, start=1):
         fig.add_trace(
-            go.Scatter(x=df[col], y=idx, mode="lines", line=dict(color="#333", width=1.2),
-                       hovertemplate=f"{col}: "+"%{x:.2f}<br>Index: %{y}<extra></extra>",
-                       name=col, showlegend=False),
+            go.Scatter(
+                x=df[col], y=idx, mode="lines",
+                line=dict(color="#333333", width=1.2),
+                hovertemplate=f"{col}: "+"%{x:.2f}<br>Index: %{y}<extra></extra>",
+                showlegend=False, name=col,
+            ), row=1, col=i
+        )
+        fig.update_xaxes(
+            title_text=col, side="top", title_standoff=10,
+            tickfont=dict(size=10),
+            showgrid=True, gridcolor="rgba(0,0,0,0.12)",
             row=1, col=i
         )
-        fig.update_xaxes(title_text=col, side="top", row=1, col=i)
-    fig.update_yaxes(autorange="reversed", title_text="Point Index", row=1, col=1)
-    fig.update_layout(height=650, margin=dict(l=10, r=10, t=40, b=10))
-    return fig
 
-# MPL fallbacks
-def cross_plot_mpl(actual, pred, title, size=(3.9,3.9)):
-    fig, ax = plt.subplots(figsize=size, dpi=100)
-    ax.scatter(actual, pred, s=14, alpha=0.85, color=COLORS["pred"])
-    lo = float(np.nanmin([actual.min(), pred.min()])); hi = float(np.nanmax([actual.max(), pred.max()]))
-    pad = 0.03 * (hi - lo if hi > lo else 1.0)
-    ax.plot([lo-pad, hi+pad], [lo-pad, hi+pad], '--', lw=1.2, color=COLORS["ref"])
-    ax.set_xlim(lo-pad, hi+pad); ax.set_ylim(lo-pad, hi+pad); ax.set_aspect('equal','box')
-    ax.set_xlabel("Actual UCS"); ax.set_ylabel("Predicted UCS"); ax.set_title(title); ax.grid(True, ls=":", alpha=0.4)
-    return fig
-
-def depth_or_index_track_mpl(df, title=None, include_actual=True):
-    depth_col = next((c for c in df.columns if 'depth' in str(c).lower()), None)
-    fig, ax = plt.subplots(figsize=(3.1, 7.2), dpi=100)
-    if depth_col is not None:
-        ax.plot(df["UCS_Pred"], df[depth_col], '-', lw=1.8, color=COLORS["pred"], label="UCS_Pred")
-        if include_actual and TARGET in df.columns:
-            ax.plot(df[TARGET], df[depth_col], ':', lw=2.0, color=COLORS["actual"], alpha=0.95, label="UCS (actual)")
-        ax.set_ylabel(depth_col); ax.set_xlabel("UCS"); ax.xaxis.set_label_position('top'); ax.xaxis.tick_top(); ax.invert_yaxis()
-    else:
-        idx = np.arange(1, len(df) + 1)
-        ax.plot(df["UCS_Pred"], idx, '-', lw=1.8, color=COLORS["pred"], label="UCS_Pred")
-        if include_actual and TARGET in df.columns:
-            ax.plot(df[TARGET], idx, ':', lw=2.0, color=COLORS["actual"], alpha=0.95, label="UCS (actual)")
-        ax.set_ylabel("Point Index"); ax.set_xlabel("UCS"); ax.xaxis.set_label_position('top'); ax.xaxis.tick_top(); ax.invert_yaxis()
-    ax.grid(True, linestyle=":", alpha=0.4); 
-    if title: ax.set_title(title, pad=8)
-    ax.legend(loc="best")
+    fig.update_yaxes(
+        autorange="reversed", title_text="Point Index", title_standoff=10,
+        tickfont=dict(size=10),
+        showgrid=True, gridcolor="rgba(0,0,0,0.12)",
+        row=1, col=1
+    )
+    fig.update_layout(
+        height=650,
+        margin=dict(l=60, r=10, t=60, b=40),
+        paper_bgcolor="#ffffff",
+        plot_bgcolor="#ffffff",
+        font=dict(size=12),
+    )
     return fig
 
-def stats_table(df: pd.DataFrame, cols: list[str]) -> pd.DataFrame:
-    cols = [c for c in cols if c in df.columns]
-    if not cols:
-        return pd.DataFrame({"Feature": [], "Min": [], "Max": [], "Mean": [], "Std": []})
-    out = df[cols].agg(['min', 'max', 'mean', 'std']).T
-    out = out.rename(columns={"min": "Min", "max": "Max", "mean": "Mean", "std": "Std"})
-    return out.reset_index().rename(columns={"index": "Feature"})
-
-# ---------- Preview dialogs ----------
-@dialog("Preview data")
-def preview_modal_dev(book: dict[str, pd.DataFrame], feature_cols: list[str]):
-    if not book:
-        st.info("No data loaded yet."); return
-    sh_train = find_sheet(book, ["Train","Training","training2","train","training"])
-    sh_test  = find_sheet(book, ["Test","Testing","testing2","test","testing"])
-    tabs, data = [], []
-    if sh_train: tabs.append("Train"); data.append(book[sh_train])
-    if sh_test:  tabs.append("Test");  data.append(book[sh_test])
-    if not tabs:
-        first_name = list(book.keys())[0]
-        tabs = [first_name]; data = [book[first_name]]
-    st.write("Use the tabs to switch between Train/Test views (if available).")
-    t_objs = st.tabs(tabs)
-    for t, df in zip(t_objs, data):
-        with t:
-            t1, t2 = st.tabs(["Tracks", "Summary"])
-            with t1:
-                if HAVE_PLOTLY:
-                    st.plotly_chart(make_index_tracks_plotly(df, feature_cols), use_container_width=True, theme=None)
-                else:
-                    st.pyplot(depth_or_index_track_mpl(df, title=None, include_actual=False), use_container_width=True)
-            with t2:
-                st.dataframe(stats_table(df, feature_cols), use_container_width=True)
-
-@dialog("Preview data")
-def preview_modal_val(book: dict[str, pd.DataFrame], feature_cols: list[str]):
-    if not book:
-        st.info("No data loaded yet."); return
-    vname = find_sheet(book, ["Validation","Validate","validation2","Val","val"]) or list(book.keys())[0]
-    df = book[vname]
-    t1, t2 = st.tabs(["Tracks", "Summary"])
-    with t1:
-        if HAVE_PLOTLY:
-            st.plotly_chart(make_index_tracks_plotly(df, feature_cols), use_container_width=True, theme=None)
-        else:
-            st.pyplot(depth_or_index_track_mpl(df, title=None, include_actual=False), use_container_width=True)
-    with t2:
-        st.dataframe(stats_table(df, feature_cols), use_container_width=True)
-
-# ---------------- Model presence ----------------
-MODEL_URL = _get_model_url()
-def ensure_model_present() -> Path:
-    for p in [DEFAULT_MODEL, *MODEL_FALLBACKS]:
-        if p.exists() and p.stat().st_size > 0:
-            return p
-    if not MODEL_URL: return None
-    try:
-        import requests
-        DEFAULT_MODEL.parent.mkdir(parents=True, exist_ok=True)
-        with st.status("Downloading model…", expanded=False):
-            with requests.get(MODEL_URL, stream=True, timeout=30) as r:
-                r.raise_for_status()
-                with open(DEFAULT_MODEL, "wb") as f:
-                    for chunk in r.iter_content(chunk_size=1<<20):
-                        if chunk: f.write(chunk)
-        return DEFAULT_MODEL
-    except Exception as e:
-        st.error(f"Failed to download model from MODEL_URL: {e}")
-        return None
-
-model_path = ensure_model_present()
-if not model_path:
-    st.error("Model not found. Upload models/ucs_rf.joblib (or set MODEL_URL in Settings → Variables).")
-    st.stop()
-
-try:
-    model = load_model(str(model_path))
-except Exception as e:
-    st.error(f"Failed to load model: {model_path}\n{e}")
-    st.stop()
-
-# Meta overrides / inference
-meta_path = MODELS_DIR / "meta.json"
-if meta_path.exists():
-    try:
-        meta = json.loads(meta_path.read_text(encoding="utf-8"))
-        FEATURES = meta.get("features", FEATURES); TARGET = meta.get("target", TARGET)
-    except Exception: pass
-else:
-    infer = infer_features_from_model(model)
-    if infer: FEATURES = infer
-
-# ---------------- Session state ----------------
-if "app_step" not in st.session_state: st.session_state.app_step = "intro"
-if "results" not in st.session_state: st.session_state.results = {}
-if "train_ranges" not in st.session_state: st.session_state.train_ranges = None
-
-for k, v in {
-    "dev_ready": False, "dev_file_loaded": False, "dev_previewed": False,
-    "dev_file_signature": None, "dev_file_bytes": b"", "dev_file_name": "",
-    "dev_file_rows": 0, "dev_file_cols": 0,
-}.items():
-    if k not in st.session_state: st.session_state[k] = v
-
-# ---------------- Hero ----------------
+# =========================
+# Session state defaults
+# =========================
+ss = st.session_state
+ss.setdefault("app_step", "dev")               # intro/dev/predict (you asked to start at dev)
+ss.setdefault("dev_bytes", None)               # raw uploaded bytes
+ss.setdefault("dev_book", None)                # parsed workbook dict
+ss.setdefault("dev_sheet_train", None)         # chosen train sheet
+ss.setdefault("dev_sheet_test", None)          # chosen test sheet
+ss.setdefault("dev_previewed", False)
+ss.setdefault("dev_ran", False)
+ss.setdefault("results", {})
+ss.setdefault("train_ranges", None)
+
+# =========================
+# Hero header
+# =========================
 st.markdown(
     f"""
     <div class="st-hero">
@@ -359,305 +339,320 @@ st.markdown(
       </div>
     </div>
     """,
-    unsafe_allow_html=True
+    unsafe_allow_html=True,
 )
 
-# ---------------- INTRO ----------------
-if st.session_state.app_step == "intro":
+# =========================
+# INTRO (kept for completeness – you said start in dev)
+# =========================
+if ss.app_step == "intro":
     st.header("Welcome!")
-    st.markdown("This software is developed by *Smart Thinking AI-Solutions Team* to estimate UCS from drilling data.")
-    st.subheader("Required Input Columns")
     st.markdown(
-        "- Q, gpm — Flow rate (gallons per minute)  \n"
-        "- SPP(psi) — Stand pipe pressure  \n"
-        "- T (kft.lbf) — Torque (thousand foot-pounds)  \n"
-        "- WOB (klbf) — Weight on bit  \n"
-        "- ROP (ft/h) — Rate of penetration"
+        "1. **Upload your data** to build the case and preview the performance of our model.\n"
+        "2. **Run Model** to compute metrics and plots.\n"
+        "3. **Proceed to Prediction** to validate on a new dataset and export results."
     )
-    st.subheader("How It Works")
-    st.markdown(
-        "1. **Upload your data to build the case and preview the performance of our model.**  \n"
-        "2. Click **Run Model** to compute metrics and plots.  \n"
-        "3. Click **Proceed to Prediction** to validate on a new dataset.  \n"
-        "4. Export results to Excel at any time."
-    )
-    if st.button("Start Showcase", type="primary", key="start_showcase"):
-        st.session_state.app_step = "dev"; st.rerun()
+    if st.button("Start", type="primary"): ss.app_step = "dev"; st.rerun()
 
-# ---------------- DEVELOPMENT ----------------
-if st.session_state.app_step == "dev":
+# =========================
+# DEVELOPMENT
+# =========================
+if ss.app_step == "dev":
+    # Sidebar controls
     st.sidebar.header("Model Development Data")
-    dev_label = "Upload Data (Excel)" if not st.session_state.dev_file_name else "Replace data (Excel)"
-    train_test_file = st.sidebar.file_uploader(dev_label, type=["xlsx","xls"], key="dev_upload")
-
-    if train_test_file is not None:
-        file_bytes = train_test_file.getvalue()
-        size = len(file_bytes)
-        sig = (train_test_file.name, size)
-        if sig != st.session_state.dev_file_signature and size > 0:
-            st.session_state.dev_file_signature = sig
-            st.session_state.dev_file_name = train_test_file.name
-            st.session_state.dev_file_bytes = file_bytes
-            _book_tmp = read_book_bytes(file_bytes)
-            if _book_tmp:
-                first_df = next(iter(_book_tmp.values()))
-                st.session_state.dev_file_rows = int(first_df.shape[0])
-                st.session_state.dev_file_cols = int(first_df.shape[1])
-                st.session_state.dev_file_loaded = True
-                st.session_state.dev_previewed = False
-                st.session_state.dev_ready = False
-
-    if st.session_state.dev_file_loaded:
-        st.sidebar.caption(
-            f"**Data loaded:** {st.session_state.dev_file_name} • "
-            f"{st.session_state.dev_file_rows} rows × {st.session_state.dev_file_cols} cols"
-        )
+    dev_file = st.sidebar.file_uploader("Replace data (Excel)", type=["xlsx","xls"], key="dev_upload")
 
-    st.sidebar.markdown('<div class="dev-actions">', unsafe_allow_html=True)
-    preview_btn = st.sidebar.button("Preview data", use_container_width=True, disabled=not st.session_state.dev_file_loaded)
-    run_btn = st.sidebar.button("Run Model", type="primary", use_container_width=True)
-    proceed_clicked = st.sidebar.button("Proceed to Prediction ▶", use_container_width=True, disabled=not st.session_state.dev_ready)
-    st.sidebar.markdown('</div>', unsafe_allow_html=True)
-
-    if proceed_clicked and st.session_state.dev_ready:
-        st.session_state.app_step = "predict"; st.rerun()
-
-    helper_top = st.container()
-    with helper_top:
-        st.subheader("Model Development")
-        if st.session_state.dev_ready:
-            st.success("Case has been built and results are displayed below.")
-        elif st.session_state.dev_file_loaded and st.session_state.dev_previewed:
-            st.info("Previewed ✓ — now click **Run Model** to build the case.")
-        elif st.session_state.dev_file_loaded:
-            st.info("📄 **Preview uploaded data** using the sidebar button, then click **Run Model**.")
-        else:
-            st.write("**Upload your data to build a case, then run the model to review development performance.**")
-
-    if preview_btn and st.session_state.dev_file_loaded and st.session_state.dev_file_bytes:
-        _book = read_book_bytes(st.session_state.dev_file_bytes)
-        st.session_state.dev_previewed = True
-        preview_modal_dev(_book, FEATURES)
-
-    if run_btn and st.session_state.dev_file_bytes:
-        with st.status("Processing…", expanded=False) as status:
-            book = read_book_bytes(st.session_state.dev_file_bytes)
-            if not book: status.update(label="Failed to read workbook.", state="error"); st.stop()
-            status.update(label="Workbook read ✓")
-            sh_train = find_sheet(book, ["Train","Training","training2","train","training"])
-            sh_test  = find_sheet(book, ["Test","Testing","testing2","test","testing"])
-            if sh_train is None or sh_test is None:
-                status.update(label="Workbook must include Train/Training/training2 and Test/Testing/testing2.", state="error"); st.stop()
-            df_tr = book[sh_train].copy(); df_te = book[sh_test].copy()
-            if not (ensure_cols(df_tr, FEATURES + [TARGET]) and ensure_cols(df_te, FEATURES + [TARGET])):
-                status.update(label="Missing required columns.", state="error"); st.stop()
-            status.update(label="Columns validated ✓"); status.update(label="Predicting…")
-
-            df_tr["UCS_Pred"] = model.predict(df_tr[FEATURES])
-            df_te["UCS_Pred"] = model.predict(df_te[FEATURES])
-            st.session_state.results["Train"] = df_tr; st.session_state.results["Test"] = df_te
-
-            st.session_state.results["metrics_train"] = {
-                "R2": r2_score(df_tr[TARGET], df_tr["UCS_Pred"]),
-                "RMSE": rmse(df_tr[TARGET], df_tr["UCS_Pred"]),
-                "MAE": mean_absolute_error(df_tr[TARGET], df_tr["UCS_Pred"]),
-            }
-            st.session_state.results["metrics_test"]  = {
-                "R2": r2_score(df_te[TARGET], df_te["UCS_Pred"]),
-                "RMSE": rmse(df_te[TARGET], df_te["UCS_Pred"]),
-                "MAE": mean_absolute_error(df_te[TARGET], df_te["UCS_Pred"]),
-            }
-
-            tr_min = df_tr[FEATURES].min().to_dict(); tr_max = df_tr[FEATURES].max().to_dict()
-            st.session_state.train_ranges = {f:(float(tr_min[f]), float(tr_max[f])) for f in FEATURES}
-
-            st.session_state.dev_ready = True
-            status.update(label="Done ✓", state="complete"); toast("Model run complete 🚀")
+    # Cache uploaded file into session (so preview doesn't clear it)
+    if dev_file is not None:
+        ss.dev_bytes = dev_file.getvalue()
+        try:
+            ss.dev_book = parse_excel_bytes(ss.dev_bytes)
+        except Exception as e:
+            st.sidebar.error(f"Failed to read workbook: {e}")
+            ss.dev_book = None
+            ss.dev_previewed = False
+            ss.dev_ran = False
+
+    # PREVIEW button (orange)
+    st.sidebar.markdown("<div id='preview-btn'>", unsafe_allow_html=True)
+    preview_click = st.sidebar.button("Preview data", use_container_width=True)
+    st.sidebar.markdown("</div>", unsafe_allow_html=True)
+
+    # RUN button (blue)
+    run_click = st.sidebar.button("Run Model", use_container_width=True)
+
+    # Proceed button (green; enabled after run)
+    st.sidebar.markdown("<div id='proceed-btn'>", unsafe_allow_html=True)
+    proceed_click = st.sidebar.button(
+        "Proceed to Prediction ▶",
+        use_container_width=True,
+        disabled=not ss.dev_ran
+    )
+    st.sidebar.markdown("</div>", unsafe_allow_html=True)
+
+    if proceed_click and ss.dev_ran:
+        ss.app_step = "predict"
+        st.rerun()
+
+    # Section heading
+    st.subheader("Model Development")
+
+    # Helper message (sticks here always)
+    helper = st.empty()
+    if ss.dev_book is None:
+        helper.markdown("<div class='helper-note'>Upload your data to build the case and preview the dataset.</div>", unsafe_allow_html=True)
+    elif not ss.dev_previewed:
+        helper.markdown("<div class='helper-note'>Data loaded ✓ — click <b>Preview data</b> to review tracks and summary.</div>", unsafe_allow_html=True)
+    elif ss.dev_previewed and not ss.dev_ran:
+        helper.markdown("<div class='helper-note'>Previewed ✓ — now click <b>Run Model</b> to build the case.</div>", unsafe_allow_html=True)
+    else:
+        helper.markdown("<div class='helper-note'>Case built ✓ — results are displayed below.</div>", unsafe_allow_html=True)
+
+    # ----------------- Preview modal -----------------
+    def preview_modal(book: dict, feature_cols: list[str]):
+        if not book: return
+        with st.expander("▼ Preview (tracks & summary)", expanded=True):
+            # Choose a sheet to preview
+            sheetnames = list(book.keys())
+            sh = st.selectbox("Sheet", options=sheetnames, index=0, key="preview_sheet_sel")
+            df = book[sh].copy()
+
+            # Tracks tab + Stats tab
+            t1, t2 = st.tabs(["Tracks", "Summary"])
+
+            with t1:
+                fig = make_index_tracks_plotly(df, feature_cols)
+                st.plotly_chart(fig, use_container_width=True, config={"displayModeBar": False})
+
+            with t2:
+                stats = df[feature_cols].describe().T[["min", "max", "mean", "std"]].rename(
+                    columns={"min":"Min", "max":"Max", "mean":"Mean", "std":"Std"}
+                )
+                st.dataframe(stats, use_container_width=True)
+
+    # If preview clicked and we have data
+    if preview_click:
+        if ss.dev_book:
+            preview_modal(ss.dev_book, FEATURES)
+            ss.dev_previewed = True
+            ss.dev_ran = False
             st.rerun()
+        else:
+            st.warning("Please upload an Excel file first.")
 
-    if ("Train" in st.session_state.results) or ("Test" in st.session_state.results):
-        tab1, tab2 = st.tabs(["Training", "Testing"])
-        if "Train" in st.session_state.results:
-            with tab1:
-                df = st.session_state.results["Train"]; m = st.session_state.results["metrics_train"]
-                c1,c2,c3 = st.columns(3)
-                c1.metric("R²", f"{m['R2']:.4f}"); c2.metric("RMSE", f"{m['RMSE']:.4f}"); c3.metric("MAE", f"{m['MAE']:.4f}")
-                left, right = st.columns([0.9, 0.55])
-                with left:
-                    if HAVE_PLOTLY:
-                        st.plotly_chart(cross_plotly(df[TARGET], df["UCS_Pred"], "Training: Actual vs Predicted"),
-                                        use_container_width=True, theme=None)
-                    else:
-                        st.pyplot(cross_plot_mpl(df[TARGET], df["UCS_Pred"], "Training: Actual vs Predicted"),
-                                  use_container_width=True)
-                with right:
-                    if HAVE_PLOTLY:
-                        st.plotly_chart(track_plotly(df, include_actual=True), use_container_width=True, theme=None)
-                    else:
-                        st.pyplot(depth_or_index_track_mpl(df, title=None, include_actual=True),
-                                  use_container_width=True)
-        if "Test" in st.session_state.results:
-            with tab2:
-                df = st.session_state.results["Test"]; m = st.session_state.results["metrics_test"]
-                c1,c2,c3 = st.columns(3)
-                c1.metric("R²", f"{m['R2']:.4f}"); c2.metric("RMSE", f"{m['RMSE']:.4f}"); c3.metric("MAE", f"{m['MAE']:.4f}")
-                left, right = st.columns([0.9, 0.55])
-                with left:
-                    if HAVE_PLOTLY:
-                        st.plotly_chart(cross_plotly(df[TARGET], df["UCS_Pred"], "Testing: Actual vs Predicted"),
-                                        use_container_width=True, theme=None)
-                    else:
-                        st.pyplot(cross_plot_mpl(df[TARGET], df["UCS_Pred"], "Testing: Actual vs Predicted"),
-                                  use_container_width=True)
-                with right:
-                    if HAVE_PLOTLY:
-                        st.plotly_chart(track_plotly(df, include_actual=True), use_container_width=True, theme=None)
-                    else:
-                        st.pyplot(depth_or_index_track_mpl(df, title=None, include_actual=True),
-                                  use_container_width=True)
+    # If run clicked and we have data
+    if run_click:
+        if not ss.dev_book:
+            st.warning("Please upload and preview your data first.")
+        else:
+            # Try to find common sheet names
+            names = list(ss.dev_book.keys())
+            def find_sheet(book, alts):
+                lo = {k.lower(): k for k in book.keys()}
+                for nm in alts:
+                    if nm.lower() in lo: return lo[nm.lower()]
+                return None
+
+            sh_train = find_sheet(ss.dev_book, ["Train","Training","training2","train","training"]) or names[0]
+            sh_test  = find_sheet(ss.dev_book, ["Test","Testing","testing2","test","testing"]) or (names[1] if len(names)>1 else names[0])
+            ss.dev_sheet_train, ss.dev_sheet_test = sh_train, sh_test
+
+            df_tr = ss.dev_book[sh_train].copy()
+            df_te = ss.dev_book[sh_test].copy()
+
+            ok = ensure_required_columns(df_tr, FEATURES+[TARGET]) and ensure_required_columns(df_te, FEATURES+[TARGET])
+            if ok:
+                df_tr["UCS_Pred"] = model.predict(df_tr[FEATURES])
+                df_te["UCS_Pred"] = model.predict(df_te[FEATURES])
+
+                from sklearn.metrics import r2_score, mean_absolute_error
+                ss.results["Train"] = df_tr
+                ss.results["Test"]  = df_te
+                ss.results["metrics_train"] = {
+                    "R2": r2_score(df_tr[TARGET], df_tr["UCS_Pred"]),
+                    "RMSE": rmse(df_tr[TARGET], df_tr["UCS_Pred"]),
+                    "MAE": mean_absolute_error(df_tr[TARGET], df_tr["UCS_Pred"]),
+                }
+                ss.results["metrics_test"]  = {
+                    "R2": r2_score(df_te[TARGET], df_te["UCS_Pred"]),
+                    "RMSE": rmse(df_te[TARGET], df_te["UCS_Pred"]),
+                    "MAE": mean_absolute_error(df_te[TARGET], df_te["UCS_Pred"]),
+                }
 
-        st.markdown("---")
-        sheets = {}; rows = []
-        if "Train" in st.session_state.results:
-            sheets["Train_with_pred"] = st.session_state.results["Train"]
-            rows.append({"Split":"Train", **{k:round(v,6) for k,v in st.session_state.results["metrics_train"].items()}})
-        if "Test" in st.session_state.results:
-            sheets["Test_with_pred"] = st.session_state.results["Test"]
-            rows.append({"Split":"Test", **{k:round(v,6) for k,v in st.session_state.results["metrics_test"].items()}})
-        summary_df = pd.DataFrame(rows) if rows else None
-        try:
-            data_bytes = export_workbook(sheets, summary_df)
-            st.download_button("Export Development Results to Excel",
-                data=data_bytes, file_name="UCS_Dev_Results.xlsx",
-                mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
-        except RuntimeError as e:
-            st.warning(str(e))
+                tr_min = df_tr[FEATURES].min().to_dict()
+                tr_max = df_tr[FEATURES].max().to_dict()
+                ss.train_ranges = {f:(float(tr_min[f]), float(tr_max[f])) for f in FEATURES}
 
-# ---------------- PREDICTION ----------------
-if st.session_state.app_step == "predict":
+                ss.dev_ran = True
+                helper.markdown("<div class='helper-note'>Case built ✓ — results are displayed below.</div>", unsafe_allow_html=True)
+            else:
+                ss.dev_ran = False
+
+    # Show results if available
+    if ss.dev_ran and ("Train" in ss.results or "Test" in ss.results):
+        ttr, tte = st.tabs(["Training", "Testing"])
+
+        if "Train" in ss.results:
+            with ttr:
+                m = ss.results["metrics_train"]
+                c1,c2,c3 = st.columns([1,1,1])
+                c1.metric("R²", f"{m['R2']:.4f}")
+                c2.metric("RMSE", f"{m['RMSE']:.4f}")
+                c3.metric("MAE", f"{m['MAE']:.4f}")
+                l, r = st.columns([0.55, 0.45])
+                with l:
+                    st.plotly_chart(cross_plotly(ss.results["Train"][TARGET], ss.results["Train"]["UCS_Pred"]),
+                                    use_container_width=True, config={"displayModeBar": False})
+                with r:
+                    st.plotly_chart(track_plotly(ss.results["Train"], include_actual=True),
+                                    use_container_width=True, config={"displayModeBar": False})
+
+        if "Test" in ss.results:
+            with tte:
+                m = ss.results["metrics_test"]
+                c1,c2,c3 = st.columns([1,1,1])
+                c1.metric("R²", f"{m['R2']:.4f}")
+                c2.metric("RMSE", f"{m['RMSE']:.4f}")
+                c3.metric("MAE", f"{m['MAE']:.4f}")
+                l, r = st.columns([0.55, 0.45])
+                with l:
+                    st.plotly_chart(cross_plotly(ss.results["Test"][TARGET], ss.results["Test"]["UCS_Pred"]),
+                                    use_container_width=True, config={"displayModeBar": False})
+                with r:
+                    st.plotly_chart(track_plotly(ss.results["Test"], include_actual=True),
+                                    use_container_width=True, config={"displayModeBar": False})
+
+# =========================
+# PREDICTION
+# =========================
+if ss.app_step == "predict":
     st.sidebar.header("Prediction (Validation)")
-    validation_file = st.sidebar.file_uploader("Upload Validation Excel", type=["xlsx","xls"], key="val_upload")
-    if validation_file is not None:
-        _book_tmp = read_book_bytes(validation_file.getvalue())
-        if _book_tmp:
-            first_df = next(iter(_book_tmp.values()))
-            st.sidebar.caption(f"**Data loaded:** {validation_file.name} • {first_df.shape[0]} rows × {first_df.shape[1]} cols")
-
-    st.sidebar.markdown('<div class="val-actions">', unsafe_allow_html=True)
-    preview_val_btn = st.sidebar.button("Preview data", use_container_width=True, disabled=(validation_file is None))
-    predict_btn = st.sidebar.button("Predict", type="primary", use_container_width=True)
-    st.sidebar.button("⬅ Back", on_click=lambda: st.session_state.update(app_step="dev"), use_container_width=True)
-    st.sidebar.markdown('</div>', unsafe_allow_html=True)
+    val_file = st.sidebar.file_uploader("Upload Validation Excel", type=["xlsx","xls"], key="val_upload")
+    predict_click = st.sidebar.button("Predict", use_container_width=True)
+    back_click = st.sidebar.button("⬅ Back", use_container_width=True)
+
+    if back_click:
+        ss.app_step = "dev"; st.rerun()
 
     st.subheader("Prediction")
-    st.write("Upload a new dataset to generate UCS predictions and evaluate performance on unseen data.")
-
-    if preview_val_btn and validation_file is not None:
-        _book = read_book_bytes(validation_file.getvalue())
-        preview_modal_val(_book, FEATURES)
-
-    if predict_btn and validation_file is not None:
-        with st.status("Predicting…", expanded=False) as status:
-            vbook = read_book_bytes(validation_file.getvalue())
-            if not vbook: status.update(label="Could not read the Validation Excel.", state="error"); st.stop()
-            status.update(label="Workbook read ✓")
-            vname = find_sheet(vbook, ["Validation","Validate","validation2","Val","val"]) or list(vbook.keys())[0]
+    st.markdown("Upload a new dataset to generate UCS predictions and evaluate performance on unseen data.")
+    st.success("Predictions ready ✓" if "Validate" in ss.results else "Waiting for input…")
+
+    if predict_click and val_file is not None:
+        try:
+            vbook = parse_excel_bytes(val_file.getvalue())
+        except Exception as e:
+            st.error(f"Could not read the Validation Excel: {e}")
+            vbook = {}
+
+        if vbook:
+            # Pick first sheet by default
+            vname = list(vbook.keys())[0]
             df_val = vbook[vname].copy()
-            if not ensure_cols(df_val, FEATURES): status.update(label="Missing required columns.", state="error"); st.stop()
-            status.update(label="Columns validated ✓")
-            df_val["UCS_Pred"] = model.predict(df_val[FEATURES])
-            st.session_state.results["Validate"] = df_val
-
-            ranges = st.session_state.train_ranges; oor_table = None; oor_pct = 0.0
-            if ranges:
-                viol = {f: (df_val[f] < ranges[f][0]) | (df_val[f] > ranges[f][1]) for f in FEATURES}
-                any_viol = pd.DataFrame(viol).any(axis=1); oor_pct = float(any_viol.mean()*100.0)
-                if any_viol.any():
-                    offenders = df_val.loc[any_viol, FEATURES].copy()
-                    offenders["Violations"] = pd.DataFrame(viol).loc[any_viol].apply(
-                        lambda r: ", ".join([c for c,v in r.items() if v]), axis=1)
-                    offenders.index = offenders.index + 1; oor_table = offenders
-
-            metrics_val = None
-            if TARGET in df_val.columns:
-                metrics_val = {
-                    "R2": r2_score(df_val[TARGET], df_val["UCS_Pred"]),
-                    "RMSE": rmse(df_val[TARGET], df_val["UCS_Pred"]),
-                    "MAE": mean_absolute_error(df_val[TARGET], df_val["UCS_Pred"])
-                }
-            st.session_state.results["metrics_val"] = metrics_val
-            st.session_state.results["summary_val"] = {
-                "n_points": len(df_val),
-                "pred_min": float(df_val["UCS_Pred"].min()),
-                "pred_max": float(df_val["UCS_Pred"].max()),
-                "oor_pct": oor_pct
-            }
-            st.session_state.results["oor_table"] = oor_table
-            status.update(label="Predictions ready ✓", state="complete")
-
-    if "Validate" in st.session_state.results:
-        st.subheader("Validation Results")
-        sv = st.session_state.results["summary_val"]; oor_table = st.session_state.results.get("oor_table")
 
-        if sv["oor_pct"] > 0:
-            st.warning("Some validation inputs fall outside the **training min–max** ranges. Interpret predictions with caution.")
+            if ensure_required_columns(df_val, FEATURES):
+                df_val["UCS_Pred"] = model.predict(df_val[FEATURES])
+                ss.results["Validate"] = df_val
+
+                # Out-of-range check vs training ranges
+                ranges = ss.train_ranges; oor_table = None; oor_pct = 0.0
+                if ranges:
+                    viol = {f: (df_val[f] < ranges[f][0]) | (df_val[f] > ranges[f][1]) for f in FEATURES}
+                    any_viol = pd.DataFrame(viol).any(axis=1); oor_pct = float(any_viol.mean()*100.0)
+                    if any_viol.any():
+                        offenders = df_val.loc[any_viol, FEATURES].copy()
+                        offenders["Violations"] = pd.DataFrame(viol).loc[any_viol].apply(
+                            lambda r: ", ".join([c for c,v in r.items() if v]), axis=1)
+                        offenders.index = offenders.index + 1; oor_table = offenders
+
+                from sklearn.metrics import r2_score, mean_absolute_error
+                metrics_val = None
+                if TARGET in df_val.columns:
+                    metrics_val = {
+                        "R2": r2_score(df_val[TARGET], df_val["UCS_Pred"]),
+                        "RMSE": rmse(df_val[TARGET], df_val["UCS_Pred"]),
+                        "MAE": mean_absolute_error(df_val[TARGET], df_val["UCS_Pred"]),
+                    }
+                ss.results["metrics_val"] = metrics_val
+                ss.results["summary_val"] = {
+                    "n_points": len(df_val),
+                    "pred_min": float(df_val["UCS_Pred"].min()),
+                    "pred_max": float(df_val["UCS_Pred"].max()),
+                    "oor_pct": oor_pct
+                }
+                ss.results["oor_table"] = oor_table
+                st.experimental_rerun()
 
+    # Show prediction results
+    if "Validate" in ss.results:
+        sv = ss.results["summary_val"]; oor_table = ss.results.get("oor_table")
         c1,c2,c3,c4 = st.columns(4)
-        c1.metric("points", f"{sv['n_points']}"); c2.metric("Pred min", f"{sv['pred_min']:.2f}")
-        c3.metric("Pred max", f"{sv['pred_max']:.2f}"); c4.metric("OOR %", f"{sv['oor_pct']:.1f}%")
+        c1.metric("# points", f"{sv['n_points']}")
+        c2.metric("Pred min", f"{sv['pred_min']:.2f}")
+        c3.metric("Pred max", f"{sv['pred_max']:.2f}")
+        c4.metric("OOR %", f"{sv['oor_pct']:.1f}%")
+        if sv["oor_pct"] > 0:
+            st.warning("Some validation rows contain inputs outside the Training min–max ranges. Review the table below.")
 
-        left, right = st.columns([0.9, 0.55])
+        left, right = st.columns([0.55, 0.45])
         with left:
-            if TARGET in st.session_state.results["Validate"].columns:
-                if HAVE_PLOTLY:
-                    st.plotly_chart(cross_plotly(st.session_state.results["Validate"][TARGET],
-                                                 st.session_state.results["Validate"]["UCS_Pred"],
-                                                 "Validation: Actual vs Predicted"),
-                                    use_container_width=True, theme=None)
-                else:
-                    st.pyplot(cross_plot_mpl(st.session_state.results["Validate"][TARGET],
-                                             st.session_state.results["Validate"]["UCS_Pred"],
-                                             "Validation: Actual vs Predicted"),
-                              use_container_width=True)
+            if TARGET in ss.results["Validate"].columns:
+                st.plotly_chart(
+                    cross_plotly(ss.results["Validate"][TARGET], ss.results["Validate"]["UCS_Pred"]),
+                    use_container_width=True, config={"displayModeBar": False}
+                )
             else:
                 st.info("Actual UCS values are not available in the validation data. Cross-plot cannot be generated.")
         with right:
-            if HAVE_PLOTLY:
-                st.plotly_chart(track_plotly(st.session_state.results["Validate"],
-                                             include_actual=(TARGET in st.session_state.results["Validate"].columns)),
-                                use_container_width=True, theme=None)
-            else:
-                st.pyplot(depth_or_index_track_mpl(st.session_state.results["Validate"], title=None,
-                                                   include_actual=(TARGET in st.session_state.results["Validate"].columns)),
-                          use_container_width=True)
+            st.plotly_chart(
+                track_plotly(ss.results["Validate"], include_actual=(TARGET in ss.results["Validate"].columns)),
+                use_container_width=True, config={"displayModeBar": False}
+            )
 
         if oor_table is not None:
             st.write("*Out-of-range rows (vs. Training min–max):*")
             st.dataframe(oor_table, use_container_width=True)
 
+        # Export
+        def export_workbook(sheets_dict, summary_df=None):
+            try:
+                import openpyxl
+            except Exception:
+                raise RuntimeError("Export requires openpyxl. Please add it to requirements.txt.")
+            buf = io.BytesIO()
+            with pd.ExcelWriter(buf, engine="openpyxl") as xw:
+                for name, frame in sheets_dict.items():
+                    frame.to_excel(xw, sheet_name=name[:31], index=False)
+                if summary_df is not None:
+                    summary_df.to_excel(xw, sheet_name="Summary", index=False)
+            return buf.getvalue()
+
         st.markdown("---")
-        sheets = {"Validate_with_pred": st.session_state.results["Validate"]}
+        sheets_to_save = {"Validate_with_pred": ss.results["Validate"]}
         rows = []
         for name, key in [("Train","metrics_train"), ("Test","metrics_test"), ("Validate","metrics_val")]:
-            m = st.session_state.results.get(key)
+            m = ss.results.get(key)
             if m: rows.append({"Split": name, **{k: round(v,6) for k,v in m.items()}})
         summary_df = pd.DataFrame(rows) if rows else None
         try:
-            data_bytes = export_workbook(sheets, summary_df)
+            data_bytes = export_workbook(sheets_to_save, summary_df)
             st.download_button("Export Validation Results to Excel",
                 data=data_bytes, file_name="UCS_Validation_Results.xlsx",
                 mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
         except RuntimeError as e:
             st.warning(str(e))
 
-# ---------------- Footer ----------------
+# =========================
+# Footer
+# =========================
 st.markdown("---")
 st.markdown(
-    """
-    <div style='text-align:center; color:#6b7280; line-height:1.6'>
-      ST_GeoMech_UCS • © Smart Thinking<br/>
-      <strong>Visit our website:</strong> <a href='https://www.smartthinking.com.sa' target='_blank'>smartthinking.com.sa</a>
-    </div>
-    """,
+    "<div style='text-align:center; color:#6b7280;'>"
+    "ST_GeoMech_UCS • © Smart Thinking"
+    "</div>"
+    "<div style='text-align:center; color:#6b7280;'>"
+    "Visit our Website: "
+    "<a href='https://www.smartthinking.com.sa' target='_blank'>smartthinking.com.sa</a>"
+    "</div>",
     unsafe_allow_html=True
 )