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

# -------------------------
# Locator page (chat-style)
# -------------------------
def locator_ui():
    st.header("🌍 Locator — define AOI (GeoJSON or coordinates)")
    idx, sdict = active_site()
    site_name = sdict["Site Name"]
    st.markdown(f"**Active site:** {site_name}")
    st.info("You can paste GeoJSON for your area of interest or enter a center point (lat, lon). The app will save the AOI to the active site and try to display a map.")

    # Chat-like: ask for GeoJSON or coordinates
    choice = st.radio("Provide:", ["GeoJSON (polygon)", "Coordinates (lat, lon)"], index=0, key=f"loc_choice_{site_name}")
    if choice.startswith("GeoJSON"):
        geo_text = st.text_area("Paste GeoJSON (Polygon or MultiPolygon)", value=sdict.get("map_snapshot") or "", key=f"geojson_area_{site_name}", height=180)
        if st.button("Save AOI and show map", key=f"save_aoi_{site_name}"):
            if not geo_text.strip():
                st.error("No GeoJSON provided.")
            else:
                try:
                    gj = json.loads(geo_text)
                    sdict["map_snapshot"] = gj
                    ss["site_descriptions"][idx] = sdict
                    st.success("AOI saved.")
                    # Try to display on simple map: compute centroid and show with st.map
                    # Find centroid of polygon(s)
                    def centroid_of_geojson(gj_obj):
                        coords = []
                        if gj_obj.get("type") == "FeatureCollection":
                            for f in gj_obj.get("features", []):
                                geom = f.get("geometry", {})
                                coords.extend(_extract_coords_from_geom(geom))
                        else:
                            geom = gj_obj if "geometry" not in gj_obj else gj_obj["geometry"]
                            coords.extend(_extract_coords_from_geom(geom))
                        if not coords:
                            return None
                        arr = np.array(coords)
                        return float(arr[:,1].mean()), float(arr[:,0].mean())
                    def _extract_coords_from_geom(geom):
                        if not geom:
                            return []
                        t = geom.get("type","")
                        if t == "Polygon":
                            return [tuple(pt) for pt in geom.get("coordinates", [])[0]]
                        if t == "MultiPolygon":
                            pts=[]
                            for poly in geom.get("coordinates", []):
                                pts.extend([tuple(pt) for pt in poly[0]])
                            return pts
                        if t == "Point":
                            return [tuple(geom.get("coordinates",[]))]
                        return []
                    cent = centroid_of_geojson(gj)
                    if cent:
                        latc, lonc = cent
                        sdict["lat"] = latc; sdict["lon"] = lonc
                        ss["site_descriptions"][idx] = sdict
                        st.map(pd.DataFrame({"lat":[latc],"lon":[lonc]}))
                    else:
                        st.info("Could not compute centroid for map preview, but AOI saved.")
                except Exception as e:
                    st.error(f"Invalid GeoJSON: {e}")
    else:
        # Coordinates mode
        lat = st.number_input("Latitude", value=float(sdict.get("lat") or 0.0), key=f"loc_lat_{site_name}")
        lon = st.number_input("Longitude", value=float(sdict.get("lon") or 0.0), key=f"loc_lon_{site_name}")
        if st.button("Save coordinates and show map", key=f"save_coords_{site_name}"):
            sdict["lat"] = float(lat); sdict["lon"] = float(lon)
            ss["site_descriptions"][idx] = sdict
            st.success("Coordinates saved.")
            st.map(pd.DataFrame({"lat":[lat],"lon":[lon]}))

    # If EE is available, offer to fetch raster/time series (placeholder)
    st.markdown("---")
    if EE_READY:
        if st.button("Fetch Earth Engine data (soil profile / climate / flood / seismic) — experimental"):
            st.info("Earth Engine available — fetching (placeholder).")
            # Placeholder: real implementation would call ee.Dataset/time series and store results
            try:
                # example: add a placeholder soil profile
                sdict["Soil Profile"] = "Placeholder Earth Engine soil profile data (EE initialized)."
                sdict["Flood Data"] = "Placeholder flood history (20 years) from EE."
                sdict["Seismic Data"] = "Placeholder seismic history (20 years) from EE."
                ss["site_descriptions"][idx] = sdict
                st.success("Earth Engine data fetched and saved to site (placeholder).")
            except Exception as e:
                st.error(f"EE fetch failed: {e}")
    else:
        st.info("Earth Engine not initialized in this runtime. To enable, set SERVICE_ACCOUNT and EARTH_ENGINE_KEY secrets and ensure earthengine-api is installed.")

