Py.Cafe

import io, json
import numpy as np
import pandas as pd
import streamlit as st
import plotly.graph_objects as go
import tifffile as tiff

st.set_page_config(page_title="DEM Raise Simulator (Lite)", layout="wide")
st.title("DEM Raise Simulator — Pyodide-safe")

st.markdown(
    "Upload a **pre-clipped DEM (GeoTIFF, float32)** and one or more **mask TIFFs** "
    "(same width/height as DEM, values 0/1). Each mask can be applied as a "
    "**SET to target elevation** or an **ADD of delta feet**."
)

# ------------------ Sidebar inputs ------------------
with st.sidebar:
    dem_file = st.file_uploader("DEM (GeoTIFF, float32)", type=["tif", "tiff"])
    st.markdown("### Masks (same grid as DEM)")
    n_masks = st.number_input("How many masks?", min_value=0, max_value=12, value=2, step=1)

    mask_specs = []
    for i in range(n_masks):
        st.markdown(f"**Mask {i+1}**")
        mf = st.file_uploader(f"Mask {i+1} (0/1 TIFF)", type=["tif", "tiff"], key=f"mask{i}")
        mode = st.selectbox("Mode", ["set", "add"], key=f"mode{i}",
                            help="set: raise up to target elevation; add: add delta feet")
        value = st.number_input("Value (ft)", value=16.0, step=0.5, key=f"val{i}")
        mask_specs.append(dict(file=mf, mode=mode, value=float(value)))

    px_area_ft2 = st.number_input("Cell area (ft²) [optional, for volumes]", min_value=0.0, value=0.0,
                                  help="If unknown, leave 0 to report area in number of cells only.")
    run = st.button("Run simulation", type="primary")

# ------------------ Helpers ------------------
def summarize(vals, px_area_ft2):
    out = {}
    if vals.size == 0:
        return dict(n_cells=0, min=np.nan, median=np.nan, mean=np.nan, p95=np.nan, max=np.nan,
                    area_ft2=0.0, vol_ft3=0.0)
    out["n_cells"] = int(vals.size)
    out["min"] = float(np.min(vals))
    out["median"] = float(np.median(vals))
    out["mean"] = float(np.mean(vals))
    out["p95"] = float(np.percentile(vals, 95))
    out["max"] = float(np.max(vals))
    if px_area_ft2 > 0:
        out["area_ft2"] = float(vals.size) * px_area_ft2
        out["vol_ft3"] = float(np.sum(vals) * px_area_ft2)
    else:
        out["area_ft2"] = 0.0
        out["vol_ft3"] = 0.0
    return out

def to_tiff_bytes(arr):
    bio = io.BytesIO()
    # Note: we don’t preserve georeferencing here (no rasterio). This is a plain TIFF.
    tiff.imwrite(bio, arr.astype(np.float32))
    bio.seek(0)
    return bio

