app.py

# app.py — GeoMate V2 (single-file)
# Save this as app.py in your HuggingFace Space (or local folder)

# 0) Page config (must be first Streamlit command)
import streamlit as st
st.set_page_config(page_title="GeoMate V2", page_icon="🌍", layout="wide", initial_sidebar_state="expanded")

# 1) Standard imports
import os
import io
import json
import math
import base64
import tempfile
from datetime import datetime
from typing import Dict, Any, Tuple, List, Optional
import streamlit as st
from streamlit_folium import st_folium

# Visualization & PDF
import matplotlib.pyplot as plt
from reportlab.lib.pagesizes import A4
from reportlab.lib import colors
from reportlab.lib.units import mm
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, Image as RLImage, PageBreak
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle

# Ensure icon fonts load (fix desktop icon display for option_menu)
st.markdown("""
<!-- Load icon fonts used by streamlit_option_menu -->
<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/bootstrap-icons@1.11.3/font/bootstrap-icons.css">
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.0/css/all.min.css">
""", unsafe_allow_html=True)

# Optional imports handled gracefully
try:
    import geemap
    import ee
    EE_AVAILABLE = True
except Exception:
    EE_AVAILABLE = False

try:
    from fpdf import FPDF
    FPDF_AVAILABLE = True
except Exception:
    FPDF_AVAILABLE = False

try:
    import faiss
    FAISS_AVAILABLE = True
except Exception:
    FAISS_AVAILABLE = False

try:
    import pytesseract
    from PIL import Image
    OCR_AVAILABLE = True
except Exception:
    OCR_AVAILABLE = False

# Groq client import — we will require key
try:
    from groq import Groq
    GROQ_AVAILABLE = True
except Exception:
    GROQ_AVAILABLE = False

# 2) Secrets check (strict)
REQUIRED_SECRETS = ["GROQ_API_KEY", "SERVICE_ACCOUNT", "EARTH_ENGINE_KEY"]
missing = [s for s in REQUIRED_SECRETS if not os.environ.get(s)]
if missing:
    st.sidebar.error(f"Missing required secrets: {', '.join(missing)}. Please add these to your HF Space secrets.")
    st.error("Required secrets missing. Please set GROQ_API_KEY, SERVICE_ACCOUNT, and EARTH_ENGINE_KEY in Secrets and reload the app.")
    st.stop()

# If Groq lib missing, still stop because user requested Groq usage
if not GROQ_AVAILABLE:
    st.sidebar.error("Python package 'groq' not installed. Add it to requirements.txt and redeploy.")
    st.error("Missing required library 'groq'. Please add to requirements and redeploy.")
    st.stop()

# 3) Global constants & helper functions
MAX_SITES = 4


# ----------------------------
# Soil Recognizer Page (Integrated 6-Class ResNet18)
# ----------------------------
import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as T
from PIL import Image
import streamlit as st

# ----------------------------
# Load Soil Model (6 Classes)
# ----------------------------
@st.cache_resource
def load_soil_model(path="soil_best_model.pth"):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    try:
        model = models.resnet18(pretrained=False)
        num_ftrs = model.fc.in_features
        model.fc = nn.Linear(num_ftrs, 6)  # 6 soil classes

        # Load checkpoint
        state_dict = torch.load(path, map_location=device)
        model.load_state_dict(state_dict)
        model = model.to(device)
        model.eval()
        return model, device
    except Exception as e:
        st.error(f"⚠️ Could not load soil model: {e}")
        return None, device

soil_model, device = load_soil_model()

# ----------------------------
# Soil Classes & Transform
# ----------------------------
SOIL_CLASSES = ["Clay", "Gravel", "Loam", "Peat", "Sand", "Silt"]

transform = T.Compose([
    T.Resize((224, 224)),
    T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406],
                [0.229, 0.224, 0.225])
])

# ----------------------------
# Prediction Function
# ----------------------------
def predict_soil(img: Image.Image):
    if soil_model is None:
        return "Model not loaded", {}

    img = img.convert("RGB")
    inp = transform(img).unsqueeze(0).to(device)

    with torch.no_grad():
        logits = soil_model(inp)
        probs = torch.softmax(logits[0], dim=0)

    top_idx = torch.argmax(probs).item()
    predicted_class = SOIL_CLASSES[top_idx]

    result = {SOIL_CLASSES[i]: float(probs[i]) for i in range(len(SOIL_CLASSES))}
    return predicted_class, result

# ----------------------------
# Soil Recognizer Page
# ----------------------------
def soil_recognizer_page():
    st.header("🖼️ Soil Recognizer (ResNet18)")

    site = st.session_state["sites"]  # your existing site getter
    if site is None:
        st.warning("⚠️ No active site selected. Please add or select a site from the sidebar.")
        return

    uploaded = st.file_uploader("Upload soil image", type=["jpg", "jpeg", "png"])
    if uploaded is not None:
        img = Image.open(uploaded)
        st.image(img, caption="Uploaded soil image", use_column_width=True)

        predicted_class, confidence_scores = predict_soil(img)
        st.success(f"✅ Predicted: **{predicted_class}**")

        st.subheader("Confidence Scores")
        for cls, score in confidence_scores.items():
            st.write(f"{cls}: {score:.2%}")

        if st.button("Save to site"):
            # Save predicted soil class into Soil Profile field
            st.session_state["sites"][st.session_state["active_site"]]["Soil Class"] = predicted_class
            st.session_state["sites"][st.session_state["active_site"]]["Soil Recognizer Confidence"] = confidence_scores[predicted_class]
            save_active_site(site)
            st.success("Saved prediction to active site memory.")

# Pre-defined dropdown text mappings (as you requested) — exact text with mapping numbers for logic backend
DILATANCY_OPTIONS = [
    "1. Quick to slow",
    "2. None to very slow",
    "3. Slow",
    "4. Slow to none",
    "5. None",
    "6. Null?"
]
TOUGHNESS_OPTIONS = [
    "1. None",
    "2. Medium",
    "3. Slight?",
    "4. Slight to Medium?",
    "5. High",
    "6. Null?"
]
DRY_STRENGTH_OPTIONS = [
    "1. None to slight",
    "2. Medium to high",
    "3. Slight to Medium",
    "4. High to very high",
    "5. Null?"
]

# Map option text to numeric codes used in your USCS logic
DILATANCY_MAP = {DILATANCY_OPTIONS[i]: i+1 for i in range(len(DILATANCY_OPTIONS))}
TOUGHNESS_MAP = {TOUGHNESS_OPTIONS[i]: i+1 for i in range(len(TOUGHNESS_OPTIONS))}
DRY_STRENGTH_MAP = {DRY_STRENGTH_OPTIONS[i]: i+1 for i in range(len(DRY_STRENGTH_OPTIONS))}

# Engineering characteristics dictionary (expanded earlier; trimmed to representative entries but detailed)
ENGINEERING_CHARACTERISTICS = {
    "Gravel": {
        "Settlement": "None",
        "Quicksand": "Impossible",
        "Frost-heaving": "None",
        "Groundwater_lowering": "Possible",
        "Cement_grouting": "Possible",
        "Silicate_bitumen_injections": "Unsuitable",
        "Compressed_air": "Possible (see notes)"
    },
    "Coarse sand": {
        "Settlement": "None",
        "Quicksand": "Impossible",
        "Frost-heaving": "None",
        "Groundwater_lowering": "Possible",
        "Cement_grouting": "Possible only if very coarse",
        "Silicate_bitumen_injections": "Suitable",
        "Compressed_air": "Suitable"
    },
    "Medium sand": {
        "Settlement": "None",
        "Quicksand": "Unlikely",
        "Frost-heaving": "None",
        "Groundwater_lowering": "Suitable",
        "Cement_grouting": "Impossible",
        "Silicate_bitumen_injections": "Suitable",
        "Compressed_air": "Suitable"
    },
    "Fine sand": {
        "Settlement": "None",
        "Quicksand": "Liable",
        "Frost-heaving": "None",
        "Groundwater_lowering": "Suitable",
        "Cement_grouting": "Impossible",
        "Silicate_bitumen_injections": "Not possible in very fine sands",
        "Compressed_air": "Suitable"
    },
    "Silt": {
        "Settlement": "Occurs",
        "Quicksand": "Liable (very coarse silts may behave differently)",
        "Frost-heaving": "Occurs",
        "Groundwater_lowering": "Generally not suitable (electro-osmosis possible)",
        "Cement_grouting": "Impossible",
        "Silicate_bitumen_injections": "Impossible",
        "Compressed_air": "Suitable"
    },
    "Clay": {
        "Settlement": "Occurs",
        "Quicksand": "Impossible",
        "Frost-heaving": "None",
        "Groundwater_lowering": "Impossible (generally)",
        "Cement_grouting": "Only in stiff fissured clay",
        "Silicate_bitumen_injections": "Impossible",
        "Compressed_air": "Used for support only in special cases"
    }
}

# USCS & AASHTO verbatim logic (function)
from math import floor