# -------------------------
# GeoMate Ask (RAG chatbot) — simplified RAG integration
# -------------------------
def run_llm_completion(prompt: str, model: str = "llama3-8b-8192") -> str:
    """Minimal LLM wrapper: uses Groq if available; else returns a dummy but structured answer."""
    if GROQ_OK and GROQ_KEY:
        try:
            client = Groq(api_key=GROQ_KEY)
            comp = client.chat.completions.create(model=model, messages=[{"role":"user", "content": prompt}], temperature=0.2, max_tokens=800)
            return comp.choices[0].message.content
        except Exception as e:
            return f"(Groq error) {e}"
    else:
        # Dummy local response for demonstration
        return f"(Dummy LLM) I received your prompt and would reply here. Model: {model}\n\nPrompt excerpt:\n{prompt[:400]}"

def rag_ui():
    st.header("🤖 GeoMate Ask — RAG Chatbot (per-site memory)")
    idx, sdict = active_site()
    site_name = sdict["Site Name"]
    st.markdown(f"**Active site:** {site_name}")
    # Prepare chat history per site
    ss.setdefault("rag_memory", {})
    chat = ss["rag_memory"].setdefault(site_name, [])
    # Show chat
    for turn in chat[-40:]:
        role, text = turn.get("role"), turn.get("text")
        if role == "user":
            st.markdown(f"**You:** {text}")
        else:
            st.markdown(f"**GeoMate:** {text}")

    # user input
    user_prompt = st.text_input("Ask GeoMate (technical)", key=f"rag_input_{site_name}")
    col1, col2 = st.columns([3,1])
    if col1.button("Send", key=f"rag_send_{site_name}") and user_prompt.strip():
        # append user
        chat.append({"role":"user", "text":user_prompt, "ts": datetime.now().isoformat()})
        ss["rag_memory"][site_name] = chat
        # Build RAG prompt using stored site data + user prompt
        context = {"site": sdict}
        full_prompt = f"Site data (json):\n{json.dumps(context, indent=2)}\n\nUser question:\n{user_prompt}\nPlease answer technically."
        # Run LLM
        with st.spinner("Running LLM..."):
            resp = run_llm_completion(full_prompt, model=ss.get("llm_model"))
        # Append bot response
        chat.append({"role":"assistant", "text":resp, "ts": datetime.now().isoformat()})
        ss["rag_memory"][site_name] = chat
        # Intelligent extraction: try to pick up numeric engineering fields (simple heuristics)
        update_site_description_from_chat(resp, site_name)
        safe_rerun()

# small NLP-ish extractor placeholder
def update_site_description_from_chat(text: str, site_name: str):
    """Naive extraction: looks for keywords like 'bearing' and a numeric value followed by units."""
    idx = None
    for i, s in enumerate(ss["site_descriptions"]):
        if s["Site Name"] == site_name:
            idx = i; break
    if idx is None:
        return
    # naive patterns
    lowered = text.lower()
    site = ss["site_descriptions"][idx]
    # look for 'bearing' followed by number (psf, kpa)
    import re
    m = re.search(r"bearing.*?([0-9]{2,6})\s*(psf|kpa|kpa\.)?", lowered)
    if m:
        val = m.group(1)
        site["Load Bearing Capacity"] = m.group(0)
        ss["site_descriptions"][idx] = site

# -------------------------
# Reports page (two types)
# -------------------------
def build_classification_pdf_bytes(site_dict: dict) -> bytes:
    """Return bytes of a classification-only PDF for a single site. Try ReportLab then FPDF fallback."""
    text = site_dict.get("classifier_decision_path") or "No classification decision path available."
    inputs = site_dict.get("classifier_inputs", {})
    title = f"GeoMate Classification Report — {site_dict['Site Name']}"
    # Try ReportLab
    if REPORTLAB_OK:
        buf = io.BytesIO()
        doc = SimpleDocTemplate(buf, pagesize=A4)
        styles = getSampleStyleSheet()
        elems = []
        elems.append(Paragraph(title, styles["Title"]))
        elems.append(Spacer(1,6))
        elems.append(Paragraph("Classification result and explanation:", styles["Heading2"]))
        elems.append(Paragraph(text.replace("\n","<br/>"), styles["BodyText"]))
        elems.append(Spacer(1,6))
        elems.append(Paragraph("Inputs:", styles["Heading3"]))
        for k,v in inputs.items():
            elems.append(Paragraph(f"{k}: {v}", styles["BodyText"]))
        doc.build(elems)
        pdf_bytes = buf.getvalue(); buf.close()
        return pdf_bytes
    elif FPDF_OK:
        pdf = FPDF()
        pdf.add_page()
        pdf.set_font("Arial", "B", 14)
        pdf.cell(0, 10, title, ln=1)
        pdf.set_font("Arial", "", 11)
        pdf.multi_cell(0, 6, 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)
        out = pdf.output(dest="S").encode("latin-1")
        return out
    else:
        # fallback: simple text file disguised as pdf (not ideal)
        return ("Classification:\n"+text+"\n\nInputs:\n"+json.dumps(inputs, indent=2)).encode("utf-8")