# ------------------ Core ------------------
if run:
    if dem_file is None:
        st.error("Please upload a DEM GeoTIFF.")
        st.stop()

    # Load DEM
    dem = tiff.imread(dem_file)  # expects 2D float32
    if dem.ndim != 2:
        st.error("DEM must be a single 2D band (height x width).")
        st.stop()

    H, W = dem.shape
    modified = dem.copy().astype(np.float32)
    raise_depth = np.full_like(modified, np.nan, dtype=np.float32)

    layer_rows = []

    # Apply masks in listed order (order = priority; later rows overwrite earlier in 'set')
    for spec in mask_specs:
        if spec["file"] is None:
            continue
        mask = tiff.imread(spec["file"])
        if mask.shape != (H, W):
            st.error(f"Mask shape {mask.shape} does not match DEM shape {(H, W)}.")
            st.stop()

        mask_bool = mask > 0
        mode = spec["mode"].lower()
        val = float(spec["value"])

        if mode == "set":
            old = modified[mask_bool]
            delta = val - old
            pos = delta > 0
            new_vals = old.copy()
            new_vals[pos] = val
            modified[mask_bool] = new_vals

            layer_raise = np.full_like(modified, np.nan, dtype=np.float32)
            layer_raise[mask_bool] = np.where(pos, delta, np.nan)

        elif mode == "add":
            layer_raise = np.full_like(modified, np.nan, dtype=np.float32)
            layer_raise[mask_bool] = val
            modified[mask_bool] = modified[mask_bool] + val
        else:
            st.error(f"Unknown mode: {mode}")
            st.stop()

        # accumulate only positive raises
        rd = np.nan_to_num(raise_depth, nan=0.0)
        add = np.nan_to_num(layer_raise, nan=0.0)
        rd += np.maximum(add, 0.0)
        rd[rd == 0] = np.nan
        raise_depth = rd

        vals = layer_raise[~np.isnan(layer_raise)]
        stats = summarize(vals, px_area_ft2)
        stats.update(dict(name=spec["file"].name, mode=mode, value=val))
        layer_rows.append(stats)

    # Summaries
    st.subheader("Per-layer summary")
    if layer_rows:
        df = pd.DataFrame(layer_rows)
        if px_area_ft2 > 0:
            df["yd3"] = df["vol_ft3"] / 27.0
            df["m3"] = df["vol_ft3"] * 0.0283168
        st.dataframe(df, use_container_width=True)
    else:
        st.info("No masks applied.")

    st.subheader("Overall summary")
    all_vals = raise_depth[~np.isnan(raise_depth)]
    overall = summarize(all_vals, px_area_ft2)
    if px_area_ft2 > 0:
        yd3 = overall["vol_ft3"]/27.0
        m3  = overall["vol_ft3"]*0.0283168
        st.markdown(
            f"Cells: {overall['n_cells']:,} | "
            f"Raise ft: min {overall['min']:.2f} | med {overall['median']:.2f} | "
            f"mean {overall['mean']:.2f} | p95 {overall['p95']:.2f} | max {overall['max']:.2f} | "
            f"Area: {overall['area_ft2']:,.0f} ft² | Volume: {overall['vol_ft3']:,.0f} ft³ "
            f"({yd3:,.0f} yd³ / {m3:,.0f} m³)"
        )
    else:
        st.markdown(
            f"Cells: {overall['n_cells']:,} | "
            f"Raise ft: min {overall['min']:.2f} | med {overall['median']:.2f} | "
            f"mean {overall['mean']:.2f} | p95 {overall['p95']:.2f} | max {overall['max']:.2f}"
        )

    # 3D plots
    st.subheader("3D views")
    mod_flip = modified[:, ::-1]
    fig = go.Figure()
    fig.add_trace(go.Surface(
        z=mod_flip, colorscale="Earth", name="Terrain",
        showscale=True, cmin=float(np.nanmin(mod_flip)), cmax=float(np.nanmax(mod_flip))
    ))
    fig.update_layout(
        title="3D Terrain (modified)",
        scene=dict(zaxis=dict(title="Elevation (ft)"),
                   camera=dict(eye=dict(x=2, y=1.5, z=1.2)),
                   aspectmode="manual", aspectratio=dict(x=1, y=1, z=0.4)),
        height=700
    )
    st.plotly_chart(fig, use_container_width=True)

    rd_flip = raise_depth[:, ::-1]
    if np.isfinite(np.nanmax(rd_flip)) and np.nanmax(rd_flip) > 0:
        fig2 = go.Figure()
        fig2.add_trace(go.Surface(
            z=np.where(np.isnan(rd_flip), 0, rd_flip),
            colorscale="Blues", showscale=True, name="Raise Depth (ft)",
            cmin=0, cmax=float(np.nanmax(rd_flip))
        ))
        fig2.update_layout(
            title="3D – Total Raise Depth",
            scene=dict(zaxis=dict(title="Feet of Raise"),
                       camera=dict(eye=dict(x=1.8, y=1.4, z=1.2)),
                       aspectmode="manual", aspectratio=dict(x=1, y=1, z=0.35)),
            height=650
        )
        st.plotly_chart(fig2, use_container_width=True)

    # Downloads (plain TIFF without georeferencing)
    st.subheader("Downloads")
    st.download_button("Modified DEM (TIFF)", to_tiff_bytes(modified), "modified_dem.tif", "image/tiff")
    st.download_button("Raise Depth (TIFF)", to_tiff_bytes(raise_depth), "raise_depth.tif", "image/tiff")

else:
    st.info("Load a DEM, add masks, then click **Run simulation**.")