app.py

# app.py — GeoMate V2 (single-file Streamlit app)
# ------------------------------------------------
# Paste this entire file into your Hugging Face Space as app.py
# ------------------------------------------------

from __future__ import annotations
import os
import io
import json
import math
from math import floor
from datetime import datetime
from typing import Any, Dict, List, Tuple, Optional

# Streamlit must be imported early and set_page_config must be first Streamlit command
import streamlit as st

# Page config (first Streamlit call)
st.set_page_config(page_title="GeoMate V2", page_icon="🌍", layout="wide", initial_sidebar_state="expanded")

# Standard libs for data & plotting
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import traceback

# Optional heavy libs — import safely
try:
    import faiss    # may fail in some environments
except Exception:
    faiss = None

try:
    import reportlab
    from reportlab.lib import colors
    from reportlab.lib.pagesizes import A4
    from reportlab.lib.units import mm
    from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, PageBreak
    from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
    REPORTLAB_OK = True
except Exception:
    REPORTLAB_OK = False

# FPDF fallback for simple PDFs
try:
    from fpdf import FPDF
    FPDF_OK = True
except Exception:
    FPDF_OK = False

# Groq client optional
try:
    from groq import Groq
    GROQ_OK = True
except Exception:
    GROQ_OK = False

# Earth Engine & geemap optional
try:
    import ee
    import geemap
    EE_OK = True
except Exception:
    ee = None
    geemap = None
    EE_OK = False

# OCR optional (pytesseract / easyocr)
OCR_TESSERACT = False
try:
    import pytesseract
    from PIL import Image
    OCR_TESSERACT = True
except Exception:
    OCR_TESSERACT = False

# Helpful alias for session state
ss = st.session_state

# -------------------------
# Helper: safe rerun
# -------------------------
def safe_rerun():
    """Try to rerun app. Prefer st.rerun(); fallback to experimental or simple stop."""
    try:
        st.rerun()
    except Exception:
        try:
            st.experimental_rerun()
        except Exception:
            # last resort: stop execution so user can click/refresh
            st.stop()

# -------------------------
# Secrets: HF-friendly access
# -------------------------
def get_secret(name: str) -> Optional[str]:
    """Read secrets from environment variables first, then streamlit secrets."""
    v = os.environ.get(name)
    if v:
        return v
    try:
        v2 = st.secrets.get(name)
        if v2:
            # st.secrets may contain nested dicts (for JSON keys)
            if isinstance(v2, dict) or isinstance(v2, list):
                # convert to JSON string
                return json.dumps(v2)
            return str(v2)
    except Exception:
        pass
    return None

# Required secret names (user asked)
GROQ_KEY = get_secret("GROQ_API_KEY")
SERVICE_ACCOUNT = get_secret("SERVICE_ACCOUNT")
EARTH_ENGINE_KEY = get_secret("EARTH_ENGINE_KEY")  # JSON content (string) or path; optional
# We'll consider them optional; pages will show clear errors if missing
HAVE_GROQ = bool(GROQ_KEY and GROQ_OK)
HAVE_SERVICE_ACCOUNT = bool(SERVICE_ACCOUNT)
HAVE_EE_KEY = bool(EARTH_ENGINE_KEY)
# EE readiness flag we'll set after attempted init
EE_READY = False

# Attempt to initialize Earth Engine if credentials provided and ee module available
if EE_OK and (SERVICE_ACCOUNT or EARTH_ENGINE_KEY):
    try:
        # If EARTH_ENGINE_KEY is a JSON string, write to temp file and use service account auth
        key_file = None
        if EARTH_ENGINE_KEY:
            # attempt to parse as json content
            try:
                parsed = json.loads(EARTH_ENGINE_KEY)
                # write to /tmp/geomate_ee_key.json
                key_file = "/tmp/geomate_ee_key.json"
                with open(key_file, "w") as f:
                    json.dump(parsed, f)
            except Exception:
                # maybe EARTH_ENGINE_KEY is a path already on disk (unlikely on HF)
                key_file = EARTH_ENGINE_KEY if os.path.exists(EARTH_ENGINE_KEY) else None
        if key_file and SERVICE_ACCOUNT:
            try:
                # Use oauth2client if available
                from oauth2client.service_account import ServiceAccountCredentials
                creds = ServiceAccountCredentials.from_json_keyfile_name(key_file, scopes=['https://www.googleapis.com/auth/earthengine'])
                ee.Initialize(creds)
                EE_READY = True
            except Exception:
                # fallback: try ee.Initialize with service_account
                try:
                    # This call may fail depending on ee versions; keep safe
                    ee.Initialize()
                    EE_READY = True
                except Exception:
                    EE_READY = False
        else:
            # try simple ee.Initialize()
            try:
                ee.Initialize()
                EE_READY = True
            except Exception:
                EE_READY = False
    except Exception:
        EE_READY = False