def classify_uscs_aashto(inputs: Dict[str, Any]) -> Tuple[str, str, int, Dict[str, str], str]:
    """
    Verbatim USCS + AASHTO classifier based on the logic you supplied.
    inputs: dictionary expected keys:
      opt: 'y' or 'n'
      P2 (float): % passing #200 (0.075 mm)
      P4 (float): % passing #4 (4.75 mm)
      D60, D30, D10 (float mm) - can be 0 if unknown
      LL, PL (float)
      nDS, nDIL, nTG (int) mapped from dropdowns
    Returns:
      result_text (markdown), aashto_str, GI, engineering_characteristics (dict), uscs_str
    """
    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
        chars = {"summary":"Highly organic peat — large settlement, low strength, not suitable for foundations without improvement."}
        res_text = f"According to USCS, the soil is **{uscs}** — {uscs_expl}\n\nAccording to AASHTO, the soil is **{aashto}**."
        return res_text, aashto, GI, chars, uscs

    # parse numeric inputs with defaults
    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 if (LL is not None and PL is not None) else 0.0

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

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

    # === AASHTO (verbatim) ===
    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
    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} (Group Index = {GI})"

    # engineering characteristics pick
    char_summary = {}
    found_key = None
    for key in ENGINEERING_CHARACTERISTICS:
        if key.lower() in uscs.lower() or key.lower() in uscs_expl.lower():
            found_key = key
            break
    if found_key:
        char_summary = ENGINEERING_CHARACTERISTICS[found_key]
    else:
        # fallback selection by starting letter
        if uscs.startswith("G") or uscs.startswith("S"):
            char_summary = ENGINEERING_CHARACTERISTICS.get("Coarse sand", {})
        else:
            char_summary = ENGINEERING_CHARACTERISTICS.get("Silt", {})

    res_text_lines = [
        f"According to USCS, the soil is **{uscs}** — {uscs_expl}",
        f"According to AASHTO, the soil is **{aashto_expl}**",
        "",
        "Engineering characteristics (summary):"
    ]
    for k,v in char_summary.items():
        res_text_lines.append(f"- **{k}**: {v}")

    result_text = "\n".join(res_text_lines)
    return result_text, aashto_expl, GI, char_summary, uscs

# Helper: GSD interpolation to find diameters D10,D30,D60
def compute_gsd_metrics(diams: List[float], passing: List[float]) -> Dict[str, float]:
    """
    diams: list of diameters in mm (descending)
    passing: corresponding % passing (0-100)
    returns D10, D30, D60, Cu, Cc
    """
    # ensure descending diam, convert to float arrays
    if len(diams) < 2 or len(diams) != len(passing):
        raise ValueError("Diameters and passing arrays must match and have at least 2 items.")
    # linear interpolation on log(d)
    import numpy as np
    d = np.array(diams)
    p = np.array(passing)
    # make sure p is decreasing or increasing? passing decreases as diameter decreases, but we will handle general interpolation by sorting by diameter descending
    order = np.argsort(-d)
    d = d[order]
    p = p[order]
    # drop duplicates etc
    # function to find Dx = diameter at which passing = x (percent)
    def find_D(x):
        if x <= p.min():
            return float(d[p.argmin()])
        if x >= p.max():
            return float(d[p.argmax()])
        # linear interpolation on p vs log(d)
        from math import log, exp
        ld = np.log(d)
        # interpolate ld as function of p
        ld_interp = np.interp(x, p[::-1], ld[::-1])  # reverse because interp expects ascending x
        return float(math.exp(ld_interp))
    D10 = find_D(10.0)
    D30 = find_D(30.0)
    D60 = find_D(60.0)
    Cu = D60 / D10 if D10 > 0 else 0.0
    Cc = (D30 ** 2) / (D10 * D60) if (D10 > 0 and D60 > 0) else 0.0
    return {"D10":D10, "D30":D30, "D60":D60, "Cu":Cu, "Cc":Cc}



# 4) Session state initialization
if "sites" not in st.session_state:
    # initialize with a default site
    st.session_state["sites"] = [{
            "Site Name": None,
            "Project Name": "Project",
            "Site ID": None,
            "Soil Class": None,
            "Soil Recognizer Confidence": None,
            "Coordinates": "",
            "lat": None,
            "lon": None,
            "Project Description": "",
            # ---------------------------
            # Site Characterization
            # ---------------------------
            "Topography": None,              # manual topo entry
            "Drainage": None,                # manual drainage notes
            "Current Land Use": None,        # can be linked to Environmental Data
            "Regional Geology": None,        # manual geology notes
            # ---------------------------
            # Investigations & Lab
            # ---------------------------
            "Field Investigation": [],
            "Laboratory Results": [],
            "GSD": None,
            "USCS": None,
            "AASHTO": None,
            "GI": None,
            # ---------------------------
            # Geotechnical Parameters
            # ---------------------------
            "Load Bearing Capacity": None,
            "Skin Shear Strength": None,
            "Relative Compaction": None,
            "Rate of Consolidation": None,
            "Nature of Construction": None,
            # ---------------------------
            # Earth Engine Data
            # ---------------------------
            "Soil Profile": {                # SoilGrids (multi-parameter)
            "Clay": None,                # e.g. % clay at 200 cm
            "Sand": None,                # % sand
            "Silt": None,                # % silt
            "OrganicCarbon": None,       # % organic carbon
            "pH": None                   # soil pH if available
            },
            "Topo Data": None,               # Avg elevation (SRTM DEM)
            "Seismic Data": None,            # PGA/g (GEM hazard)
            "Flood Data": None,              # JRC Surface Water occurrence
            "Environmental Data": {          # Landcover, vegetation, urban, etc.
            "Landcover Stats": None,     # histogram by class
            "Forest Loss": None,         # future add: Hansen dataset
            "Urban Fraction": None       # optional calc from landcover
            },
            "Weather Data": {                # daily/monthly climate summaries
            "Rainfall": None,
            "Temperature": None,
            "Humidity": None
            },
            "Atmospheric Data": {            # optional: pollution, aerosols
            "AerosolOpticalDepth": None,
            "NO2": None,
            "CO": None
            },
            # ---------------------------
            # Map & Visualization
            # ---------------------------
            "map_snapshot": None,
            # ---------------------------
            # AI / Reporting
            # ---------------------------
            "chat_history": [],
            "classifier_inputs": {},
            "classifier_decision": None,
            "report_convo_state": 0,
            "report_missing_fields": [],
            "report_answers": {}
             }]
        

if "active_site" not in st.session_state:
    st.session_state["active_site"] = 0

if "llm_model" not in st.session_state:
    st.session_state["llm_model"] = "groq/compound"

# Groq client (simple wrapper)
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
def groq_generate(prompt: str, model: str = None, max_tokens: int = 512) -> str:
    """Call Groq. If call fails, return an explanatory text."""
    try:
        client = Groq(api_key=GROQ_API_KEY)
        model_name = model or st.session_state["llm_model"]
        completion = client.chat.completions.create(
            model=model_name,
            messages=[{"role":"user","content":prompt}],
            temperature=0.2,
            max_tokens=max_tokens
        )
        text = completion.choices[0].message.content
        return text
    except Exception as e:
        return f"[LLM error or offline: {e}]"

# 5) UI helper: nice CSS for chat bubbles & page styling
st.markdown("""
<style>
/* Background and card styling */
body { background: #0b0b0b; color: #e9eef6; }
.stApp > .main > .block-container { padding-top: 18px; }

/* Landing and cards */
.gm-card { background: linear-gradient(180deg, rgba(255,122,0,0.04), rgba(255,122,0,0.02)); border-radius:12px; padding:14px; border:1px solid rgba(255,122,0,0.06);}
.gm-cta { background: linear-gradient(90deg,#ff7a00,#ff3a3a); color:white; padding:10px 14px; border-radius:10px; font-weight:700; }

/* Chat bubbles */
.chat-bot { background: #0f1720; border-left:4px solid #FF7A00; padding:10px 12px; border-radius:12px; margin:6px 0; color:#e9eef6; }
.chat-user { background: #1a1f27; padding:10px 12px; border-radius:12px; margin:6px 0; color:#cfe6ff; text-align:right;}
.small-muted { color:#9aa7bf; font-size:12px; }
</style>
""", unsafe_allow_html=True)

# 6) Sidebar: navigation, site selector, model selector
from streamlit_option_menu import option_menu

