app2.py

# app.py
import os
import io
import json
import tempfile
import torch
import requests
import numpy as np
import pandas as pd
import gradio as gr
from PIL import Image
import torchvision.transforms as T
from geopy.geocoders import Nominatim

# --- Import model ---
from src.model import CompactGeoEmbed

# --- Config ---


MODEL_PATH = "model/geo_model.pth"
GEE_KEY_PATH = os.environ.get("GEE_KEY_PATH", "keys/gee_service_account.json")
DEVICE = torch.device("cpu")
TF_SIZE = (120, 120)

# --- Load local model ---
def load_model():
    model = CompactGeoEmbed(32, 96)
    try:
        state = torch.load(MODEL_PATH, map_location=DEVICE)
        model.load_state_dict(state)
        print("✅ Model loaded successfully from local path.")
    except Exception as e:
        print("⚠️ Model load failed:", e)
    model.to(DEVICE).eval()
    return model

MODEL = load_model()
tf = T.Compose([T.Resize(TF_SIZE), T.ToTensor()])

# --- Reverse geocode (city, country info) ---
def reverse_geocode(lat, lon):
    try:
        geolocator = Nominatim(user_agent="geo-risk-app", timeout=10)
        loc = geolocator.reverse((lat, lon), language="en")
        return loc.address if loc else "Unknown location"
    except Exception as e:
        print("⚠️ Reverse geocode failed:", e)
        return "Unknown location"

# --- Earth Engine init with service account ---
def init_gee():
    try:
        import ee
    except Exception as e:
        print("⚠️ earthengine-api not installed:", e)
        return False

    # Try to read credentials from Hugging Face secret first
    gee_secret = os.environ.get("GEE_KEY_JSON")

    if gee_secret:
        try:
            svc = json.loads(gee_secret)
            credentials = ee.ServiceAccountCredentials(svc["client_email"], key_data=gee_secret)
            ee.Initialize(credentials)
            print("✅ Earth Engine initialized from Hugging Face secret.")
            return True
        except Exception as e:
            print("⚠️ GEE secret init failed:", e)

    # Fallback to local file if running locally
    if os.path.exists(GEE_KEY_PATH):
        try:
            with open(GEE_KEY_PATH, "r") as f:
                svc = json.load(f)
            credentials = ee.ServiceAccountCredentials(svc["client_email"], GEE_KEY_PATH)
            ee.Initialize(credentials)
            print("✅ Earth Engine initialized from local key.")
            return True
        except Exception as e:
            print("⚠️ Local GEE initialization failed:", e)

    print("⚠️ GEE credentials not found (neither secret nor file).")
    return False


GEE_READY = init_gee()

# --- Fetch satellite & elevation tiles from GEE ---
def fetch_gee_images(lat, lon):
    try:
        if not GEE_READY:
            raise RuntimeError("GEE not initialized")

        import ee

        p = ee.Geometry.Point([lon, lat])
        region = p.buffer(1000).bounds()

        s2 = (
            ee.ImageCollection("COPERNICUS/S2_SR_HARMONIZED")
            .filterBounds(p)
            .filterDate("2024-01-01", "2024-12-31")
            .sort("CLOUDY_PIXEL_PERCENTAGE")
            .first()
            .select(["B4", "B3", "B2"])
        )
        s2_url = s2.visualize(min=0, max=3000).getThumbURL(
            {"region": region, "dimensions": f"{TF_SIZE[0]}x{TF_SIZE[1]}", "format": "png"}
        )
        s2_img = Image.open(io.BytesIO(requests.get(s2_url, timeout=10).content)).convert("RGB")

        srtm = ee.Image("USGS/SRTMGL1_003")
        elev_url = srtm.visualize(min=0, max=3000).getThumbURL(
            {"region": region, "dimensions": f"{TF_SIZE[0]}x{TF_SIZE[1]}", "format": "png"}
        )
        elev_img = Image.open(io.BytesIO(requests.get(elev_url, timeout=10).content)).convert("L")

        return s2_img, elev_img
    except Exception as e:
        print("⚠️ GEE fetch failed:", e)
        rgb = Image.new("RGB", TF_SIZE, (127, 127, 127))
        elev = Image.new("L", TF_SIZE, 127)
        return rgb, elev

