import gradio as gr
import rasterio
import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objs as go
from gradio.exceptions import Error
from rasterio.errors import RasterioIOError

def process_dem(dem_file):
    # dem_file can be a tempfile-like object, or a str path (depending on gr.File type)
    if isinstance(dem_file, str):
        path = dem_file
    else:
        # Gradio hands a TemporaryFile object; use its .name
        path = getattr(dem_file, "name", None)

    if not path:
        raise Error("No file received. Please upload a GeoTIFF (.tif/.tiff).")

    # Robust validation by CONTENT (not extension)
    try:
        with rasterio.open(path) as src:
            # Accept common GeoTIFF drivers
            if src.driver not in ("GTiff", "COG"):
                raise Error(f"Unsupported raster driver: {src.driver}. Please upload a GeoTIFF.")
            dem = src.read(1).astype(float)
    except RasterioIOError as e:
        raise Error(f"Could not read the file as a GeoTIFF. Details: {e}")

    nrows, ncols = dem.shape
    Z = dem

    # --- SLOPE ---
    dy, dx = np.gradient(Z)
    slope = np.sqrt(dx**2 + dy**2)

    # --- RISK MASK (top 5% steepest) ---
    threshold = np.percentile(slope, 95)
    risk_mask = slope > threshold

    # --- 2D RISK MAP ---
    fig2d, ax = plt.subplots(figsize=(8, 6))
    c = ax.imshow(slope, cmap="hot", origin="upper")
    ax.contour(risk_mask, levels=[0.5], colors="blue", linewidths=0.8)
    plt.colorbar(c, ax=ax, label="Slope (steepness)")
    ax.set_title("Slope Risk Map (Hot = Steep, Blue = Risk zones)")
    ax.set_xlabel("Column Index (X)")
    ax.set_ylabel("Row Index (Y)")
    risk_map_path = "risk_map.png"
    plt.savefig(risk_map_path, dpi=150, bbox_inches="tight")
    plt.close(fig2d)

    # --- INTERACTIVE 3D DEM (Plotly) ---
    step = max(1, nrows // 200)
    fig3d = go.Figure()

    fig3d.add_trace(go.Surface(
        z=Z[::step, ::step],
        colorscale="Earth",
        showscale=True,
        opacity=0.9,
        contours=dict(z=dict(show=True, usecolormap=True, highlightcolor="black", project_z=True))
    ))

    fig3d.add_trace(go.Surface(
        z=np.where(risk_mask[::step, ::step], Z[::step, ::step], np.nan),
        surfacecolor=np.ones_like(Z[::step, ::step]),
        colorscale=[[0, "purple"], [1, "purple"]],
        showscale=False,
        opacity=0.6
    ))

    fig3d.update_layout(
        title="Interactive 3D DEM with Contours & Steep Slope Highlight",
        scene=dict(
            xaxis_title="X (grid cols)",
            yaxis_title="Y (grid rows)",
            zaxis_title="Elevation (m)",
            aspectmode="data"
        )
    )

    return risk_map_path, fig3d


# --- GRADIO APP ---
demo = gr.Interface(
    fn=process_dem,
    # 🔑 Remove file_types to bypass extension-based rejection entirely
    inputs=gr.File(label="Upload DEM (GeoTIFF .tif/.tiff)"),
    outputs=[
        gr.Image(type="filepath", label="2D Slope Risk Map"),
        gr.Plot(label="Interactive 3D DEM (with Contours & Risk Zones)")
    ],
    title="3D DEM & Landslide Risk Visualizer",
    description="Upload a GeoTIFF DEM file to see a 2D slope risk map and an interactive 3D DEM with contours & steep slope zones highlighted."
)

if __name__ == "__main__":
    demo.launch()