with st.sidebar:
    st.markdown("<h2 style='color:#FF8C00;margin:8px 0'>GeoMate V2</h2>", unsafe_allow_html=True)
    # LLM model selector
    st.session_state["llm_model"] = st.selectbox("Select LLM model", options=[
        "meta-llama/llama-4-maverick-17b-128e-instruct",
        "llama-3.1-8b-instant",
        "meta-llama/llama-guard-4-12b",
        "llama-3.3-70b-versatile",
        "groq/compound"
    ], index=0)
    st.markdown("---")

    # Site management controls
    st.markdown("### Project Sites")
    site_names = [s.get("Site Name", f"Site {i}") for i,s in enumerate(st.session_state["sites"])]
    # Add new site input
    new_site_name = st.text_input("New site name", value="", key="new_site_name_input")
    if st.button("➕ Add / Create Site"):
        if new_site_name.strip() == "":
            st.warning("Enter a name for the new site.")
        elif len(st.session_state["sites"]) >= MAX_SITES:
            st.error(f"Maximum of {MAX_SITES} sites allowed.")
        else:
            idx = len(st.session_state["sites"])
            #idx = len(st.session_state["sites"]) + 1
            st.session_state["sites"].append({
            "Site Name": new_site_name.strip(),
            "Project Name": "Project - " + new_site_name.strip(),
            "Site ID": idx,
            "Soil Class": None,
            "Soil Recognizer Confidence": None,
            "Coordinates": "",
            "lat": None,
            "lon": None,
            "Project Description": "",
            # ---------------------------
            # Site Characterization
            # ---------------------------
            "Topography": None,              # manual topo entry
            "Drainage": None,                # manual drainage notes
            "Current Land Use": None,        # can be linked to Environmental Data
            "Regional Geology": None,        # manual geology notes
            # ---------------------------
            # Investigations & Lab
            # ---------------------------
            "Field Investigation": [],
            "Laboratory Results": [],
            "GSD": None,
            "USCS": None,
            "AASHTO": None,
            "GI": None,
            # ---------------------------
            # Geotechnical Parameters
            # ---------------------------
            "Load Bearing Capacity": None,
            "Skin Shear Strength": None,
            "Relative Compaction": None,
            "Rate of Consolidation": None,
            "Nature of Construction": None,
            # ---------------------------
            # Earth Engine Data
            # ---------------------------
            "Soil Profile": {                # SoilGrids (multi-parameter)
            "Clay": None,                # e.g. % clay at 200 cm
            "Sand": None,                # % sand
            "Silt": None,                # % silt
            "OrganicCarbon": None,       # % organic carbon
            "pH": None                   # soil pH if available
            },
            "Topo Data": None,               # Avg elevation (SRTM DEM)
            "Seismic Data": None,            # PGA/g (GEM hazard)
            "Flood Data": None,              # JRC Surface Water occurrence
            "Environmental Data": {          # Landcover, vegetation, urban, etc.
            "Landcover Stats": None,     # histogram by class
            "Forest Loss": None,         # future add: Hansen dataset
            "Urban Fraction": None       # optional calc from landcover
            },
            "Weather Data": {                # daily/monthly climate summaries
            "Rainfall": None,
            "Temperature": None,
            "Humidity": None
            },
            "Atmospheric Data": {            # optional: pollution, aerosols
            "AerosolOpticalDepth": None,
            "NO2": None,
            "CO": None
            },
            # ---------------------------
            # Map & Visualization
            # ---------------------------
            "map_snapshot": None,
            # ---------------------------
            # AI / Reporting
            # ---------------------------
            "chat_history": [],
            "classifier_inputs": {},
            "classifier_decision": None,
            "report_convo_state": 0,
            "report_missing_fields": [],
            "report_answers": {}
             })


            st.success(f"Site '{new_site_name.strip()}' created.")
            st.session_state["active_site"] = idx
            st.rerun()

    # Active site selector
    if site_names:
        active_index = st.selectbox("Active Site", options=list(range(len(site_names))),
                                    format_func=lambda x: site_names[x], index=st.session_state["active_site"])
        st.session_state["active_site"] = active_index
    st.markdown("---")
    st.write("Active Site JSON (live)")
    st.json(st.session_state["sites"][st.session_state["active_site"]])

    st.markdown("---")
    st.markdown("© GeoMate • Advanced geotechnical copilot", unsafe_allow_html=True)

# 7) Pages implementation
def landing_page():
    # Background hero image
    BACKGROUND_URL = "/app/background_placeholder.jpg"  # Replace with your image or URL

    # -------------------- HERO SECTION --------------------
    st.markdown(f"""
    <div style="
        background-image: url('{BACKGROUND_URL}');
        background-size: cover;
        background-position: center;
        padding: 64px 32px;
        border-radius: 16px;
        margin-bottom: 24px;
        position: relative;
        box-shadow: 0 6px 24px rgba(0,0,0,0.4);
        ">
      <div style="background: rgba(15,15,15,0.65);
                  padding:32px;
                  border-radius:14px;
                  max-width:980px;">
        <h1 style='color:#FF8C00; font-size:42px; margin:0 0 12px'>
            🌍 GeoMate V2
        </h1>
        <p style='color:#f2f4f7; font-size:17px; margin:0 0 14px; line-height:1.6'>
          Your AI Geotechnical Copilot — soil recognition, classification, locator (Google Earth Engine), 
          RAG-powered Q&A, and dynamic professional reports.
        </p>
        <div style='margin-top:10px; color:#d6d6d6; font-size:14px'>
          🚀 Quick: Classifier • GSD • Locator • RAG • OCR • Reports
        </div>
      </div>
    </div>
    """, unsafe_allow_html=True)

    # -------------------- BRAND HEADER --------------------
    st.markdown("""
    <div style="display:flex;align-items:center;gap:16px; margin-bottom:14px">
        <div style="width:80px;height:80px;
                    border-radius:16px;
                    background:linear-gradient(135deg,#ff7a00,#ff3a3a);
                    display:flex;align-items:center;justify-content:center;
                    box-shadow:0 6px 18px rgba(0,0,0,0.55)">
            <span style='font-size:36px'>🛰️</span>
        </div>
        <div>
            <h1 style='margin:0;color:#FF8C00;font-size:30px'>GeoMate V2</h1>
            <div style='color:#888; font-size:14px'>
                AI geotechnical copilot — smarter, faster & field-ready
            </div>
        </div>
    </div>
    """, unsafe_allow_html=True)
    st.markdown("---")

    # -------------------- CONTENT GRID --------------------
    col1, col2 = st.columns([2,1])

    with col1:
        st.markdown("<div class='gm-card'>", unsafe_allow_html=True)
        st.write("GeoMate is built to help geotechnical engineers streamline workflows:")
        st.markdown("""
        - 🧪 **Soil Recognition & Classification** (USCS, AASHTO)  
        - 📈 **Grain Size Distribution (GSD) Curve plotting**  
        - 🌍 **Earth Engine data fetch for sites**  
        - 🤖 **RAG + LLM Q&A** for textbook-driven answers  
        - 📑 **Professional report generation** (PDF-ready)  
        - 📷 **OCR-based input from site drawings**  
        """)
        st.markdown("</div>", unsafe_allow_html=True)

        # Quick Actions as styled cards
        st.markdown("### ⚡ Quick Actions")
        qa_cols = st.columns(3)
        quick_actions = [
            ("🧪", "Classifier", "Classifier"),
            ("📈", "GSD Curve", "GSD"),
            ("🌍", "Locator", "Locator"),
            ("🤖", "GeoMate Ask", "RAG"),
            ("📷", "OCR", "OCR"),
            ("📑", "Reports", "Reports"),
        ]
        for i, (icon, label, page) in enumerate(quick_actions):
            with qa_cols[i % 3]:
                if st.button(f"{icon} {label}", use_container_width=True):
                    st.session_state["page"] = page
                    st.rerun()

    # -------------------- SITE SUMMARY --------------------
    with col2:
        site = st.session_state["sites"][st.session_state["active_site"]]
        st.markdown("""
        <div style="
            background: #1e1e1e;
            border-radius:14px;
            padding:20px;
            text-align:center;
            box-shadow:0 4px 14px rgba(0,0,0,0.5);
        ">
          <h3 style='color:#FF8C00; margin:0 0 10px'>📍 Live Site Summary</h3>
        """, unsafe_allow_html=True)

        st.write(f"🏗️ **Site:** {site.get('Site Name')}")
        st.write(f"🧱 **USCS:** {site.get('USCS', '—')}")
        st.write(f"🛣️ **AASHTO:** {site.get('AASHTO', '—')}")
        st.write(f"📊 **GSD saved:** {'✅ Yes' if site.get('GSD') else '❌ No'}")

        st.markdown("</div>", unsafe_allow_html=True)