else:
    EE_READY = False

# -------------------------
# Session state initialization
# -------------------------
# site_descriptions: list of dicts (max 4). We'll store as list for site ordering but also index by name.
if "site_descriptions" not in ss:
    # default single site
    ss["site_descriptions"] = []
if "active_site_index" not in ss:
    ss["active_site_index"] = 0
if "llm_model" not in ss:
    ss["llm_model"] = "llama3-8b-8192"
if "page" not in ss:
    ss["page"] = "Landing"
if "rag_memory" not in ss:
    ss["rag_memory"] = {}  # per-site chat memory
if "classifier_states" not in ss:
    ss["classifier_states"] = {}  # per-site classifier step & inputs

# Utility: create default site structure
def make_empty_site(name: str = "Site 1") -> dict:
    return {
        "Site Name": name,
        "Site Coordinates": "",
        "lat": None,
        "lon": None,
        "Load Bearing Capacity": None,
        "Skin Shear Strength": None,
        "Relative Compaction": None,
        "Rate of Consolidation": None,
        "Nature of Construction": None,
        "Soil Profile": None,
        "Flood Data": None,
        "Seismic Data": None,
        "GSD": None,
        "USCS": None,
        "AASHTO": None,
        "GI": None,
        "classifier_inputs": {},
        "classifier_decision_path": "",
        "chat_history": [],
        "report_convo_state": 0,
        "map_snapshot": None,
        "classifier_state": 0,
        "classifier_chat": []
    }

# ensure at least one site exists
if len(ss["site_descriptions"]) == 0:
    ss["site_descriptions"].append(make_empty_site("Home"))

# -------------------------
# Sidebar: site management & LLM model selection
# -------------------------
from streamlit_option_menu import option_menu

def sidebar_ui():
    st.sidebar.markdown("<div style='text-align:center'><h2 style='color:#FF8C00;margin:6px 0'>GeoMate V2</h2></div>", unsafe_allow_html=True)
    st.sidebar.markdown("---")

    # Model selector (persist)
    st.sidebar.subheader("LLM Model")
    model_options = ["llama3-8b-8192", "gemma-7b-it", "mixtral-8x7b-32768"]
    ss["llm_model"] = st.sidebar.selectbox("Select LLM model", model_options, index=model_options.index(ss.get("llm_model","llama3-8b-8192")), key="llm_select")

    st.sidebar.markdown("---")
    st.sidebar.subheader("Project Sites (max 4)")
    # show sites and allow add/remove, choose active site
    colA, colB = st.sidebar.columns([2,1])
    with colA:
        new_site_name = st.text_input("New site name", value="", key="new_site_name_input")
    with colB:
        if st.button("➕", key="add_site_btn"):
            # add new site up to 4
            if len(ss["site_descriptions"]) >= 4:
                st.sidebar.warning("Maximum 4 sites allowed.")
            else:
                name = new_site_name.strip() or f"Site {len(ss['site_descriptions'])+1}"
                ss["site_descriptions"].append(make_empty_site(name))
                ss["active_site_index"] = len(ss["site_descriptions"]) - 1
                safe_rerun()

    # list of site names
    site_names = [s["Site Name"] for s in ss["site_descriptions"]]
    idx = st.sidebar.radio("Active Site", options=list(range(len(site_names))), format_func=lambda i: site_names[i], index=ss.get("active_site_index",0), key="active_site_radio")
    ss["active_site_index"] = idx

    # Remove site
    if st.sidebar.button("🗑️ Remove Active Site", key="remove_site_btn"):
        if len(ss["site_descriptions"]) <= 1:
            st.sidebar.warning("Cannot remove last site.")
        else:
            ss["site_descriptions"].pop(ss["active_site_index"])
            ss["active_site_index"] = max(0, ss["active_site_index"] - 1)
            safe_rerun()

    st.sidebar.markdown("---")
    st.sidebar.subheader("Active Site JSON")
    with st.sidebar.expander("Show site JSON", expanded=False):
        st.code(json.dumps(ss["site_descriptions"][ss["active_site_index"]], indent=2), language="json")

    st.sidebar.markdown("---")
    # Secrets indicator (not blocking)
    st.sidebar.subheader("Service Status")
    col1, col2 = st.sidebar.columns(2)
    col1.markdown("LLM:")
    col2.markdown("✅" if HAVE_GROQ else "⚠️ (no Groq)")
    col1.markdown("Earth Engine:")
    col2.markdown("✅" if EE_READY else "⚠️ (not initialized)")
    st.sidebar.markdown("---")
    # Navigation menu
    pages = ["Landing", "Soil Recognizer", "Soil Classifier", "GSD Curve", "Locator", "GeoMate Ask", "Reports"]
    icons = ["house", "image", "flask", "bar-chart", "geo-alt", "robot", "file-earmark-text"]
    choice = option_menu(None, pages, icons=icons, menu_icon="cast", default_index=pages.index(ss.get("page","Landing")), orientation="vertical", styles={
        "container": {"padding": "0px"},
        "nav-link-selected": {"background-color": "#FF7A00"},
    })
    if choice and choice != ss.get("page"):
        ss["page"] = choice
        safe_rerun()

