Update app.py

app.py

@@ -5,11 +5,11 @@ import pandas as pd
 import numpy as np
 import joblib
 import matplotlib
-matplotlib.use("Agg")  # fallback only
+matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error
 
-# Try Plotly for interactivity
+# Plotly (for interactivity)
 HAVE_PLOTLY = True
 try:
     import plotly.graph_objects as go
@@ -17,20 +17,16 @@ try:
 except Exception:
     HAVE_PLOTLY = False
 
-# =========================
-# Defaults
-# =========================
+# ---------------- 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", "orange": "#f59e0b", "green": "#198754"}
+COLORS = {"pred": "#1f77b4", "actual": "#f2b702", "ref": "#5a5a5a"}
 
-# =========================
-# Page / Theme
-# =========================
+# ---------------- Page / Theme ----------------
 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(
@@ -39,35 +35,36 @@ st.markdown(
     .stApp { background: #FFFFFF; }
     section[data-testid="stSidebar"] { background: #F6F9FC; }
     .block-container { padding-top: .5rem; padding-bottom: .5rem; }
-    .stButton>button{ background:#0d6efd; color:#fff; font-weight:bold; border-radius:8px; border:none; padding:10px 24px; }
+
+    /* 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 header */
+    /* 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; }
 
-    /* Colorize sidebar button groups we wrap with custom classes */
-    section[data-testid="stSidebar"] .dev-actions .stButton:nth-of-type(1) button { background: #f59e0b !important; } /* Preview - orange */
-    section[data-testid="stSidebar"] .dev-actions .stButton:nth-of-type(2) button { background: #0d6efd !important; } /* Run - blue */
-    section[data-testid="stSidebar"] .dev-actions .stButton:nth-of-type(3) button { background: #198754 !important; } /* Proceed - green */
+    /* 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) */
 
-    section[data-testid="stSidebar"] .val-actions .stButton:nth-of-type(1) button { background: #f59e0b !important; } /* Preview - orange */
-    section[data-testid="stSidebar"] .val-actions .stButton:nth-of-type(2) button { background: #0d6efd !important; } /* Predict - blue */
+    .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) */
 
-    /* Disabled look */
-    section[data-testid="stSidebar"] .dev-actions .stButton button:disabled,
-    section[data-testid="stSidebar"] .val-actions .stButton button:disabled { filter: grayscale(40%); opacity:.6; }
+    .dev-actions .stButton button:disabled,
+    .val-actions .stButton button:disabled{ filter: grayscale(40%); opacity:.6; }
     </style>
     """,
     unsafe_allow_html=True
 )
 
-# =========================
-# Helpers
-# =========================
+# ---------------- Helpers ----------------
 try:
     dialog = st.dialog
 except AttributeError:
@@ -132,7 +129,20 @@ def inline_logo(path="logo.png") -> str:
     except Exception:
         return ""
 
-# -------- Plotting (Plotly first, Matplotlib fallback) --------
+def export_workbook(sheets_dict: dict[str, pd.DataFrame], summary_df: pd.DataFrame|None):
+    try:
+        import openpyxl  # noqa
+    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):
     lo = float(np.nanmin([actual.min(), pred.min()]))
     hi = float(np.nanmax([actual.max(), pred.max()]))
@@ -149,10 +159,8 @@ def cross_plotly(actual, pred, title):
         mode="lines", line=dict(dash="dash", width=1.5, color=COLORS["ref"]),
         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_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")
     return fig
@@ -160,11 +168,9 @@ def cross_plotly(actual, pred, title):
 def track_plotly(df, include_actual=True):
     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
+        y = df[depth_col]; y_label = depth_col
     else:
-        y = np.arange(1, len(df) + 1)
-        y_label = "Point Index"
+        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",
@@ -181,10 +187,8 @@ def track_plotly(df, include_actual=True):
         ))
     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)
-    )
+    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))
     return fig
 
 def make_index_tracks_plotly(df: pd.DataFrame, cols: list[str]):
@@ -196,12 +200,13 @@ def make_index_tracks_plotly(df: pd.DataFrame, cols: list[str]):
         fig.update_layout(height=200, margin=dict(l=10,r=10,t=10,b=10))
         return fig
     n = len(cols)
-    fig = make_subplots(rows=1, cols=n, shared_y=True, horizontal_spacing=0.05)
+    fig = make_subplots(rows=1, cols=n, shared_yaxes=True, horizontal_spacing=0.05)  # <-- FIX
     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),
+                       hovertemplate=f"{col}: "+"%{x:.2f}<br>Index: %{y}<extra></extra>",
+                       name=col, showlegend=False),
             row=1, col=i
         )
         fig.update_xaxes(title_text=col, side="top", row=1, col=i)
@@ -209,7 +214,7 @@ def make_index_tracks_plotly(df: pd.DataFrame, cols: list[str]):
     fig.update_layout(height=650, margin=dict(l=10, r=10, t=40, b=10))
     return fig
 
-# Fallbacks (kept if Plotly missing)
+# 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"])
@@ -239,7 +244,6 @@ def depth_or_index_track_mpl(df, title=None, include_actual=True):
     ax.legend(loc="best")
     return fig
 
-# ---------- Preview modal helpers ----------
 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:
@@ -248,6 +252,7 @@ def stats_table(df: pd.DataFrame, cols: list[str]) -> pd.DataFrame:
     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:
@@ -288,11 +293,8 @@ def preview_modal_val(book: dict[str, pd.DataFrame], feature_cols: list[str]):
     with t2:
         st.dataframe(stats_table(df, feature_cols), use_container_width=True)
 
-# =========================
-# Model presence
-# =========================
+# ---------------- 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:
@@ -334,23 +336,19 @@ else:
     infer = infer_features_from_model(model)
     if infer: FEATURES = infer
 
-# =========================
-# Session state
-# =========================
+# ---------------- 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_preview_request": False,
-    "dev_file_bytes": b"", "dev_file_name": "", "dev_file_rows": 0, "dev_file_cols": 0,
+    "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 header
-# =========================
+# ---------------- Hero ----------------
 st.markdown(
     f"""
     <div class="st-hero">
@@ -361,12 +359,10 @@ st.markdown(
       </div>
     </div>
     """,
-    unsafe_allow_html=True,
+    unsafe_allow_html=True
 )
 
-# =========================
-# INTRO PAGE
-# =========================
+# ---------------- INTRO ----------------
 if st.session_state.app_step == "intro":
     st.header("Welcome!")
     st.markdown("This software is developed by *Smart Thinking AI-Solutions Team* to estimate UCS from drilling data.")
@@ -388,20 +384,15 @@ if st.session_state.app_step == "intro":
     if st.button("Start Showcase", type="primary", key="start_showcase"):
         st.session_state.app_step = "dev"; st.rerun()
 
-# =========================
-# MODEL DEVELOPMENT
-# =========================
+# ---------------- DEVELOPMENT ----------------
 if st.session_state.app_step == "dev":
     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")
 
-    # Persist upload
     if train_test_file is not None:
-        try:
-            file_bytes = train_test_file.getvalue(); size = len(file_bytes)
-        except Exception:
-            file_bytes = b""; size = 0
+        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
@@ -422,7 +413,6 @@ if st.session_state.app_step == "dev":
             f"{st.session_state.dev_file_rows} rows × {st.session_state.dev_file_cols} cols"
         )
 
-    # Button group with wrapper to color via CSS
     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)
@@ -432,7 +422,6 @@ if st.session_state.app_step == "dev":
     if proceed_clicked and st.session_state.dev_ready:
         st.session_state.app_step = "predict"; st.rerun()
 
-    # Pinned helper
     helper_top = st.container()
     with helper_top:
         st.subheader("Model Development")
@@ -486,7 +475,6 @@ if st.session_state.app_step == "dev":
             status.update(label="Done ✓", state="complete"); toast("Model run complete 🚀")
             st.rerun()
 
-    # Results
     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:
@@ -545,9 +533,7 @@ if st.session_state.app_step == "dev":
         except RuntimeError as e:
             st.warning(str(e))
 
-# =========================
-# PREDICTION (Validation)
-# =========================
+# ---------------- PREDICTION ----------------
 if st.session_state.app_step == "predict":
     st.sidebar.header("Prediction (Validation)")
     validation_file = st.sidebar.file_uploader("Upload Validation Excel", type=["xlsx","xls"], key="val_upload")
@@ -664,9 +650,7 @@ if st.session_state.app_step == "predict":
         except RuntimeError as e:
             st.warning(str(e))
 
-# =========================
-# Footer
-# =========================
+# ---------------- Footer ----------------
 st.markdown("---")
 st.markdown(
     """