# Soil Classifier page (conversational, step-by-step)
def soil_classifier_page():
    st.header("🧪 Soil Classifier — Conversational (USCS & AASHTO)")
    site = st.session_state["sites"][st.session_state["active_site"]]

    # conversation state machine: steps list
    steps = [
        {"id":"intro", "bot":"Hello — I am the GeoMate Soil Classifier. Ready to start?"},
        {"id":"organic", "bot":"Is the soil at this site organic (contains high organic matter, feels spongy or has odour)?", "type":"choice", "choices":["No","Yes"]},
        {"id":"P2", "bot":"Please enter the percentage passing the #200 sieve (0.075 mm). Example: 12", "type":"number"},
        {"id":"P4", "bot":"What is the percentage passing the sieve no. 4 (4.75 mm)? (enter 0 if unknown)", "type":"number"},
        {"id":"hasD", "bot":"Do you know the D10, D30 and D60 diameters (in mm)?", "type":"choice","choices":["No","Yes"]},
        {"id":"D60", "bot":"Enter D60 (diameter in mm corresponding to 60% passing).", "type":"number"},
        {"id":"D30", "bot":"Enter D30 (diameter in mm corresponding to 30% passing).", "type":"number"},
        {"id":"D10", "bot":"Enter D10 (diameter in mm corresponding to 10% passing).", "type":"number"},
        {"id":"LL", "bot":"What is the liquid limit (LL)?", "type":"number"},
        {"id":"PL", "bot":"What is the plastic limit (PL)?", "type":"number"},
        {"id":"dry", "bot":"Select the observed dry strength of the fine soil (if applicable).", "type":"select", "options":DRY_STRENGTH_OPTIONS},
        {"id":"dilat", "bot":"Select the observed dilatancy behaviour.", "type":"select", "options":DILATANCY_OPTIONS},
        {"id":"tough", "bot":"Select the observed toughness.", "type":"select", "options":TOUGHNESS_OPTIONS},
        {"id":"confirm", "bot":"Would you like me to classify now?", "type":"choice", "choices":["No","Yes"]}
    ]

    if "classifier_step" not in st.session_state:
        st.session_state["classifier_step"] = 0
    if "classifier_inputs" not in st.session_state:
        st.session_state["classifier_inputs"] = dict(site.get("classifier_inputs", {}))

    step_idx = st.session_state["classifier_step"]

    # chat history display
    st.markdown("<div class='gm-card'>", unsafe_allow_html=True)
    st.markdown("<div class='chat-bot'>{}</div>".format("GeoMate: Hello — soil classifier ready. Use the controls below to answer step-by-step."), unsafe_allow_html=True)
    # Show stored user answers sequentially for context
    # render question up to current step
    for i in range(step_idx+1):
        s = steps[i]
        # show bot prompt
        st.markdown(f"<div class='chat-bot'>{s['bot']}</div>", unsafe_allow_html=True)
        # show user answer if exists in classifier_inputs
        key = s["id"]
        val = st.session_state["classifier_inputs"].get(key)
        if val is not None:
            st.markdown(f"<div class='chat-user'>{val}</div>", unsafe_allow_html=True)

    st.markdown("</div>", unsafe_allow_html=True)

    # Render input widget for current step
    current = steps[step_idx]
    step_id = current["id"]
    proceed = False
    user_answer = None

    cols = st.columns([1,1,1])
    with cols[0]:
        if current.get("type") == "choice":
            choice = st.radio(current["bot"], options=current["choices"], index=0, key=f"cls_{step_id}")
            user_answer = choice
        elif current.get("type") == "number":
            # numeric input without +/- spinner (we use text_input and validate)
            raw = st.text_input(current["bot"], value=str(st.session_state["classifier_inputs"].get(step_id,"")), key=f"cls_{step_id}_num")
            # validate numeric
            try:
                if raw.strip() == "":
                    user_answer = None
                else:
                    user_answer = float(raw)
            except:
                st.warning("Please enter a valid number (e.g., 12 or 0).")
                user_answer = None
        elif current.get("type") == "select":
            opts = current.get("options", [])
            sel = st.selectbox(current["bot"], options=opts, index=0, key=f"cls_{step_id}_sel")
            user_answer = sel
        else:
            # just a message step — proceed
            user_answer = None

    # controls: Next / Back
    coln, colb, colsave = st.columns([1,1,1])
    with coln:
        if st.button("➡️ Next", key=f"next_{step_id}"):
            # store answer if provided
            if current.get("type") == "number":
                if user_answer is None:
                    st.warning("Please enter a numeric value or enter 0 if unknown.")
                else:
                    st.session_state["classifier_inputs"][step_id] = user_answer
                    st.session_state["classifier_step"] = min(step_idx+1, len(steps)-1)
                    st.rerun()
            elif current.get("type") in ("choice","select"):
                st.session_state["classifier_inputs"][step_id] = user_answer
                st.session_state["classifier_step"] = min(step_idx+1, len(steps)-1)
                st.rerun()
            else:
                # message-only step
                st.session_state["classifier_step"] = min(step_idx+1, len(steps)-1)
                st.rerun()
    with colb:
        if st.button("⬅️ Back", key=f"back_{step_id}"):
            st.session_state["classifier_step"] = max(0, step_idx-1)
            st.rerun()
    with colsave:
        if st.button("💾 Save & Classify now", key="save_and_classify"):
            # prepare inputs in required format for classify_uscs_aashto
            ci = st.session_state["classifier_inputs"].copy()
            # Normalize choices into expected codes
            if isinstance(ci.get("dry"), str):
                ci["nDS"] = DRY_STRENGTH_MAP.get(ci.get("dry"), 5)
            if isinstance(ci.get("dilat"), str):
                ci["nDIL"] = DILATANCY_MAP.get(ci.get("dilat"), 6)
            if isinstance(ci.get("tough"), str):
                ci["nTG"] = TOUGHNESS_MAP.get(ci.get("tough"), 6)
            # map 'Yes'/'No' for organic and hasD
            ci["opt"] = "y" if ci.get("organic","No")=="Yes" or ci.get("organic",ci.get("organic"))=="Yes" else ci.get("organic","n")
            # our field names in CI may differ: convert organic stored under 'organic' step to 'opt'
            if "organic" in ci:
                ci["opt"] = "y" if ci["organic"]=="Yes" else "n"
            # map D entries: D60 etc may be present
            # call classification
            try:
                res_text, aashto, GI, chars, uscs = classify_uscs_aashto(ci)
            except Exception as e:
                st.error(f"Classification error: {e}")
                res_text = f"Error during classification: {e}"
                aashto = "N/A"; GI = 0; chars = {}; uscs = "N/A"
            # save into active site
            site["USCS"] = uscs
            site["AASHTO"] = aashto
            site["GI"] = GI
            site["classifier_inputs"] = ci
            site["classifier_decision"] = res_text
            st.success("Classification complete. Results saved to site.")
            st.write("### Classification Results")
            st.markdown(res_text)
            # Keep classifier_step at end so user can review
            st.session_state["classifier_step"] = len(steps)-1

# GSD Curve Page
def gsd_page():
    st.header("📈 Grain Size Distribution (GSD) Curve")
    site = st.session_state["sites"][st.session_state["active_site"]]
    st.markdown("Enter diameters (mm) and % passing (comma-separated). Use descending diameters (largest to smallest).")
    diam_input = st.text_area("Diameters (mm) comma-separated", value=site.get("GSD",{}).get("diameters","75,50,37.5,25,19,12.5,9.5,4.75,2,0.85,0.425,0.25,0.18,0.15,0.075") if site.get("GSD") else "75,50,37.5,25,19,12.5,9.5,4.75,2,0.85,0.425,0.25,0.18,0.15,0.075")
    pass_input = st.text_area("% Passing comma-separated", value=site.get("GSD",{}).get("passing","100,98,96,90,85,78,72,65,55,45,35,25,18,14,8") if site.get("GSD") else "100,98,96,90,85,78,72,65,55,45,35,25,18,14,8")
    if st.button("Compute GSD & Save"):
        try:
            diams = [float(x.strip()) for x in diam_input.split(",") if x.strip()]
            passing = [float(x.strip()) for x in pass_input.split(",") if x.strip()]
            metrics = compute_gsd_metrics(diams, passing)
            # plot
            fig, ax = plt.subplots(figsize=(7,4))
            ax.semilogx(diams, passing, marker='o')
            ax.set_xlabel("Particle size (mm)")
            ax.set_ylabel("% Passing")
            ax.invert_xaxis()
            ax.grid(True, which='both', linestyle='--', linewidth=0.5)
            ax.set_title("Grain Size Distribution")
            st.pyplot(fig)
            # save into site
            site["GSD"] = {"diameters":diams, "passing":passing, **metrics}
            st.success(f"Saved GSD for site. D10={metrics['D10']:.4g} mm, D30={metrics['D30']:.4g} mm, D60={metrics['D60']:.4g} mm")
        except Exception as e:
            st.error(f"GSD error: {e}")

# OCR Page
def ocr_page():
    st.header("📷 OCR — extract values from an image")
    site = st.session_state["sites"][st.session_state["active_site"]]
    if not OCR_AVAILABLE:
        st.warning("OCR dependencies not available (pytesseract/PIL). Add pytesseract and pillow to requirements to enable OCR.")
    uploaded = st.file_uploader("Upload an image (photo of textbook question or sieve data)", type=["png","jpg","jpeg"])
    if uploaded:
        if OCR_AVAILABLE:
            try:
                img = Image.open(uploaded)
                st.image(img, caption="Uploaded", use_column_width=True)
                text = pytesseract.image_to_string(img)
                st.text_area("Extracted text", value=text, height=180)
                # Basic parsing: try to find LL, PL, D10 etc via regex
                import re
                found = {}
                for key in ["LL","PL","D10","D30","D60","P2","P4","CBR"]:
                    pattern = re.compile(rf"{key}[:=]?\s*([0-9]+\.?[0-9]*)", re.I)
                    m = pattern.search(text)
                    if m:
                        found[key] = float(m.group(1))
                        site.setdefault("classifier_inputs",{})[key] = float(m.group(1))
                if found:
                    st.success(f"Parsed values: {found}")
                    st.write("Values saved into classifier inputs.")
                else:
                    st.info("No clear numeric matches found automatically.")
            except Exception as e:
                st.error(f"OCR failed: {e}")
        else:
            st.warning("OCR not available in this deployment.")