# -------------------------
# Landing UI
# -------------------------
def landing_ui():
    st.markdown(
        """
        <style>
        .hero {
            background: linear-gradient(135deg,#0f0f0f 0%, #060606 100%);
            border-radius: 14px;
            padding: 20px;
            border: 1px solid rgba(255,122,0,0.08);
        }
        .glow-btn {
            background: linear-gradient(90deg,#ff7a00,#ff3a3a);
            color: white;
            padding: 10px 18px;
            border-radius: 10px;
            font-weight:700;
            box-shadow: 0 6px 24px rgba(255,122,0,0.12);
            border: none;
        }
        </style>
        """, unsafe_allow_html=True)
    st.markdown("<div class='hero'>", unsafe_allow_html=True)
    st.markdown("<h1 style='color:#FF8C00;margin:0'>🌍 GeoMate V2</h1>", unsafe_allow_html=True)
    st.markdown("<p style='color:#ddd'>AI copilot for geotechnical engineering — soil recognition, classification, locator, RAG, professional reports.</p>", unsafe_allow_html=True)
    st.markdown("</div>", unsafe_allow_html=True)
    st.markdown("---")
    st.write("Quick actions")
    c1, c2, c3 = st.columns(3)
    if c1.button("🖼️ Soil Recognizer"):
        ss["page"] = "Soil Recognizer"; safe_rerun()
    if c2.button("🧪 Soil Classifier"):
        ss["page"] = "Soil Classifier"; safe_rerun()
    if c3.button("📊 GSD Curve"):
        ss["page"] = "GSD Curve"; safe_rerun()
    c4, c5, c6 = st.columns(3)
    if c4.button("🌍 Locator"):
        ss["page"] = "Locator"; safe_rerun()
    if c5.button("🤖 GeoMate Ask"):
        ss["page"] = "GeoMate Ask"; safe_rerun()
    if c6.button("📑 Reports"):
        ss["page"] = "Reports"; safe_rerun()

# -------------------------
# Utility: active site helpers
# -------------------------
def active_site() -> Tuple[str, dict]:
    idx = ss.get("active_site_index", 0)
    idx = max(0, min(idx, len(ss["site_descriptions"]) - 1))
    ss["active_site_index"] = idx
    site = ss["site_descriptions"][idx]
    return idx, site

def save_site_field(field: str, value: Any):
    idx, site = active_site()
    site[field] = value
    ss["site_descriptions"][idx] = site

