# app.py
import streamlit as st
import folium
from streamlit_folium import st_folium
from model import predict_metals
import pandas as pd

st.set_page_config(page_title="GeoMetals Predictor", layout="wide")

st.title("🧪 GeoMetals Exposure Predictor")
st.write("Click a location on the map and select stormwater runoff status to predict heavy metal concentrations.")

# Map setup
st.subheader("📍 Select Location")
m = folium.Map(location=[33.45, -112.07], zoom_start=10)  # Metro Phoenix default center

marker = folium.LatLngPopup()
m.add_child(marker)

map_data = st_folium(m, width=700, height=500)

# User location input
lat, lon = None, None
if map_data["last_clicked"]:
    lat = map_data["last_clicked"]["lat"]
    lon = map_data["last_clicked"]["lng"]
    st.success(f"Selected Coordinates: ({lat:.5f}, {lon:.5f})")

# Stormwater input
stormwater = st.radio("Stormwater Runoff Present?", ("Yes", "No")) == "Yes"

# Submit and predict
if st.button("🔍 Predict Heavy Metal Concentrations") and lat and lon:
    stormwater_type = "stormwater" if stormwater else "no_stormwater"
    input_data = pd.DataFrame([{
        "lat": lat,
        "lon": lon,
        "type": stormwater_type
    }])
    prediction = predict_metals(input_data)

    st.subheader("🔬 Predicted Heavy Metal Concentrations (ppm)")
    for metal, value in prediction.items():
        warning = ""
        if metal == "Fe_ppm" and value > 300:
            warning = "⚠️ Above EPA soil screening level for Iron"
        elif metal == "Cr_ppm" and value > 210:
            warning = "⚠️ Chromium exceeds safe threshold"
        elif metal == "Mn_ppm" and value > 500:
            warning = "⚠️ High Manganese levels"
        elif metal == "Mo_ppm" and value > 40:
            warning = "⚠️ Elevated Molybdenum"
        elif metal == "In_ppm" and value > 0.5:
            warning = "⚠️ Indium concern"
        elif metal == "Ta_ppm" and value > 1:
            warning = "⚠️ Tantalum elevated"
        st.write(f"**{metal.replace('_ppm', '')}: {value:.2f} ppm** {warning}")

    st.subheader("📉 Legacy Risk Score (Experimental)")
    score = sum(prediction.values()) / len(prediction)
    st.metric("Aggregate Risk Score", f"{score:.2f}")
else:
    st.info("Click a location on the map and hit 'Predict'.")