# Locator Page (with Earth Engine auth at top)
# Locator Page (with Earth Engine auth at top)

import os
import json
import streamlit as st
import geemap.foliumap as geemap
import ee
import matplotlib.pyplot as plt
from datetime import datetime
from io import BytesIO
import base64
import folium

import tempfile

def export_map_snapshot(m, width=800, height=600):
    """
    Export geemap Map object to PNG snapshot (returns bytes).
    - m: geemap.Map
    - width, height: dimensions of snapshot
    """
    try:
        # geemap has a built-in screenshot method
        tmpfile = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        m.screenshot(filename=tmpfile.name, region=None, dimensions=(width, height))
        with open(tmpfile.name, "rb") as f:
            return f.read()
    except Exception as e:
        st.warning(f"Map snapshot failed: {e}")
        return None


def locator_page():
    """
    Robust locator page:
    - Uses initialize_ee() auth routine (expects EARTHENGINE_TOKEN / SERVICE_ACCOUNT in env)
    - Shows interactive map with basemaps and overlays
    - Captures ROI drawn on the map
    - Computes summaries and saves them into active site and soil_json
    """

    st.title("🌍 GeoMate Interactive Earth Explorer")
    st.markdown(
        "Draw a polygon (or rectangle) on the map using the drawing tool. "
        "Then press **Compute Summaries** to compute soil clay, elevation, seismic, flood occurrence, landcover, and NDVI."
    )

    # ----------------------------
    # EE Auth
    # ----------------------------
    EARTHENGINE_TOKEN = os.getenv("EARTHENGINE_TOKEN")
    SERVICE_ACCOUNT = os.getenv("SERVICE_ACCOUNT")

    def initialize_ee():
        if "ee_initialized" in st.session_state and st.session_state["ee_initialized"]:
            return True
        if EARTHENGINE_TOKEN and SERVICE_ACCOUNT:
            try:
                creds = ee.ServiceAccountCredentials(email=SERVICE_ACCOUNT, key_data=EARTHENGINE_TOKEN)
                ee.Initialize(creds)
                st.session_state["ee_initialized"] = True
                return True
            except Exception as e:
                st.warning(f"Service account init failed: {e}, falling back...")
        try:
            ee.Initialize()
            st.session_state["ee_initialized"] = True
            return True
        except Exception:
            try:
                ee.Authenticate()
                ee.Initialize()
                st.session_state["ee_initialized"] = True
                return True
            except Exception as e:
                st.error(f"Earth Engine auth failed: {e}")
                return False

    if not initialize_ee():
        st.stop()

    # ----------------------------
    # Safe reducers
    # ----------------------------
    def safe_get_reduce(region, image, band, scale=1000, default=None, max_pixels=int(1e7)):
        try:
            rr = image.reduceRegion(ee.Reducer.mean(), region, scale=scale, maxPixels=max_pixels)
            val = rr.get(band)
            return float(val.getInfo()) if val else default
        except Exception:
            return default

    def safe_reduce_histogram(region, image, band, scale=1000, max_pixels=int(1e7)):
        try:
            rr = image.reduceRegion(ee.Reducer.frequencyHistogram(), region, scale=scale, maxPixels=max_pixels)
            hist = rr.get(band)
            return hist.getInfo() if hist else {}
        except Exception:
            return {}

    def safe_time_series(region, collection, band, start, end, reducer=ee.Reducer.mean(), scale=1000, max_pixels=int(1e7)):
        try:
            def per_image(img):
                date = img.date().format("YYYY-MM-dd")
                val = img.reduceRegion(reducer, region, scale=scale, maxPixels=max_pixels).get(band)
                return ee.Feature(None, {"date": date, "val": val})
            feats = collection.filterDate(start, end).map(per_image).filter(ee.Filter.notNull(["val"])).getInfo()
            pts = []
            for f in feats.get("features", []):
                p = f.get("properties", {})
                if p.get("val") is not None:
                    pts.append((p.get("date"), float(p.get("val"))))
            return pts
        except Exception:
            return []

    # ----------------------------
    # Map setup
    # ----------------------------
    m = geemap.Map(center=[28.0, 72.0], zoom=5, plugin_Draw=True, draw_export=True, locate_control=True)

    basemaps = [
        "HYBRID", "ROADMAP", "TERRAIN", "SATELLITE",
        "Esri.WorldImagery", "Esri.WorldTopoMap", "Esri.WorldShadedRelief",
        "Esri.NatGeoWorldMap", "Esri.OceanBasemap",
        "CartoDB.Positron", "CartoDB.DarkMatter",
        "Stamen.Terrain", "Stamen.Watercolor",
        "OpenStreetMap", "Esri.WorldGrayCanvas", "Esri.WorldStreetMap"
    ]
    for b in basemaps:
        try:
            m.add_basemap(b)
        except Exception:
            pass

    # ----------------------------
    # Datasets (DEM, Soil, Seismic, Flood, Landcover, NDVI)
    # ----------------------------
    
    # --- DEM ---
    try:
        dem = ee.Image("NASA/NASADEM_HGT/001")  # NASADEM ~30m global
        dem_band_name = "elevation"
        st.info("Using NASADEM dataset for elevation (30m).")
    except Exception:
        try:
            dem = ee.Image("USGS/SRTMGL1_003")  # SRTM ~30m fallback
            dem_band_name = "elevation"
            st.warning("Fallback to SRTM DEM (30m).")
        except Exception:
            dem = None
            dem_band_name = None
            st.error("No DEM dataset available.")
    
    # --- Soil Clay Fraction ---
    soil_img = None
    soil_band = None
    chosen_soil_band = None
    try:
        soil_img = ee.Image("OpenLandMap/SOL/SOL_CLAY-WFRACTION_USDA-3A1A1A_M/v02")
        bands = soil_img.bandNames().getInfo()
        # interactive soil band selector
        chosen_soil_band = st.selectbox(
            "Select soil depth / clay band",
            options=bands,
            index=bands.index("b200") if "b200" in bands else 0
        )
        st.info(f"Using OpenLandMap Clay Fraction — Band: {chosen_soil_band}")
    except Exception:
        try:
            soil_img = ee.Image("projects/soilgrids-isric/clay_mean")
            bands = soil_img.bandNames().getInfo()
            chosen_soil_band = st.selectbox(
                "Select soil depth (SoilGrids)",
                options=bands,
                index=0
            )
            st.warning(f"Fallback to SoilGrids Clay dataset — Band: {chosen_soil_band}")
        except Exception:
            soil_img = None
            chosen_soil_band = None
            st.error("No soil dataset available.")
    
    # --- Seismic Hazard ---
    try:
        seismic_img = ee.Image("SEDAC/GSHAPSeismicHazard")
        seismic_band = "gshap"
        st.info("Using SEDAC Global Seismic Hazard dataset.")
    except Exception:
        seismic_img = None
        seismic_band = None
        st.error("No seismic hazard dataset available.")
    
    # --- Flood Occurrence ---
    try:
        water = ee.Image("JRC/GSW1_4/GlobalSurfaceWater")
        water_band = "occurrence"
        st.info("Using JRC Global Surface Water Occurrence (1984–2020).")
    except Exception:
        water = None
        water_band = None
        st.error("No flood dataset available.")
    
    # --- Landcover ---
    try:
        landcover = ee.Image("ESA/WorldCover/v200")
        lc_band = "Map"
        st.info("Using ESA WorldCover v200 (2020, 10m).")
    except Exception:
        landcover = None
        lc_band = None
        st.error("No landcover dataset available.")
    
    # --- NDVI ---
    try:
        ndvi_col = ee.ImageCollection("MODIS/061/MOD13A2").select("NDVI")
        st.info("Using MODIS NDVI (16-day, 1km, global).")
    except Exception:
        ndvi_col = None
        st.error("No NDVI dataset available.")
    
    # ----------------------------
    # Add Layers to Map
    # ----------------------------
    if dem:
        try:
            m.addLayer(
                dem,
                {"min": 0, "max": 4000, "palette": ["blue", "green", "brown", "white"]},
                "DEM / Elevation"
            )
        except Exception:
            pass
    
    if soil_img and chosen_soil_band:
        try:
            m.addLayer(
                soil_img.select(chosen_soil_band),
                {"min": 0.0, "max": 0.6, "palette": ["#ffffcc", "#c2e699", "#78c679", "#31a354"]},
                f"Soil Clay Fraction ({chosen_soil_band})"
            )
        except Exception:
            pass
    
    if seismic_img:
        try:
            m.addLayer(
                seismic_img,
                {"min": 0, "max": 1, "palette": ["white", "yellow", "red"]},
                "Seismic Hazard"
            )
        except Exception:
            pass
    
    if water:
        try:
            m.addLayer(
                water.select(water_band),
                {"min": 0, "max": 100, "palette": ["white", "blue"]},
                "Flood Occurrence (%)"
            )
        except Exception:
            pass
    
    if landcover:
        try:
            m.addLayer(
                landcover,
                {
                    "min": 10,
                    "max": 100,
                    "palette": [
                        "#006400", "#ffbb22", "#ffff4c", "#f096ff", "#fa0000",
                        "#b4b4b4", "#f0f0f0", "#0064c8", "#0096a0", "#00cf75"
                    ]
                },
                "Landcover (ESA WorldCover)"
            )
        except Exception:
            pass
    
    try:
        countries = ee.FeatureCollection("USDOS/LSIB_SIMPLE/2017")
        m.addLayer(
            countries.style(**{"color": "black", "fillColor": "00000000", "width": 1}),
            {},
            "Country Boundaries"
        )
    except Exception:
        pass
    
    # ----------------------------
    # --- Render the map ---
    m.to_streamlit(height=600, responsive=True)
    
    # --- ROI capture ---
    roi = None
    
    # Case 1: New drawing in this run
    if m.user_roi:
        try:
            roi = m.user_roi
            st.session_state["roi_geojson"] = roi.toGeoJSONString()
            st.success("✅ ROI captured from user drawing")
        except Exception as e:
            st.warning(f"Could not parse m.user_roi: {e}")
    
    elif m.draw_features:
        try:
            last_feature = m.draw_features[-1]
            geom = last_feature.get("geometry") or last_feature
            roi = ee.Geometry(geom)
            st.session_state["roi_geojson"] = json.dumps(last_feature)
            st.success("✅ ROI captured from draw_features")
        except Exception as e:
            st.warning(f"Could not parse m.draw_features: {e}")
    
    # Case 2: Restore ROI from session_state if nothing new drawn
    elif "roi_geojson" in st.session_state:
        try:
            stored = json.loads(st.session_state["roi_geojson"])
            geom = stored.get("geometry") if isinstance(stored, dict) else stored
            roi = ee.Geometry(geom)
            st.info("♻️ ROI restored from session state")
        except Exception as e:
            st.warning(f"Could not restore ROI: {e}")
    
    # --- Button to trigger computation ---
    if st.button("Compute Summaries"):
        if roi is None:
            st.error("⚠️ No ROI found. Please draw a polygon/rectangle/circle and try again.")
        else:
            st.success("🚀 ROI found — computing summaries...")
          


        chosen_soil_band = None
        if soil_img:
            bands = soil_img.bandNames().getInfo()
            chosen_soil_band = st.selectbox("Soil band to analyze", bands, index=bands.index(soil_band) if soil_band in bands else 0)

        soil_val = safe_get_reduce(roi, soil_img.select(chosen_soil_band), chosen_soil_band, 1000) if soil_img and chosen_soil_band else None
        elev_val = safe_get_reduce(roi, dem, dem_band_name, 1000)
        seismic_val = safe_get_reduce(roi, seismic_img, seismic_band, 5000) if seismic_img else None
        flood_val = safe_get_reduce(roi, water.select(water_band), water_band, 30) if water else None
        lc_stats = safe_reduce_histogram(roi, landcover, lc_band, 30) if landcover else {}
        ndvi_ts = []
        if ndvi_col:
            end = datetime.utcnow().strftime("%Y-%m-%d")
            start = (datetime.utcnow().replace(year=datetime.utcnow().year-2)).strftime("%Y-%m-%d")
            ndvi_ts = safe_time_series(roi, ndvi_col, "NDVI", start, end)

        # Save results
        active = st.session_state.get("active_site", 0)
        if "sites" in st.session_state:
            site = st.session_state["sites"][active]
            site["ROI"] = roi.getInfo()
            site["Soil Profile"] = f"{soil_val} ({chosen_soil_band})" if soil_val else "N/A"
            site["Topo Data"] = f"{elev_val} m" if elev_val else "N/A"
            site["Seismic Data"] = seismic_val
            site["Flood Data"] = flood_val
            site["Environmental Data"] = {"Landcover": lc_stats, "NDVI": ndvi_ts}
        st.session_state["soil_json"] = {
            "Soil": soil_val, "Soil Band": chosen_soil_band,
            "Elevation": elev_val, "Seismic": seismic_val,
            "Flood": flood_val, "Landcover Stats": lc_stats,
            "NDVI TS": ndvi_ts
        }
        # --- Save map snapshot ---
        map_bytes = export_map_snapshot(m)
        if map_bytes:
            st.session_state["last_map_snapshot"] = map_bytes
            if "sites" in st.session_state:
                st.session_state["sites"][active]["map_snapshot"] = map_bytes
            st.image(map_bytes, caption="Map Snapshot", use_column_width=True)

        # Display
        st.subheader("📊 Summary")
        st.write(f"**Soil:** {soil_val}")
        st.write(f"**Elevation:** {elev_val}")
        st.write(f"**Seismic:** {seismic_val}")
        st.write(f"**Flood:** {flood_val}")
        st.json(lc_stats)
        if ndvi_ts:
            d,v = zip(*ndvi_ts)
            fig, ax = plt.subplots()
            ax.plot(d, v, marker="o"); ax.set_title("NDVI"); ax.set_xlabel("Date")
            st.pyplot(fig)

