import gradio as gr
import rasterio
import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objs as go
import logging

# Configure logging for better debugging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def process_dem(dem_path):
    """
    Analyzes a DEM file to generate a 2D slope risk map and a 3D interactive plot.
    """
    if dem_path is None:
        return None, None
        
    try:
        # --- OPEN DEM ---
        logger.info(f"Processing file: {dem_path.name}")
        with rasterio.open(dem_path.name) as src:
            dem = src.read(1).astype(float)
            profile = src.profile
            logger.info("Successfully opened DEM file.")
    except rasterio.errors.RasterioIOError as e:
        logger.error(f"Rasterio error: Failed to open or read the DEM file. Error: {e}")
        # Return an error message to be displayed by Gradio
        raise gr.Error("Failed to process the DEM file. Please ensure it is a valid GeoTIFF (.tif) file.")
    except Exception as e:
        logger.error(f"An unexpected error occurred during file processing: {e}")
        raise gr.Error(f"An unexpected error occurred: {e}")

    nrows, ncols = dem.shape
    Z = dem

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

    # --- RISK MASK ---
    try:
        threshold = np.percentile(slope, 95)  # Top 5% steepest slopes
        risk_mask = slope > threshold
    except IndexError:
        # Handle case where slope array is empty or too small
        logger.warning("Slope array is empty, skipping percentile calculation.")
        risk_mask = np.zeros_like(slope, dtype=bool)

    # --- 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)
    logger.info("Generated 2D risk map.")

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

    # Base DEM surface
    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))
    ))

    # Risk overlay (purple)
    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"
        )
    )
    logger.info("Generated 3D Plotly figure.")

    return risk_map_path, fig3d

# --- GRADIO APP ---
demo = gr.Interface(
    fn=process_dem,
    inputs=gr.File(label="Upload DEM (.tif)", file_types=[".tif"]),
    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()