def build_full_phase1_pdf_bytes(site_list: List[dict], ext_refs: List[str]) -> bytes:
    """
    Build full geotechnical report (Phase 1) combining site data, classifier results, GSD and maps (if any).
    This is a streamlined generation: for production you would expand each analysis section.
    """
    title = "GeoMate — Full Geotechnical Investigation Report"
    # Use ReportLab if available
    if REPORTLAB_OK:
        buf = io.BytesIO()
        doc = SimpleDocTemplate(buf, pagesize=A4, leftMargin=20*mm, rightMargin=20*mm, topMargin=20*mm)
        styles = getSampleStyleSheet()
        elems = []
        elems.append(Paragraph(title, styles["Title"]))
        elems.append(Spacer(1,8))
        elems.append(Paragraph(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M')}", styles["Normal"]))
        elems.append(Spacer(1,12))
        # Summary
        elems.append(Paragraph("SUMMARY", styles["Heading2"]))
        elems.append(Paragraph("This full geotechnical report was generated by GeoMate V2. The following pages contain site descriptions, classification results and recommendations.", styles["BodyText"]))
        elems.append(PageBreak())
        # For each site:
        for s in site_list:
            elems.append(Paragraph(f"SITE: {s.get('Site Name','Unnamed')}", styles["Heading1"]))
            elems.append(Paragraph(f"Location: {s.get('Site Coordinates','Not provided')}", styles["BodyText"]))
            elems.append(Spacer(1,6))
            # classifier results
            elems.append(Paragraph("Classification", styles["Heading2"]))
            elems.append(Paragraph(s.get("classifier_decision_path","No classification available.").replace("\n","<br/>"), styles["BodyText"]))
            elems.append(Spacer(1,6))
            # GSD table
            if s.get("GSD"):
                elems.append(Paragraph("GSD Summary", styles["Heading2"]))
                g = s["GSD"]
                tdata = [["D10 (mm)", "D30 (mm)", "D60 (mm)", "Cu", "Cc"]]
                tdata.append([str(g.get("D10")), str(g.get("D30")), str(g.get("D60")), str(g.get("Cu")), str(g.get("Cc"))])
                t = Table(tdata, colWidths=[30*mm]*5)
                t.setStyle(TableStyle([("GRID",(0,0),(-1,-1),0.5,colors.grey),("BACKGROUND",(0,0),(-1,0),colors.HexColor("#1F4E79")),("TEXTCOLOR",(0,0),(-1,0),colors.white)]))
                elems.append(t)
            elems.append(PageBreak())
        # Add external references
        elems.append(Paragraph("External References", styles["Heading2"]))
        for r in ext_refs:
            elems.append(Paragraph(r, styles["BodyText"]))
        doc.build(elems)
        pdf_bytes = buf.getvalue(); buf.close()
        return pdf_bytes
    elif FPDF_OK:
        pdf = FPDF()
        pdf.add_page()
        pdf.set_font("Arial", "B", 16)
        pdf.cell(0, 10, title, ln=1)
        pdf.set_font("Arial", "", 11)
        pdf.ln(4)
        for s in site_list:
            pdf.set_font("Arial", "B", 13)
            pdf.cell(0, 8, f"Site: {s.get('Site Name','')}", ln=1)
            pdf.set_font("Arial", "", 11)
            pdf.multi_cell(0, 6, s.get("classifier_decision_path","No classification data."))
            pdf.ln(4)
            if s.get("GSD"):
                g = s["GSD"]
                pdf.cell(0, 6, f"GSD D10={g.get('D10')}, D30={g.get('D30')}, D60={g.get('D60')}", ln=1)
            pdf.add_page()
        out = pdf.output(dest="S").encode("latin-1")
        return out
    else:
        # fallback text
        parts = [title, "Generated: "+datetime.now().isoformat()]
        for s in site_list:
            parts.append("SITE: "+s.get("Site Name",""))
            parts.append("Classification:\n"+(s.get("classifier_decision_path") or "No data"))
            parts.append("GSD: "+(json.dumps(s.get("GSD") or {})))
        parts.append("REFERENCES: "+json.dumps(ext_refs))
        return ("\n\n".join(parts)).encode("utf-8")

# Reports UI
def reports_ui():
    st.header("📑 Reports — Classification-only & Full Geotechnical Report")
    idx, sdict = active_site()
    st.subheader("Classification-only Report")
    st.markdown("Generates a PDF containing the classification result and the decision path for the selected site.")
    if st.button("Generate Classification PDF", key=f"gen_cls_pdf_{sdict['Site Name']}"):
        pdf_bytes = build_classification_pdf_bytes(sdict)
        st.download_button("Download Classification PDF", data=pdf_bytes, file_name=f"{sdict['Site Name']}_classification.pdf", mime="application/pdf")
    st.markdown("---")
    st.subheader("Full Geotechnical Report")
    st.markdown("Chatbot will gather missing parameters and produce a full Phase 1 report for selected sites (up to 4).")
    # multi-select sites
    all_site_names = [s["Site Name"] for s in ss["site_descriptions"]]
    chosen = st.multiselect("Select sites to include", options=all_site_names, default=all_site_names, key="report_sites_select")
    ext_refs_text = st.text_area("External references (one per line)", key="ext_refs")
    # If user wants to have the bot gather missing params, start convo
    if st.button("Start conversational data gather for Full Report", key="start_report_convo"):
        # We'll run a simple in-app conversational loop: for brevity, ask a set of required params per site
        required_questions = [
            ("Load Bearing Capacity", "What is the soil bearing capacity (e.g., 2000 psf or 'Don't know')?"),
            ("Skin Shear Strength", "Provide skin shear strength (kPa) if known, else 'Don't know'"),
            ("Relative Compaction", "Relative compaction (%)"),
            ("Rate of Consolidation", "Rate of consolidation (e.g., cv in m2/year)"),
            ("Nature of Construction", "Nature of construction (e.g., 2-storey residence)"),
            ("Other", "Any other relevant notes or 'None'")
        ]
        # For each chosen site, ask interactively. We'll perform a sequential loop inside the app (chatbot-style rudimentary)
        for site_name in chosen:
            st.info(f"Collecting data for site: {site_name}")
            # Retrieve site dict
            site_idx = [i for i,s in enumerate(ss["site_descriptions"]) if s["Site Name"]==site_name][0]
            site_obj = ss["site_descriptions"][site_idx]
            for field, q in required_questions:
                # If field exists and not None, skip confirmation ask
                current_val = site_obj.get(field)
                if current_val:
                    st.write(f"{field} (existing): {current_val}")
                    keep = st.radio(f"Keep existing {field} for {site_name}?", ["Keep","Replace"], key=f"keep_{site_name}_{field}")
                    if keep == "Keep":
                        continue
                ans = st.text_input(q, key=f"report_q_{site_name}_{field}")
                if ans.strip().lower() in ["don't know","dont know","skip","n","no","unknown",""]:
                    site_obj[field] = None
                else:
                    site_obj[field] = ans.strip()
                ss["site_descriptions"][site_idx] = site_obj
            st.success(f"Data saved for {site_name}.")
        st.info("Conversational gather complete. Use Generate Full Report PDF button to create the PDF.")

    if st.button("Generate Full Report PDF", key="gen_full_pdf_btn"):
        # collect selected site dicts
        site_objs = [s for s in ss["site_descriptions"] if s["Site Name"] in chosen]
        ext_refs = [r.strip() for r in ext_refs_text.splitlines() if r.strip()]
        pdf_bytes = build_full_phase1_pdf_bytes(site_objs, ext_refs)
        st.download_button("Download Full Geotechnical Report", data=pdf_bytes, file_name=f"GeoMate_Full_Report_{datetime.now().strftime('%Y%m%d')}.pdf", mime="application/pdf")
# -------------------------
# Soil Recognizer (placeholder)
# -------------------------
def soil_recognizer_ui():
    st.header("🖼️ Soil Recognizer (image-based)")
    st.info("Upload a soil image (photo of sample or field). If an offline model 'soil_best_model.pth' exists in repo, it will be used. Otherwise this page acts as a placeholder for integrating your ML model or API.")
    uploaded = st.file_uploader("Upload soil image (jpg/png)", type=["jpg","jpeg","png"], key="sr_upload")
    if uploaded:
        st.image(uploaded, caption="Uploaded image", use_column_width=True)
        # Placeholder inference
        st.warning("Model inference not configured in this Space. To enable, upload 'soil_best_model.pth' and implement model loading code here.")
        if OCR_TESSERACT:
            st.info("Attempting OCR of image (to extract printed text)")
            try:
                img = Image.open(uploaded)
                text = pytesseract.image_to_string(img)
                st.text_area("Extracted text (OCR)", value=text, height=200)
            except Exception as e:
                st.error(f"OCR failed: {e}")

# -------------------------
# Main UI runner
# -------------------------
def main():
    sidebar_ui()
    # Top-level content area routing
    page = ss.get("page", "Landing")
    if page == "Landing":
        landing_ui()
    elif page == "Soil Recognizer":
        soil_recognizer_ui()
    elif page == "Soil Classifier":
        soil_classifier_ui()
    elif page == "GSD Curve":
        gsd_curve_ui()
    elif page == "Locator":
        locator_ui()
    elif page == "GeoMate Ask":
        rag_ui()
    elif page == "Reports":
        reports_ui()
    else:
        st.write("Page not found.")

if __name__ == "__main__":
    main()