# GeoMate Ask (RAG) — simple chat with memory per site and auto-extract numeric values
import re, json, pickle
import streamlit as st
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings

# -------------------
# Load FAISS DB once
# -------------------
@st.cache_resource
def load_faiss():
    # Adjust path to where you unzip faiss_books_db.zip
    faiss_dir = "faiss_books_db"
    # embeddings must match the one you used when creating index
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

    with open(f"{faiss_dir}/index.pkl", "rb") as f:
        data = pickle.load(f)
    vectorstore = FAISS.load_local(faiss_dir, embeddings, allow_dangerous_deserialization=True)
    return vectorstore

vectorstore = load_faiss()


# -------------------
# RAG Chat Page
# -------------------
def rag_page():
    st.header("🤖 GeoMate Ask (RAG + Groq)")
    site = st.session_state["sites"][st.session_state["active_site"]]

    # --- Ensure Site ID exists ---
    if site.get("Site ID") is None:
        site_id = st.session_state["sites"].index(site)
        site["Site ID"] = site_id
    else:
        site_id = site["Site ID"]

    # --- Initialize rag_history properly ---
    if "rag_history" not in st.session_state:
        st.session_state["rag_history"] = {}
    if site_id not in st.session_state["rag_history"]:
        st.session_state["rag_history"][site_id] = []

    # --- Display chat history ---
    hist = st.session_state["rag_history"][site_id]
    for entry in hist:
        who, text = entry.get("who"), entry.get("text")
        if who == "bot":
            st.markdown(f"<div class='chat-bot'>{text}</div>", unsafe_allow_html=True)
        else:
            st.markdown(f"<div class='chat-user'>{text}</div>", unsafe_allow_html=True)

    # --- User input ---
    user_msg = st.text_input("You:", key=f"rag_input_{site_id}")
    if st.button("Send", key=f"rag_send_{site_id}"):
        if not user_msg.strip():
            st.warning("Enter a message.")
        else:
            # Save user msg
            st.session_state["rag_history"][site_id].append(
                {"who": "user", "text": user_msg}
            )

            # --- Retrieve from FAISS ---
            docs = vectorstore.similarity_search(user_msg, k=3)
            context_text = "\n".join([d.page_content for d in docs])

            # --- Build context for LLM ---
            context = {
                "site": {
                    k: v
                    for k, v in site.items()
                    if k in [
                        "Site Name",
                        "lat",
                        "lon",
                        "USCS",
                        "AASHTO",
                        "GI",
                        "Load Bearing Capacity",
                        "Soil Profile",
                        "Flood Data",
                        "Seismic Data",
                    ]
                },
                "chat_history": st.session_state["rag_history"][site_id],
            }

            prompt = (
                f"You are GeoMate AI, an expert geotechnical assistant.\n\n"
                f"Relevant references:\n{context_text}\n\n"
                f"Site context: {json.dumps(context)}\n\n"
                f"User: {user_msg}\n\n"
                f"Answer concisely, include citations [ref:source]. "
                f"If user provides numeric engineering values, return them in the format: [[FIELD: value unit]]."
            )

            # Call your Groq model (replace placeholder with real call)
            resp = groq_generate(prompt, model=st.session_state["llm_model"], max_tokens=500)
            
            # Save bot reply
            st.session_state["rag_history"][site_id].append({"who": "bot", "text": resp})

            # Display reply
            st.markdown(f"<div class='chat-bot'>{resp}</div>", unsafe_allow_html=True)

            # Extract bracketed numeric values
            matches = re.findall(
                r"\[\[([A-Za-z0-9 _/-]+):\s*([0-9.+-eE]+)\s*([A-Za-z%\/]*)\]\]", resp
            )
            for m in matches:
                field, val, unit = m[0].strip(), m[1].strip(), m[2].strip()
                if "bearing" in field.lower():
                    site["Load Bearing Capacity"] = f"{val} {unit}"
                elif "skin" in field.lower():
                    site["Skin Shear Strength"] = f"{val} {unit}"
                elif "compaction" in field.lower():
                    site["Relative Compaction"] = f"{val} {unit}"

            st.success(
                "Response saved ✅ with citations and recognized numeric fields auto-stored in site data."
            )