# -------------------------
# USCS & AASHTO verbatim logic (function)
# -------------------------
# Uses exact logic from your script and mapping of descriptor strings to numbers
def uscs_aashto_from_inputs(inputs: Dict[str,Any]) -> Tuple[str,str,str,int,Dict[str,str]]:
    """
    Return: (result_text, uscs_symbol, aashto_symbol, GI, char_summary)
    """
    # Engineering characteristics dictionary (detailed-ish)
    ENGINEERING_CHARACTERISTICS = {
        "Gravel": {
            "Settlement": "None",
            "Quicksand": "Impossible",
            "Frost-heaving": "None",
            "Groundwater_lowering": "Possible",
            "Cement_grouting": "Possible",
            "Silicate_bitumen_injections": "Unsuitable",
            "Compressed_air": "Possible"
        },
        "Coarse sand": {"Settlement":"None","Quicksand":"Impossible","Frost-heaving":"None"},
        "Medium sand": {"Settlement":"None","Quicksand":"Unlikely"},
        "Fine sand": {"Settlement":"None","Quicksand":"Liable"},
        "Silt": {"Settlement":"Occurs","Quicksand":"Liable","Frost-heaving":"Occurs"},
        "Clay": {"Settlement":"Occurs","Quicksand":"Impossible"}
    }

    opt = str(inputs.get("opt","n")).lower()
    if opt == 'y':
        uscs = "Pt"
        uscs_expl = "Peat / organic soil — compressible, high organic content; poor engineering properties for load-bearing without special treatment."
        aashto = "Organic (special handling)"
        GI = 0
        result_text = f"According to USCS, the soil is {uscs} — {uscs_expl}\nAccording to AASHTO, the soil is {aashto}."
        return result_text, uscs, aashto, GI, {"summary":"Organic peat: large settlement, low strength."}

    # read numeric inputs safely
    P2 = float(inputs.get("P2", 0.0))
    P4 = float(inputs.get("P4", 0.0))
    D60 = float(inputs.get("D60", 0.0))
    D30 = float(inputs.get("D30", 0.0))
    D10 = float(inputs.get("D10", 0.0))
    LL = float(inputs.get("LL", 0.0))
    PL = float(inputs.get("PL", 0.0))
    PI = LL - PL

    Cu = (D60 / D10) if (D10 > 0 and D60 > 0) else 0
    Cc = ((D30 ** 2) / (D10 * D60)) if (D10 > 0 and D30 > 0 and D60 > 0) else 0

    uscs = "Unknown"; uscs_expl = ""
    if P2 <= 50:
        # Coarse-Grained
        if P4 <= 50:
            # Gravels
            if Cu and Cc:
                if Cu >= 4 and 1 <= Cc <= 3:
                    uscs, uscs_expl = "GW", "Well-graded gravel (good engineering properties, high strength, good drainage)."
                else:
                    uscs, uscs_expl = "GP", "Poorly-graded gravel (less favorable gradation)."
            else:
                if PI < 4 or PI < 0.73 * (LL - 20):
                    uscs, uscs_expl = "GM", "Silty gravel (fines may reduce permeability and strength)."
                elif PI > 7 and PI > 0.73 * (LL - 20):
                    uscs, uscs_expl = "GC", "Clayey gravel (higher plasticity)."
                else:
                    uscs, uscs_expl = "GM-GC", "Gravel with mixed silt/clay fines."
        else:
            # Sands
            if Cu and Cc:
                if Cu >= 6 and 1 <= Cc <= 3:
                    uscs, uscs_expl = "SW", "Well-graded sand (good compaction and drainage)."
                else:
                    uscs, uscs_expl = "SP", "Poorly-graded sand (uniform or gap-graded)."
            else:
                if PI < 4 or PI <= 0.73 * (LL - 20):
                    uscs, uscs_expl = "SM", "Silty sand (low-plasticity fines)."
                elif PI > 7 and PI > 0.73 * (LL - 20):
                    uscs, uscs_expl = "SC", "Clayey sand (clayey fines present)."
                else:
                    uscs, uscs_expl = "SM-SC", "Transition between silty sand and clayey sand."
    else:
        # Fine-grained soils
        nDS = int(inputs.get("nDS", 5))
        nDIL = int(inputs.get("nDIL", 6))
        nTG = int(inputs.get("nTG", 6))
        if LL < 50:
            if 20 <= LL < 50 and PI <= 0.73 * (LL - 20):
                if nDS == 1 or nDIL == 3 or nTG == 3:
                    uscs, uscs_expl = "ML", "Silt (low plasticity)."
                elif nDS == 3 or nDIL == 3 or nTG == 3:
                    uscs, uscs_expl = "OL", "Organic silt (low plasticity)."
                else:
                    uscs, uscs_expl = "ML-OL", "Mixed silt/organic silt."
            elif 10 <= LL <= 30 and 4 <= PI <= 7 and PI > 0.72 * (LL - 20):
                if nDS == 1 or nDIL == 1 or nTG == 1:
                    uscs, uscs_expl = "ML", "Silt."
                elif nDS == 2 or nDIL == 2 or nTG == 2:
                    uscs, uscs_expl = "CL", "Clay (low plasticity)."
                else:
                    uscs, uscs_expl = "ML-CL", "Mixed silt/clay."
            else:
                uscs, uscs_expl = "CL", "Clay (low plasticity)."
        else:
            if PI < 0.73 * (LL - 20):
                if nDS == 3 or nDIL == 4 or nTG == 4:
                    uscs, uscs_expl = "MH", "Silt (high plasticity)"
                elif nDS == 2 or nDIL == 2 or nTG == 4:
                    uscs, uscs_expl = "OH", "Organic silt/clay (high plasticity)"
                else:
                    uscs, uscs_expl = "MH-OH", "Mixed high-plasticity silt/organic"
            else:
                uscs, uscs_expl = "CH", "Clay (high plasticity)"

    # AASHTO logic
    if P2 <= 35:
        if P2 <= 15 and P4 <= 30 and PI <= 6:
            aashto = "A-1-a"
        elif P2 <= 25 and P4 <= 50 and PI <= 6:
            aashto = "A-1-b"
        elif P2 <= 35 and P4 > 0:
            if LL <= 40 and PI <= 10:
                aashto = "A-2-4"
            elif LL >= 41 and PI <= 10:
                aashto = "A-2-5"
            elif LL <= 40 and PI >= 11:
                aashto = "A-2-6"
            elif LL >= 41 and PI >= 11:
                aashto = "A-2-7"
            else:
                aashto = "A-2"
        else:
            aashto = "A-3"
    else:
        if LL <= 40 and PI <= 10:
            aashto = "A-4"
        elif LL >= 41 and PI <= 10:
            aashto = "A-5"
        elif LL <= 40 and PI >= 11:
            aashto = "A-6"
        else:
            aashto = "A-7-5" if PI <= (LL - 30) else "A-7-6"

    # Group Index (GI)
    a = P2 - 35
    a = 0 if a < 0 else (40 if a > 40 else a)
    b = P2 - 15
    b = 0 if b < 0 else (40 if b > 40 else b)
    c = LL - 40
    c = 0 if c < 0 else (20 if c > 20 else c)
    d = PI - 10
    d = 0 if d < 0 else (20 if d > 20 else d)
    GI = floor(0.2 * a + 0.005 * a * c + 0.01 * b * d)

    aashto_expl = f"{aashto} (GI = {GI})"

    # Choose characteristic summary based on USCS family
    char_summary = ENGINEERING_CHARACTERISTICS.get("Silt", {})
    if uscs.startswith(("G", "S")):
        char_summary = ENGINEERING_CHARACTERISTICS.get("Coarse sand", ENGINEERING_CHARACTERISTICS.get("Gravel"))
    if uscs.startswith(("M", "C", "O", "H")):
        char_summary = ENGINEERING_CHARACTERISTICS.get("Silt", {})

    result_text = f"According to USCS, the soil is **{uscs}** — {uscs_expl}\n\nAccording to AASHTO, the soil is **{aashto_expl}**\n\nEngineering characteristics summary:\n"
    for k, v in char_summary.items():
        result_text += f"- {k}: {v}\n"

    return result_text, uscs, aashto, GI, char_summary

