app.py

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

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

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

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

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

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

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

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

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

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

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

# Helpful alias for session state
ss = st.session_state

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    return result_text, uscs, aashto, GI, char_summary

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

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

    col_up, col_manual = st.col