# -------------------
# Report fields (still needed in reports_page)
# -------------------
REPORT_FIELDS = [
    ("Load Bearing Capacity", "kPa or psf"),
    ("Skin Shear Strength", "kPa"),
    ("Relative Compaction", "%"),
    ("Rate of Consolidation", "mm/yr or days"),
    ("Nature of Construction", "text"),
    ("Borehole Count", "number"),
    ("Max Depth (m)", "m"),
    ("SPT N (avg)", "blows/ft"),
    ("CBR (%)", "%"),
    ("Allowable Bearing (kPa)", "kPa"),
]

# -------------------------------
# Imports
# -------------------------------
import io, re, json, tempfile
from datetime import datetime
from typing import Dict, Any, Optional, List

import streamlit as st
from reportlab.platypus import (
    SimpleDocTemplate, Paragraph, Spacer, PageBreak, Table, TableStyle, Image as RLImage
)
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4
from reportlab.lib.units import mm

# -------------------------------
# PDF Builder
# -------------------------------
def build_full_geotech_pdf(
    site: Dict[str, Any],
    filename: str,
    include_map_image: Optional[bytes] = None,
    ext_refs: Optional[List[str]] = None
):
    """
    Build a professional PDF report using site data + references.
    """
    styles = getSampleStyleSheet()
    title_style = ParagraphStyle(
        "title", parent=styles["Title"], fontSize=20, alignment=1,
        textColor=colors.HexColor("#FF7A00")
    )
    h1 = ParagraphStyle(
        "h1", parent=styles["Heading1"], fontSize=14,
        textColor=colors.HexColor("#1F4E79"), spaceAfter=6
    )
    body = ParagraphStyle("body", parent=styles["BodyText"], fontSize=10.5, leading=13)
    bullet = ParagraphStyle("bullet", parent=body, leftIndent=12, bulletIndent=6)

    doc = SimpleDocTemplate(
        filename, pagesize=A4,
        leftMargin=18*mm, rightMargin=18*mm,
        topMargin=18*mm, bottomMargin=18*mm
    )
    elems = []

    # --- Title page ---
    elems.append(Paragraph("GEOTECHNICAL INVESTIGATION REPORT", title_style))
    elems.append(Spacer(1, 12))
    company = site.get("Company Name", "Client / Company: Not provided")
    contact = site.get("Company Contact", "")
    elems.append(Paragraph(f"<b>{company}</b>", body))
    if contact:
        elems.append(Paragraph(contact, body))
    elems.append(Spacer(1, 12))
    elems.append(Paragraph(f"<b>Project:</b> {site.get('Project Name','-')}", body))
    elems.append(Paragraph(f"<b>Site:</b> {site.get('Site Name','-')}", body))
    elems.append(Paragraph(f"<b>Date:</b> {datetime.today().strftime('%Y-%m-%d')}", body))
    elems.append(PageBreak())

    # --- Table of contents ---
    elems.append(Paragraph("TABLE OF CONTENTS", h1))
    toc_items = [
        "1.0 Introduction",
        "2.0 Site description and geology",
        "3.0 Field investigation & laboratory testing",
        "4.0 Evaluation of geotechnical properties",
        "5.0 Provisional site classification",
        "6.0 Recommendations",
        "7.0 LLM Analysis",
        "8.0 Figures & Tables",
        "9.0 Appendices & References"
    ]
    for i, t in enumerate(toc_items, start=1):
        elems.append(Paragraph(f"{i}. {t}", body))
    elems.append(PageBreak())

    # --- Summary ---
    elems.append(Paragraph("SUMMARY", h1))
    summary_bullets = [
        f"Site: {site.get('Site Name','-')}.",
        f"General geology: {site.get('Soil Profile','Not provided')}.",
        f"Key lab tests: {', '.join([r.get('sampleId','') for r in site.get('Laboratory Results',[])]) if site.get('Laboratory Results') else 'No lab results provided.'}",
        f"Classification: USCS = {site.get('USCS','Not provided')}; AASHTO = {site.get('AASHTO','Not provided')}.",
        "Primary recommendation: See Recommendations section."
    ]
    for s in summary_bullets:
        elems.append(Paragraph(f"• {s}", bullet))
    elems.append(PageBreak())

    # --- Introduction ---
    elems.append(Paragraph("1.0 INTRODUCTION", h1))
    intro_text = site.get("Project Description", "Project description not provided.")
    elems.append(Paragraph(intro_text, body))

    # --- Site description & geology ---
    elems.append(Paragraph("2.0 SITE DESCRIPTION AND GEOLOGY", h1))
    site_geo = [
        f"Topography: {site.get('Topography','Not provided')}",
        f"Drainage: {site.get('Drainage','Not provided')}",
        f"Current land use: {site.get('Current Land Use','Not provided')}",
        f"Regional geology: {site.get('Regional Geology','Not provided')}"
    ]
    for t in site_geo:
        elems.append(Paragraph(t, body))
    elems.append(PageBreak())

    # --- Field & lab testing ---
    elems.append(Paragraph("3.0 FIELD INVESTIGATION & LABORATORY TESTING", h1))
    if site.get("Field Investigation"):
        for item in site["Field Investigation"]:
            elems.append(Paragraph(f"<b>{item.get('id','Test')}</b> — depth {item.get('depth','-')}", body))
            for layer in item.get("layers", []):
                elems.append(Paragraph(f"- {layer.get('depth','')} : {layer.get('description','')}", body))
    else:
        elems.append(Paragraph("No field investigation data supplied.", body))

    lab_rows = site.get("Laboratory Results", [])
    if lab_rows:
        elems.append(Spacer(1, 6))
        elems.append(Paragraph("Laboratory Results", h1))
        data = [["Sample ID","Material","LL","PI","Linear Shrinkage","%Clay","%Silt","%Sand","%Gravel","Expansiveness"]]
        for r in lab_rows:
            data.append([
                r.get("sampleId","-"), r.get("material","-"),
                str(r.get("liquidLimit","-")), str(r.get("plasticityIndex","-")),
                str(r.get("linearShrinkage","-")), str(r.get("percentClay","-")),
                str(r.get("percentSilt","-")), str(r.get("percentSand","-")),
                str(r.get("percentGravel","-")), r.get("potentialExpansiveness","-")
            ])
        t = Table(data, repeatRows=1, colWidths=[40*mm,40*mm,18*mm,18*mm,22*mm,20*mm,20*mm,20*mm,20*mm,30*mm])
        t.setStyle(TableStyle([
            ('BACKGROUND',(0,0),(-1,0),colors.HexColor("#1F4E79")),
            ('TEXTCOLOR',(0,0),(-1,0),colors.white),
            ('GRID',(0,0),(-1,-1),0.4,colors.grey),
            ('BOX',(0,0),(-1,-1),1,colors.HexColor("#FF7A00"))
        ]))
        elems.append(t)
    elems.append(PageBreak())

    # --- Evaluation & classification ---
    elems.append(Paragraph("4.0 EVALUATION OF GEOTECHNICAL PROPERTIES", h1))
    elems.append(Paragraph(site.get("Evaluation","Evaluation not provided."), body))
    elems.append(Paragraph("5.0 PROVISIONAL SITE CLASSIFICATION", h1))
    elems.append(Paragraph(site.get("Provisional Classification","Not provided."), body))
    elems.append(Paragraph("6.0 RECOMMENDATIONS", h1))
    elems.append(Paragraph(site.get("Recommendations","Not provided."), body))

    # --- LLM Analysis ---
    elems.append(Paragraph("7.0 LLM ANALYSIS (GeoMate)", h1))
    llm_text = site.get("LLM_Report_Text", None)
    if llm_text:
        elems.append(Paragraph(llm_text.replace("\n","\n\n"), body))
    else:
        elems.append(Paragraph("No LLM analysis saved for this site.", body))

    # --- Map snapshot ---
    if include_map_image:
        try:
            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
            tmp.write(include_map_image)
            tmp.flush()
            elems.append(PageBreak())
            elems.append(Paragraph("Map Snapshot", h1))
            elems.append(RLImage(tmp.name, width=160*mm, height=90*mm))
        except Exception:
            pass

    # --- References ---
    elems.append(PageBreak())
    elems.append(Paragraph("9.0 APPENDICES & REFERENCES", h1))
    if ext_refs:
        for r in ext_refs:
            elems.append(Paragraph(f"- {r}", body))
    else:
        elems.append(Paragraph("- No external references provided.", body))

    doc.build(elems)
    return filename