# -------------------------
# GSD Curve page (separate)
# -------------------------
def gsd_curve_ui():
    st.header("📊 Grain Size Distribution (GSD) Curve")
    _, site = active_site()
    st.markdown(f"**Active site:** {site['Site Name']}")

    st.info("Upload a CSV (two columns: diameter_mm, percent_passing) or enter diameters and % passing manually. I will compute D10, D30, D60 and save them for the active site.")

    col_up, col_manual = st.columns([1,1])
    uploaded = None
    sieve = None; passing = None
    with col_up:
        uploaded = st.file_uploader("Upload CSV (diameter_mm, percent_passing)", type=["csv","txt"], key=f"gsd_upload_{site['Site Name']}")
        if uploaded:
            try:
                df = pd.read_csv(uploaded)
                # heuristic: first numeric column = diameter, second numeric = percent passing
                numeric_cols = [c for c in df.columns if np.issubdtype(df[c].dtype, np.number)]
                if len(numeric_cols) >= 2:
                    sieve = df[numeric_cols[0]].values.astype(float)
                    passing = df[numeric_cols[1]].values.astype(float)
                else:
                    # try to coerce
                    sieve = df.iloc[:,0].astype(float).values
                    passing = df.iloc[:,1].astype(float).values
                st.success("CSV loaded.")
            except Exception as e:
                st.error(f"Error reading CSV: {e}")
                sieve = passing = None

    with col_manual:
        diam_text = st.text_area("Diameters (mm) comma-separated (e.g. 75,50,37.5,...)", key=f"gsd_diams_{site['Site Name']}")
        pass_text = st.text_area("% Passing comma-separated (same order)", key=f"gsd_pass_{site['Site Name']}")
        if diam_text.strip() and pass_text.strip():
            try:
                sieve = np.array([float(x.strip()) for x in diam_text.split(",") if x.strip()])
                passing = np.array([float(x.strip()) for x in pass_text.split(",") if x.strip()])
            except Exception as e:
                st.error(f"Invalid manual input: {e}")
                sieve = passing = None

    if sieve is None or passing is None:
        st.info("Provide GSD data above to compute D-values.")
        return

    # Sort descending by sieve
    order = np.argsort(-sieve)
    sieve = sieve[order]; passing = passing[order]
    # Ensure percent is in 0-100
    if np.any(passing < 0) or np.any(passing > 100):
        st.warning("Some % passing values are outside 0-100. Please verify.")
    # Interpolate percent -> diameter (need percent increasing)
    percent = passing.copy()
    if not np.all(np.diff(percent) >= 0):
        percent = percent[::-1]; sieve = sieve[::-1]
    # Ensure arrays are floats
    percent = percent.astype(float); sieve = sieve.astype(float)
    # interpolation function (percent -> diameter)
    def interp_d(pct: float) -> Optional[float]:
        try:
            # xp must be increasing
            return float(np.interp(pct, percent, sieve))
        except Exception:
            return None
    D10 = interp_d(10.0); D30 = interp_d(30.0); D60 = interp_d(60.0)
    Cu = (D60 / D10) if (D10 and D60 and D10>0) else None
    Cc = (D30**2)/(D10*D60) if (D10 and D30 and D60 and D10>0) else None

    # Plot
    fig, ax = plt.subplots(figsize=(7,4))
    ax.plot(sieve, passing, marker='o', label="% Passing")
    ax.set_xscale('log')
    ax.invert_xaxis()
    ax.set_xlabel("Particle diameter (mm) — log scale")
    ax.set_ylabel("% Passing")
    if D10: ax.axvline(D10, color='orange', linestyle='--', label=f"D10={D10:.4g} mm")
    if D30: ax.axvline(D30, color='red', linestyle='--', label=f"D30={D30:.4g} mm")
    if D60: ax.axvline(D60, color='blue', linestyle='--', label=f"D60={D60:.4g} mm")
    ax.grid(True, which='both', linestyle='--', linewidth=0.4)
    ax.legend()
    st.pyplot(fig)

    # Save to active site
    idx, sdict = active_site()
    sdict["GSD"] = {
        "sieve_mm": sieve.tolist(),
        "percent_passing": passing.tolist(),
        "D10": float(D10) if D10 is not None else None,
        "D30": float(D30) if D30 is not None else None,
        "D60": float(D60) if D60 is not None else None,
        "Cu": float(Cu) if Cu is not None else None,
        "Cc": float(Cc) if Cc is not None else None
    }
    ss["site_descriptions"][idx] = sdict
    st.success(f"Saved GSD to site: D10={D10}, D30={D30}, D60={D60}")
    if st.button("Copy D-values to Soil Classifier inputs (for this site)", key=f"copy_d_to_cls_{sdict['Site Name']}"):
        ss.setdefault("classifier_states", {})
        ss["classifier_states"].setdefault(sdict["Site Name"], {"step":0, "inputs":{}})
        ss["classifier_states"][sdict["Site Name"]]["inputs"]["D10"] = float(D10) if D10 is not None else 0.0
        ss["classifier_states"][sdict["Site Name"]]["inputs"]["D30"] = float(D30) if D30 is not None else 0.0
        ss["classifier_states"][sdict["Site Name"]]["inputs"]["D60"] = float(D60) if D60 is not None else 0.0
        st.info("Copied. Go to Soil Classifier page and continue.")
    st.markdown("---")
    st.caption("Tip: If your classifier asks for D-values and you don't have them, compute them here and copy them back.")