# --- Weather API ---
def fetch_weather(lat, lon):
    try:
        url = (
            f"https://api.open-meteo.com/v1/forecast?"
            f"latitude={lat}&longitude={lon}&daily=temperature_2m_max,"
            "precipitation_sum,relative_humidity_2m_mean,wind_speed_10m_max"
            "&forecast_days=1&timezone=UTC"
        )
        r = requests.get(url, timeout=8)
        d = r.json().get("daily", {})
        return {
            "temp_max": float(d.get("temperature_2m_max", [None])[0]) if d else None,
            "precip": float(d.get("precipitation_sum", [None])[0]) if d else None,
            "humidity": float(d.get("relative_humidity_2m_mean", [None])[0]) if d else None,
            "wind_speed": float(d.get("wind_speed_10m_max", [None])[0]) if d else None,
        }
    except Exception as e:
        print("⚠️ Weather fetch failed:", e)
        return {"temp_max": None, "precip": None, "humidity": None, "wind_speed": None}

# --- Preprocess image ---
def preprocess(img):
    img = img.convert("RGB")
    return tf(img).unsqueeze(0)

# --- IP-based location fallback ---
def get_ip_location():
    try:
        r = requests.get("https://ipapi.co/json", timeout=5)
        data = r.json()
        return round(float(data["latitude"]), 5), round(float(data["longitude"]), 5)
    except Exception as e:
        print("⚠️ IP location fetch failed:", e)
        return 51.5072, -0.1276  # fallback to London

# --- Inference ---
def predict(lat, lon, img):
    # fallback if user didn't fill lat/lon
    if lat is None or lon is None:
        lat, lon = get_ip_location()

    lat = round(float(lat), 5)
    lon = round(float(lon), 5)

    location_str = reverse_geocode(lat, lon)
    rgb_img, elev_img = fetch_gee_images(lat, lon)

    # user-uploaded image takes priority
    if img is not None:
        rgb_img = img.resize(TF_SIZE)

    x = preprocess(rgb_img).to(DEVICE)
    e = torch.tensor(np.array(elev_img) / 255.0, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(DEVICE)

    with torch.no_grad():
        try:
            _, _, r = MODEL(x, e)
            risk_score = float(r.item())
        except Exception as e:
            print("⚠️ Model inference failed:", e)
            risk_score = None

    weather = fetch_weather(lat, lon)
    result = {
        "Location": location_str,
        "Latitude": lat,
        "Longitude": lon,
        "Predicted_Risk_Score": round(risk_score, 4) if risk_score is not None else None,
        **weather,
    }

    df = pd.DataFrame([result])
    with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as tmp:
        tmp_path = tmp.name
    df.to_csv(tmp_path, index=False)

    return rgb_img, df, tmp_path

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("## 🌍 Geo-Risk Predictor (Local Model + GEE + Location Auto-Detect)")

    with gr.Row():
        lat = gr.Number(value=None, label="Latitude")
        lon = gr.Number(value=None, label="Longitude")
        get_loc_btn = gr.Button("📍 Use My Location")

    img = gr.Image(type="pil", label=f"Optional RGB Tile (auto-resized to {TF_SIZE[0]}×{TF_SIZE[1]})")
    run_btn = gr.Button("Run Prediction")

    rgb_preview = gr.Image(label="Satellite Image Used")
    output_df = gr.DataFrame(label="Predicted Data", interactive=False)
    file_out = gr.File(label="Download CSV")

    # main prediction button
    run_btn.click(fn=predict, inputs=[lat, lon, img], outputs=[rgb_preview, output_df, file_out])

    # JS geolocation with IP fallback
    get_loc_btn.click(
        None,
        [],
        [lat, lon],
        js="""
        async () => {
            if (navigator.geolocation) {
                try {
                    const pos = await new Promise((res, rej) =>
                        navigator.geolocation.getCurrentPosition(res, rej)
                    );
                    return [pos.coords.latitude.toFixed(5), pos.coords.longitude.toFixed(5)];
                } catch (err) {
                    console.warn("Browser geolocation failed, fallback to IP API.");
                    const ip = await fetch("https://ipapi.co/json");
                    const data = await ip.json();
                    return [data.latitude.toFixed(5), data.longitude.toFixed(5)];
                }
            } else {
                const ip = await fetch("https://ipapi.co/json");
                const data = await ip.json();
                return [data.latitude.toFixed(5), data.longitude.toFixed(5)];
            }
        }
        """,
    )

if __name__ == "__main__":
#     # demo.launch(server_name="0.0.0.0",
#     #             server_port=int(os.environ.get("PORT", 7860)),
#     #             share=True)
    
    demo.launch(share=True)


# if __name__ == "__main__":
#     # demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
#     demo.launch()