# -------------------------------
# Reports Page
# -------------------------------
def reports_page():
    st.header("📑 Reports — Classification & Full Geotechnical")
    site = st.session_state["sites"][st.session_state["active_site"]]

    # ---------------- Classification Report ----------------
    st.subheader("Classification-only report")
    if site.get("classifier_decision"):
        st.markdown("You have a saved classification for this site.")
        if st.button("Generate Classification PDF"):
            fname = f"classification_{site['Site Name'].replace(' ','_')}.pdf"
            buffer = io.BytesIO()
            doc = SimpleDocTemplate(buffer, pagesize=A4)
            elems = []
            elems.append(Paragraph("Soil Classification Report", getSampleStyleSheet()['Title']))
            elems.append(Spacer(1,6))
            elems.append(Paragraph(f"Site: {site.get('Site Name')}", getSampleStyleSheet()['Normal']))
            elems.append(Spacer(1,6))
            elems.append(Paragraph("Classification result:", getSampleStyleSheet()['Heading2']))
            elems.append(Paragraph(site.get("classifier_decision","-"), getSampleStyleSheet()['BodyText']))

            # Add FAISS citations if present in rag_history
            if "rag_history" in st.session_state and site.get("Site ID") in st.session_state["rag_history"]:
                refs = []
                for h in st.session_state["rag_history"][site["Site ID"]]:
                    if h["who"]=="bot" and "[ref:" in h["text"]:
                        for m in re.findall(r"\[ref:([^\]]+)\]", h["text"]):
                            refs.append(m)
                if refs:
                    elems.append(Spacer(1,12))
                    elems.append(Paragraph("References:", getSampleStyleSheet()['Heading2']))
                    for r in set(refs):
                        elems.append(Paragraph(f"- {r}", getSampleStyleSheet()['Normal']))

            doc.build(elems)
            buffer.seek(0)
            st.download_button("Download Classification PDF", buffer, file_name=fname, mime="application/pdf")
    else:
        st.info("No classification saved for this site yet. Use the Classifier page.")

    # ---------------- Quick Report Form ----------------
    st.markdown("### Quick report form (edit values and request LLM analysis)")
    with st.form(key="report_quick_form"):
        cols = st.columns([2,1,1])
        cols[0].markdown("**Parameter**")
        cols[1].markdown("**Value**")
        cols[2].markdown("**Unit / Notes**")

        inputs = {}
        for (fld, unit) in REPORT_FIELDS:
            c1, c2, c3 = st.columns([2,1,1])
            c1.markdown(f"**{fld}**")
            default_val = site.get(fld, "")
            inputs[fld] = c2.text_input(fld, value=str(default_val), label_visibility="collapsed", key=f"quick_{fld}")
            c3.markdown(unit)

        submitted = st.form_submit_button("Save values to site")
        if submitted:
            for fld, _ in REPORT_FIELDS:
                val = inputs.get(fld, "").strip()
                site[fld] = val if val != "" else "Not provided"
            st.success("Saved quick report values to active site.")

    # ---------------- LLM Analysis ----------------
    st.markdown("#### LLM-powered analysis")
    if st.button("Ask GeoMate (generate analysis & recommendations)"):
        context = {
            "site_name": site.get("Site Name"),
            "project": site.get("Project Name"),
            "site_summary": {
                "USCS": site.get("USCS"), "AASHTO": site.get("AASHTO"), "GI": site.get("GI"),
                "Soil Profile": site.get("Soil Profile"),
                "Key lab results": [r.get("sampleId") for r in site.get("Laboratory Results",[])]
            },
            "inputs": {fld: site.get(fld,"Not provided") for fld,_ in REPORT_FIELDS}
        }
        prompt = (
            "You are GeoMate AI, an engineering assistant. Given the following site context and "
            "engineering parameters (some may be 'Not provided'), produce:\n1) short executive summary, "
            "2) geotechnical interpretation (classification, key risks), 3) recommended remedial/improvement "
            "options and 4) short design notes. Provide any numeric outputs in the format [[FIELD: value unit]].\n\n"
            f"Context: {json.dumps(context)}\n\nAnswer concisely and professionally."
        )
        resp = groq_generate(prompt, model=st.session_state["llm_model"], max_tokens=600)
        
        st.markdown("**GeoMate analysis**")
        st.markdown(resp)

        matches = re.findall(r"\[\[([A-Za-z0-9 _/-]+):\s*([0-9.+-eE]+)\s*([A-Za-z%\/]*)\]\]", resp)
        for m in matches:
            field, val, unit = m[0].strip(), m[1].strip(), m[2].strip()
            if "bearing" in field.lower():
                site["Load Bearing Capacity"] = f"{val} {unit}"
            elif "skin" in field.lower():
                site["Skin Shear Strength"] = f"{val} {unit}"
            elif "compaction" in field.lower():
                site["Relative Compaction"] = f"{val} {unit}"

        site["LLM_Report_Text"] = resp
        st.success("LLM analysis saved to site under 'LLM_Report_Text'.")

    # ---------------- Full Report Chatbot ----------------
    st.markdown("---")
    st.subheader("Full Geotechnical Report (chatbot will gather missing fields)")
    if st.button("Start Report Chatbot"):
        st.session_state["sites"][st.session_state["active_site"]]["report_convo_state"] = 0
        st.rerun()

    state = site.get("report_convo_state", -1)
    if state >= 0:
        st.markdown("Chatbot will ask for missing fields. You can answer or type 'skip' to leave blank.")
        show_table = [(f, site.get(f, "Not provided")) for f,_ in REPORT_FIELDS]
        st.table(show_table)

        if state < len(REPORT_FIELDS):
            field, unit = REPORT_FIELDS[state]
            ans = st.text_input(f"GeoMate — Please provide '{field}' ({unit})", key=f"report_in_{state}")
            c1, c2 = st.columns([1,1])
            with c1:
                if st.button("Submit", key=f"report_submit_{state}"):
                    site[field] = ans.strip() if ans.strip() not in ("skip","don't know","dont know","na","n/a","") else "Not provided"
                    site["report_convo_state"] = state + 1
                    st.rerun()
            with c2:
                if st.button("Skip", key=f"report_skip_{state}"):
                    site[field] = "Not provided"
                    site["report_convo_state"] = state + 1
                    st.rerun()
        else:
            st.success("All report questions asked. You can generate the full report now.")
            ext_ref_text = st.text_area("Optional: External references (one per line)", value="")
            ext_refs = [r.strip() for r in ext_ref_text.splitlines() if r.strip()]

            faiss_refs = []
            if "rag_history" in st.session_state and site.get("Site ID") in st.session_state["rag_history"]:
                for h in st.session_state["rag_history"][site["Site ID"]]:
                    if h["who"]=="bot" and "[ref:" in h["text"]:
                        for m in re.findall(r"\[ref:([^\]]+)\]", h["text"]):
                            faiss_refs.append(m)
            all_refs = list(set(ext_refs + faiss_refs))

            outname = f"Full_Geotech_Report_{site.get('Site Name','site')}.pdf"
            mapimg = site.get("map_snapshot")

            build_full_geotech_pdf(site, outname, include_map_image=mapimg, ext_refs=all_refs)

            with open(outname, "rb") as f:
                st.download_button("Download Full Geotechnical Report", f, file_name=outname, mime="application/pdf")
                
# 8) Page router
if "page" not in st.session_state:
    st.session_state["page"] = "Home"

page = st.session_state["page"]

# Option menu top (main nav)
selected = option_menu(
    None,
    ["Home","Soil recognizer","Classifier","GSD","OCR","Locator","RAG","Reports"],
    icons=["house","chart","journal-code","bar-chart","camera","geo-alt","robot","file-earmark-text"],
    menu_icon="cast",
    default_index=["Home","Soil recognizer","Classifier","GSD","OCR","Locator","RAG","Reports"].index(page) if page in ["Home","Soil recognizer","Classifier","GSD","OCR","Locator","RAG","Reports"] else 0,
    orientation="horizontal",
    styles={
        "container": {"padding":"0px","background-color":"#0b0b0b"},
        "nav-link": {"font-size":"14px","color":"#cfcfcf"},
        "nav-link-selected": {"background-color":"#FF7A00","color":"white"},
    }
)
st.session_state["page"] = selected
page = selected

# Display page content
if page == "Home":
    landing_page()
elif page == "Classifier":
    soil_classifier_page()
elif page == "GSD":
    gsd_page()
elif page == "OCR":
    ocr_page()
elif page == "Locator":
    locator_page()
elif page == "RAG":
    rag_page()
elif page == "Reports":
    reports_page()
elif page == "Soil recognizer":
    soil_recognizer_page()
else:
    landing_page()

# Footer
st.markdown("<hr/>", unsafe_allow_html=True)
st.markdown("<div style='text-align:center;color:#9aa7bf'>GeoMate V2 • AI geotechnical copilot • Built for HF Spaces</div>", unsafe_allow_html=True)