# -------------------------
# Soil Classifier (chatbot-style; smart input flow)
# -------------------------
def soil_classifier_ui():
    st.header("🧪 Soil Classifier — Chatbot-style (smart inputs)")
    idx, sdict = active_site()
    site_name = sdict["Site Name"]

    # Prepare classifier state container per-site
    ss.setdefault("classifier_states", {})
    state = ss["classifier_states"].setdefault(site_name, {"step":0, "inputs":{}})

    # Pre-populate inputs from site GSD if present
    if sdict.get("GSD"):
        g = sdict["GSD"]
        state["inputs"].setdefault("D10", g.get("D10", 0.0))
        state["inputs"].setdefault("D30", g.get("D30", 0.0))
        state["inputs"].setdefault("D60", g.get("D60", 0.0))

    # Dropdown exact strings mapping per your requested table
    dil_options = [
        "1. Quick to slow",
        "2. None to very slow",
        "3. Slow",
        "4. Slow to none",
        "5. None",
        "6. Null?"
    ]
    tough_options = [
        "1. None",
        "2. Medium",
        "3. Slight?",
        "4. Slight to Medium?",
        "5. High",
        "6. Null?"
    ]
    dry_options = [
        "1. None to slight",
        "2. Medium to high",
        "3. Slight to Medium",
        "4. High to very high",
        "5. Null?"
    ]

    # Mode selector: Both / USCS only / AASHTO only
    if "classifier_mode" not in ss:
        ss["classifier_mode"] = {}
    ss["classifier_mode"].setdefault(site_name, "Both")
    mode = st.selectbox("Classification mode", ["Both", "USCS only", "AASHTO only"], index=["Both","USCS only","AASHTO only"].index(ss["classifier_mode"][site_name]), key=f"mode_{site_name}")
    ss["classifier_mode"][site_name] = mode

    # Step machine:
    step = state.get("step", 0)
    inputs = state.setdefault("inputs", {})

    def goto(n):
        state["step"] = n
        ss["classifier_states"][site_name] = state
        safe_rerun()

    def save_input(key, value):
        inputs[key] = value
        state["inputs"] = inputs
        ss["classifier_states"][site_name] = state

    # Chat style UI: show recent conversation history as simple list
    st.markdown(f"**Active site:** {site_name}")
    st.markdown("---")

    if step == 0:
        st.markdown("**GeoMate:** Hello! I'm the Soil Classifier bot. Shall we begin classification for this site?")
        c1, c2 = st.columns(2)
        if c1.button("Yes — Start", key=f"cls_start_{site_name}"):
            goto(1)
        if c2.button("Cancel", key=f"cls_cancel_{site_name}"):
            st.info("Classifier cancelled.")

    elif step == 1:
        st.markdown("**GeoMate:** Is the soil organic (contains high organic matter, spongy, odour)?")
        c1, c2 = st.columns(2)
        if c1.button("No (inorganic)", key=f"org_no_{site_name}"):
            save_input("opt", "n"); goto(2)
        if c2.button("Yes (organic)", key=f"org_yes_{site_name}"):
            save_input("opt", "y"); goto(12)  # skip to final classification for organic

    elif step == 2:
        st.markdown("**GeoMate:** What is the percentage passing the #200 sieve (0.075 mm)?")
        val = st.number_input("Percentage passing #200 (P2)", value=float(inputs.get("P2",0.0)), min_value=0.0, max_value=100.0, key=f"P2_input_{site_name}", format="%.2f")
        c1, c2, c3 = st.columns([1,1,1])
        if c1.button("Next", key=f"P2_next_{site_name}"):
            save_input("P2", float(val)); goto(3)
        if c2.button("Skip", key=f"P2_skip_{site_name}"):
            save_input("P2", 0.0); goto(3)
        if c3.button("Back", key=f"P2_back_{site_name}"):
            goto(0)

    elif step == 3:
        # decide path based on P2
        P2 = float(inputs.get("P2", 0.0))
        if P2 > 50:
            st.markdown("**GeoMate:** P2 > 50 — fine-grained soil path selected.")
            if st.button("Continue (fine-grained)", key=f"cont_fine_{site_name}"):
                goto(4)
        else:
            st.markdown("**GeoMate:** P2 <= 50 — coarse-grained soil path selected.")
            if st.button("Continue (coarse-grained)", key=f"cont_coarse_{site_name}"):
                goto(6)

    # Fine-grained branch
    elif step == 4:
        st.markdown("**GeoMate:** Enter Liquid Limit (LL).")
        val = st.number_input("Liquid Limit (LL)", value=float(inputs.get("LL",0.0)), min_value=0.0, max_value=200.0, key=f"LL_{site_name}", format="%.2f")
        col1, col2 = st.columns(2)
        if col1.button("Next", key=f"LL_next_{site_name}"):
            save_input("LL", float(val)); goto(5)
        if col2.button("Back", key=f"LL_back_{site_name}"):
            goto(3)

    elif step == 5:
        st.markdown("**GeoMate:** Enter Plastic Limit (PL).")
        val = st.number_input("Plastic Limit (PL)", value=float(inputs.get("PL",0.0)), min_value=0.0, max_value=200.0, key=f"PL_{site_name}", format="%.2f")
        col1, col2 = st.columns(2)
        if col1.button("Next", key=f"PL_next_{site_name}"):
            save_input("PL", float(val)); goto(11)  # go to descriptors step
        if col2.button("Back", key=f"PL_back_{site_name}"):
            goto(4)

    # Coarse branch: ask % passing #4 and D-values option
    elif step == 6:
        st.markdown("**GeoMate:** What is the % passing sieve no. 4 (4.75 mm)?")
        val = st.number_input("% passing #4 (P4)", value=float(inputs.get("P4",0.0)), min_value=0.0, max_value=100.0, key=f"P4_{site_name}", format="%.2f")
        c1, c2, c3 = st.columns([1,1,1])
        if c1.button("Next", key=f"P4_next_{site_name}"):
            save_input("P4", float(val)); goto(7)
        if c2.button("Compute D-values (GSD page)", key=f"P4_gsd_{site_name}"):
            st.info("Please use GSD Curve page to compute D-values and then copy them back to classifier.")
        if c3.button("Back", key=f"P4_back_{site_name}"):
            goto(3)

    elif step == 7:
        st.markdown("**GeoMate:** Do you know D60, D30, D10 diameters (mm)?")
        c1, c2, c3 = st.columns([1,1,1])
        if c1.button("Yes — enter values", key=f"dvals_yes_{site_name}"):
            goto(8)
        if c2.button("No — compute from GSD", key=f"dvals_no_{site_name}"):
            st.info("Use the GSD Curve page and then click 'Copy D-values to Soil Classifier' there.")
        if c3.button("Skip", key=f"dvals_skip_{site_name}"):
            save_input("D60", 0.0); save_input("D30", 0.0); save_input("D10", 0.0); goto(11)

    elif step == 8:
        st.markdown("**GeoMate:** Enter D60 (mm)")
        val = st.number_input("D60 (mm)", value=float(inputs.get("D60",0.0)), min_value=0.0, key=f"D60_{site_name}")
        if st.button("Next", key=f"D60_next_{site_name}"):
            save_input("D60", float(val)); goto(9)
        if st.button("Back", key=f"D60_back_{site_name}"):
            goto(7)

    elif step == 9:
        st.markdown("**GeoMate:** Enter D30 (mm)")
        val = st.number_input("D30 (mm)", value=float(inputs.get("D30",0.0)), min_value=0.0, key=f"D30_{site_name}")
        if st.button("Next", key=f"D30_next_{site_name}"):
            save_input("D30", float(val)); goto(10)
        if st.button("Back", key=f"D30_back_{site_name}"):
            goto(8)

    elif step == 10:
        st.markdown("**GeoMate:** Enter D10 (mm)")
        val = st.number_input("D10 (mm)", value=float(inputs.get("D10",0.0)), min_value=0.0, key=f"D10_{site_name}")
        if st.button("Next", key=f"D10_next_{site_name}"):
            save_input("D10", float(val)); goto(11)
        if st.button("Back", key=f"D10_back_{site_name}"):
            goto(9)

    # descriptors (for fine soils) and finishing
    elif step == 11:
        st.markdown("**GeoMate:** Fine soil descriptors (if applicable). You can skip.")
        sel_dry = st.selectbox("Dry strength", dry_options, index=min(2, len(dry_options)-1), key=f"dry_{site_name}")
        sel_dil = st.selectbox("Dilatancy", dil_options, index=0, key=f"dil_{site_name}")
        sel_tg = st.selectbox("Toughness", tough_options, index=0, key=f"tough_{site_name}")
        col1, col2, col3 = st.columns([1,1,1])
        # map to numeric - map strings to numbers i+1
        dry_map = {dry_options[i]: i+1 for i in range(len(dry_options))}
        dil_map = {dil_options[i]: i+1 for i in range(len(dil_options))}
        tough_map = {tough_options[i]: i+1 for i in range(len(tough_options))}
        if col1.button("Save & Continue", key=f"desc_save_{site_name}"):
            save_input("nDS", dry_map.get(sel_dry, 5))
            save_input("nDIL", dil_map.get(sel_dil, 6))
            save_input("nTG", tough_map.get(sel_tg, 6))
            goto(12)
        if col2.button("Skip descriptors", key=f"desc_skip_{site_name}"):
            save_input("nDS", 5); save_input("nDIL", 6); save_input("nTG", 6); goto(12)
        if col3.button("Back", key=f"desc_back_{site_name}"):
            # go to previous appropriate step
            # if LL in inputs, go to LL/PL steps else go to coarse branch
            if inputs.get("LL") is not None and inputs.get("PL") is not None:
                goto(5)
            else:
                goto(6)

    elif step == 12:
        st.markdown("**GeoMate:** Ready to classify. Review inputs and press **Classify**.")
        st.json(inputs)
        col1, col2 = st.columns([1,1])
        if col1.button("Classify", key=f"classify_now_{site_name}"):
            try:
                res_text, uscs_sym, aashto_sym, GI, char_summary = uscs_aashto_from_inputs(inputs)
                # Save to site dict
                sdict["USCS"] = uscs_sym
                sdict["AASHTO"] = aashto_sym
                sdict["GI"] = GI
                sdict["classifier_inputs"] = inputs
                sdict["classifier_decision_path"] = res_text
                # persist
                ss["site_descriptions"][ss["active_site_index"]] = sdict
                st.success("Classification complete and saved.")
                st.markdown("### Result")
                if mode == "USCS only":
                    st.markdown(f"**USCS:** {uscs_sym}")
                elif mode == "AASHTO only":
                    st.markdown(f"**AASHTO:** {aashto_sym} (GI={GI})")
                else:
                    st.markdown(res_text)
                # offer PDF export
                if FPDF_OK and st.button("Export Classification PDF", key=f"exp_cls_pdf_{site_name}"):
                    pdf = FPDF()
                    pdf.add_page()
                    pdf.set_font("Arial", "B", 14)
                    pdf.cell(0, 8, f"GeoMate Classification — {site_name}", ln=1)
                    pdf.ln(4)
                    pdf.set_font("Arial", "", 11)
                    pdf.multi_cell(0, 7, res_text)
                    pdf.ln(4)
                    pdf.set_font("Arial", "B", 12)
                    pdf.cell(0, 6, "Inputs:", ln=1)
                    pdf.set_font("Arial", "", 10)
                    for k, v in inputs.items():
                        pdf.cell(0,5, f"{k}: {v}", ln=1)
                    fn = f"{site_name.replace(' ','_')}_classification.pdf"
                    pdf.output(fn)
                    with open(fn,"rb") as f:
                        st.download_button("Download PDF", f, file_name=fn, mime="application/pdf")
                # done
            except Exception as e:
                st.error(f"Classification failed: {e}\n{traceback.format_exc()}")
        if col2.button("Back to edit", key=f"back_edit_{site_name}"):
            goto(11)

    else:
        st.warning("Classifier state reset.")
        state["step"] = 0
        ss["classifier_states"][site_name] = state