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

# 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

# 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 = get_active_site()  # 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"):
            site["Soil Profile"] = predicted_class
            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}

# PDF builder (reportlab) — create professional document similar to uploaded sample
def build_full_geotech_pdf(site: Dict[str, Any], filename: str, include_map_image: Optional[bytes]=None, ext_refs: Optional[List[str]]=None):
    """
    site: dictionary of site data
    filename: output file path
    include_map_image: bytes of image to embed (optional)
    ext_refs: list of external refs (strings)
    """
    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 = []
    # Cover
    elems.append(Paragraph("GEOTECHNICAL INVESTIGATION REPORT", title_style))
    elems.append(Spacer(1,6))
    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(Spacer(1,8))
    elems.append(Paragraph("<b>Prepared by:</b> GeoMate AI", 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())

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

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

    # 3.0 Field investigation and laboratory 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 table
    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())

    # 4.0 Evaluation & 5.0 Classification & 6.0 Recommendations
    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))

    # Map
    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

    # External refs
    if ext_refs:
        elems.append(PageBreak())
        elems.append(Paragraph("References", h1))
        for r in ext_refs:
            elems.append(Paragraph(f"- {r}", body))

    doc.build(elems)
    return filename

# 4) Session state initialization
if "sites" not in st.session_state:
    # initialize with a default site
    st.session_state["sites"] = [{
        "Soil Profile": None,
        "Site Name":"Home",
        "Project Name":"Demo Project",
        "Site ID": 0,
        "Coordinates":"",
        "lat": None,
        "lon": None,
        "Project Description":"",
        "Topography":"",
        "Drainage":"",
        "Current Land Use":"",
        "Regional Geology":"",
        "Field Investigation": [],
        "Laboratory Results": [],
        "GSD": None,
        "USCS": None,
        "AASHTO": None,
        "GI": 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,
        "Environmental Data": None,
        "Topo Data": None,
        "map_snapshot": None,
        "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"] = "meta-llama/llama-4-maverick-17b-128e-instruct"

# 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",
        "llama3-8b-8192",
        "gemma-7b-it",
        "mixtral-8x7b-32768"
    ], 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"])
            st.session_state["sites"].append({
                "Site Name": new_site_name.strip(),
                "Project Name": "Project - " + new_site_name.strip(),
                "Site ID": idx,
                "Coordinates":"",
                "lat": None,
                "lon": None,
                "Project Description":"",
                "Topography":"",
                "Drainage":"",
                "Current Land Use":"",
                "Regional Geology":"",
                "Field Investigation": [],
                "Laboratory Results": [],
                "GSD": None,
                "USCS": None,
                "AASHTO": None,
                "GI": 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,
                "Environmental Data": None,
                "Topo Data": None,
                "map_snapshot": None,
                "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():
    st.markdown("<div style='display:flex;align-items:center;gap:12px'>"
                "<div style='width:76px;height:76px;border-radius:14px;background:linear-gradient(135deg,#ff7a00,#ff3a3a);display:flex;align-items:center;justify-content:center;box-shadow:0 8px 24px rgba(0,0,0,0.6)'>"
                "<span style='font-size:34px'>🛰️</span></div>"
                "<div><h1 style='margin:0;color:#FF8C00'>GeoMate V2</h1>"
                "<div class='small-muted'>AI geotechnical copilot — soil recognition, classification, locator, RAG, and reports</div></div></div>", unsafe_allow_html=True)
    st.markdown("---")
    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: classify soils (USCS/AASHTO), plot GSD, fetch Earth Engine data, chat with a RAG-backed LLM, and generate professional geotechnical reports.")
        st.markdown("</div>", unsafe_allow_html=True)
        st.markdown("### Quick Actions")
        c1, c2, c3 = st.columns(3)
        if c1.button("🧪 Classifier"):
            st.session_state["page"] = "Classifier"; st.rerun()
        if c2.button("📈 GSD Curve"):
            st.session_state["page"] = "GSD"; st.rerun()
        if c3.button("🌍 Locator"):
            st.session_state["page"] = "Locator"; st.rerun()
        c4, c5, c6 = st.columns(3)
        if c4.button("🤖 GeoMate Ask"):
            st.session_state["page"] = "RAG"; st.rerun()
        if c5.button("📷 OCR"):
            st.session_state["page"] = "OCR"; st.rerun()
        if c6.button("📑 Reports"):
            st.session_state["page"] = "Reports"; st.rerun()
    with col2:
        st.markdown("<div class='gm-card' style='text-align:center'>", unsafe_allow_html=True)
        st.markdown("<h3 style='color:#FF8C00'>Live Site Summary</h3>", unsafe_allow_html=True)
        site = st.session_state["sites"][st.session_state["active_site"]]
        st.write(f"Site: **{site.get('Site Name')}**")
        st.write(f"USCS: {site.get('USCS')}, 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 (Earth Engine integration)
def locator_page():
    st.header("🌍 Locator — Select Area of Interest")
    site = st.session_state["sites"][st.session_state["active_site"]]
    st.markdown("You can enter coordinates manually or draw/upload a GeoJSON boundary (draw-mode not available in this minimal example).")
    lat = st.number_input("Latitude", value=site.get("lat") or 0.0, format="%.6f", key="locator_lat")
    lon = st.number_input("Longitude", value=site.get("lon") or 0.0, format="%.6f", key="locator_lon")
    site["lat"] = lat; site["lon"] = lon
    if st.button("Fetch Earth Data (EE)"):
        if not EE_AVAILABLE:
            st.error("Earth Engine/Geemap not available in this environment. Please add geemap & earthengine-api to requirements.")
        else:
            try:
                # Initialize EE if needed
                if not ee.data._credentials:
                    # try to authenticate via service account string in env
                    ee_key_json = os.environ.get("EARTH_ENGINE_KEY")
                    # if provided as JSON string, write to temp file
                    tmp = tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".json")
                    tmp.write(ee_key_json)
                    tmp.flush()
                    ee.Initialize(service_account=os.environ.get("SERVICE_ACCOUNT"), private_key_json=tmp.name)
                # sample calls (very simple placeholders)
                st.info("Querying Earth Engine for available layers (this may take a few seconds)...")
                # NOTE: Detailed real EE reducers and datasets must be added for production.
                # We'll store placeholder results and attempt to create a simple map snapshot.
                site["Soil Profile"] = f"Colluvial soils over dolomite (fetched {datetime.today().date()})"
                site["Flood Data"] = "No major floods in last 20 years (placeholder)"
                site["Seismic Data"] = "Seismic zone: Moderate; historic events: low"
                site["Environmental Data"] = "No major environmental constraints found (placeholder)"
                site["Topo Data"] = "Gentle slope"
                # Create a simple map via geemap and snapshot it
                m = geemap.Map(center=[lat, lon], zoom=10)
                m.add_basemap("SATELLITE")
                # capture map to PNG (geemap has method to export html; we use to_image if available)
                try:
                    img = m.to_image()
                    buf = io.BytesIO()
                    img.save(buf, format="PNG")
                    img_bytes = buf.getvalue()
                    site["map_snapshot"] = img_bytes
                    st.image(img_bytes, caption="Map snapshot", use_column_width=True)
                except Exception:
                    st.warning("Map snapshot not available in this environment.")
                st.success("Earth Engine data fetched and saved to site (placeholders).")
            except Exception as e:
                st.error(f"Earth Engine error: {e}")

# GeoMate Ask (RAG) — simple chat with memory per site and auto-extract numeric values
def rag_page():
    st.header("🤖 GeoMate Ask (RAG + Groq)")
    site = st.session_state["sites"][st.session_state["active_site"]]
    st.markdown("Chat with GeoMate. The LLM has memory per site for this session; any engineering values provided will be parsed and saved.")
    if "rag_history" not in st.session_state:
        st.session_state["rag_history"] = {i: [] for i in range(len(st.session_state["sites"]))}
    hist = st.session_state["rag_history"].get(site["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_msg = st.text_input("You:", key="rag_input")
    if st.button("Send", key="rag_send"):
        if not user_msg.strip():
            st.warning("Enter a message.")
        else:
            # Save user msg
            st.session_state["rag_history"][site["Site ID"]].append({"who":"user","text":user_msg})
            # Build prompt including site context
            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["Site ID"]]
            }
            prompt = f"You are GeoMate AI. Site context: {json.dumps(context)}. User: {user_msg}\nRespond professionally and concisely. If user provides numeric engineering values, return them in the format: [[FIELD: value unit]]."
            resp = groq_generate(prompt, model=st.session_state["llm_model"], max_tokens=400)
            # Save bot reply
            st.session_state["rag_history"][site["Site ID"]].append({"who":"bot","text":resp})
            # Display
            st.markdown(f"<div class='chat-bot'>{resp}</div>", unsafe_allow_html=True)
            # Try to extract bracketed fields like [[Load bearing capacity: 2000 psf]]
            import re
            matches = re.findall(r"\[\[([A-Za-z0-9 _/-]+):\s*([0-9.+-eE]+)\s*([A-Za-z%\/]*)\]\]", resp)
            for m in matches:
                field = m[0].strip()
                val = m[1].strip()
                unit = m[2].strip()
                # Map common fields to site keys
                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 and any recognized numeric fields auto-stored in the site data.")

# Reports page — conversational missing-parameter bot & PDF generation
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")
]

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

    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"
            # simple 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']))
            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.")

    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()

    # Conversational data collection
    state = site.get("report_convo_state", 0)
    if site.get("report_convo_state") is not None:
        st.markdown("Chatbot will ask for missing fields. You can answer or type 'skip' to leave blank.")
        # Show current known fields
        st.write("Current key parameters (live):")
        show_table = []
        for f,_ in REPORT_FIELDS:
            show_table.append((f, site.get(f, "Not provided")))
        st.table(show_table)
        # continue conversation step-by-step
        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}"):
                    if ans.strip().lower() in ("skip","don't know","dont know","na","n/a",""):
                        site[field] = "Not provided"
                    else:
                        site[field] = ans.strip()
                    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.")
            if st.button("Generate Full Geotechnical Report PDF"):
                # Prepare ext_refs
                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()]
                # Build PDF using reportlab builder
                outname = f"Full_Geotech_Report_{site.get('Site Name','site')}.pdf"
                # include map image bytes if available
                mapimg = site.get("map_snapshot")
                build_full_geotech_pdf(site, outname, include_map_image=mapimg, ext_refs=ext_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","Classifier","GSD","OCR","Locator","RAG","Reports"],
    icons=["house","journal-code","bar-chart","camera","geo-alt","robot","file-earmark-text"],
    menu_icon="cast",
    default_index=["Home","Classifier","GSD","OCR","Locator","RAG","Reports"].index(page) if page in ["Home","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)