"""
Advanced 3D Reconstruction from Single Images with Responsible AI Features
"""

import gradio as gr
import numpy as np
import torch
from PIL import Image, ImageDraw
from transformers import GLPNForDepthEstimation, GLPNImageProcessor
import open3d as o3d
import plotly.graph_objects as go
import matplotlib.pyplot as plt
import io
import json
import time
from pathlib import Path
import tempfile
import zipfile
import hashlib
from datetime import datetime

# ============================================================================
# RESPONSIBLE AI GUIDELINES
# ============================================================================

RESPONSIBLE_AI_NOTICE = """
## ⚠️ Responsible Use Guidelines

### Privacy & Consent
- **Do not upload images containing identifiable people without their explicit consent**
- **Do not use for surveillance, tracking, or monitoring individuals**
- Facial features may be reconstructed in 3D consider privacy implications
- Remove metadata (EXIF) that may contain location or personal information

### Ethical Use
- This tool is for **educational, research, and creative purposes only**
- **Prohibited uses:**
  - Creating deepfakes or misleading 3D content
  - Unauthorized documentation of private property
  - Circumventing security systems
  - Generating 3D models for harassment or stalking
  - Commercial use without proper rights to source images

### Limitations & Bias
- Training data is concentrated in indoor Western settings, so accuracy on other
  regions and architectural styles is not well characterised
- Likely to perform less reliably on scenes far from the training distribution
- Scale is relative, not absolute, not suitable for precision measurements
- Single viewpoint limitations, occluded areas are inferred, not captured

### Data Usage
- Images are processed locally during your session
- No images are stored or transmitted to external servers
- Processing logs contain only technical metrics, no image content
- You retain all rights to your uploaded images and generated 3D models

**By using this tool, you agree to these responsible use guidelines.**
"""

# ============================================================================
# BUILT IN SAMPLE IMAGES  (real photographs, fetched once at startup)
#
# Two real, openly licensed demo photographs are downloaded at launch and
# cached to a local folder so the sample tab loads instantly on every run.
# They are chosen to demonstrate the model's strengths and limits side by side:
#
# Sample 1 · Indoor scene (two cats on a couch)
#   The well known COCO image 000000039769, used across Hugging Face vision
#   tutorials. A real 640x480 indoor photo with clear furniture geometry and
#   near/far structure, close to the NYU Depth V2 distribution GLPN was
#   trained on. Start here to see a clean reconstruction.
#
# Sample 2 · Outdoor street scene
#   A real driving/street photo from Niantic's monodepth2 repo. Open sky and a
#   receding road stress test GLPN's known weak spots. Run this, then switch to
#   DPT (High Quality) and compare the metrics, outlier % and the watertight
#   flag reveal which model handled the scene better.
#
# Sources (both from openly licensed research repositories, fetched via raw
# GitHub URLs so they work behind most corporate proxies):
#   - COCO 000000039769  -> huggingface/transformers test fixtures
#   - monodepth2 sample  -> nianticlabs/monodepth2 assets
#
# If the machine is offline, a labelled placeholder is generated instead so the
# app still launches without crashing.
# ============================================================================

import urllib.request as _urlreq

_SAMPLE_SPECS = {
    "indoor": {
        "url": "https://raw.githubusercontent.com/huggingface/transformers/main/tests/fixtures/tests_samples/COCO/000000039769.png",
        "filename": "sample_indoor_scene.png",
        "label": "indoor scene",
    },
    "outdoor": {
        "url": "https://raw.githubusercontent.com/nianticlabs/monodepth2/master/assets/test_image.jpg",
        "filename": "sample_outdoor_street.jpg",
        "label": "outdoor street",
    },
}

# Local cache folder for the downloaded samples (also used as gr.Image values).
_SAMPLE_DIR = Path(tempfile.gettempdir()) / "depthforge_samples"
_SAMPLE_DIR.mkdir(parents=True, exist_ok=True)


def _placeholder(label: str) -> Image.Image:
    """Simple labelled placeholder used only when a download fails (offline)."""
    img = Image.new("RGB", (640, 480), (235, 230, 220))
    d = ImageDraw.Draw(img)
    d.rectangle([8, 8, 631, 471], outline=(150, 140, 120), width=2)
    d.text((24, 24), f"sample unavailable offline:\n{label}", fill=(90, 80, 60))
    d.text((24, 440), "connect to the internet and restart to load the real photo",
           fill=(120, 110, 90))
    return img


def _load_sample(key: str) -> Image.Image:
    """
    Return a real sample photo as a PIL image, downloading + caching on first
    run. Falls back to a placeholder if the machine has no internet access.
    """
    spec = _SAMPLE_SPECS[key]
    dest = _SAMPLE_DIR / spec["filename"]

    # use the cached copy if we already downloaded it this session / earlier
    if dest.exists() and dest.stat().st_size > 0:
        try:
            return Image.open(dest).convert("RGB")
        except Exception:
            pass  # corrupt cache -> re-download below

    try:
        req = _urlreq.Request(spec["url"], headers={"User-Agent": "DepthForge/1.0"})
        with _urlreq.urlopen(req, timeout=15) as r:
            dest.write_bytes(r.read())
        print(f"\u2713 sample '{key}' downloaded -> {dest}")
        return Image.open(dest).convert("RGB")
    except Exception as e:
        print(f"\u26a0 could not download sample '{key}' ({e}); using placeholder")
        return _placeholder(spec["label"])


# Fetch once at startup so the sample tab is instant thereafter
SAMPLE_INDOOR  = _load_sample("indoor")
SAMPLE_OUTDOOR = _load_sample("outdoor")

# ============================================================================
# PRIVACY & SAFETY FUNCTIONS
# ============================================================================

def check_image_safety(image):
    warnings = []
    width, height = image.size
    if width * height > 10_000_000:
        warnings.append("⚠️ Very large image: consider resizing to improve processing speed")
    if max(width, height) / min(width, height) > 3:
        warnings.append("⚠️ Unusual aspect ratio detected: ensure image doesn't contain unintended content")
    try:
        exif = image.getexif()
        if exif and any(k for k in exif.keys() if k in [34853, 0x8825]):
            warnings.append("⚠️ GPS location data detected in image: consider removing EXIF data for privacy")
    except Exception:
        pass
    return True, "\n".join(warnings) if warnings else None


def generate_session_id():
    return hashlib.sha256(str(datetime.now()).encode()).hexdigest()[:16]


def content_policy_check(image):
    w, h = image.size
    if w < 100 or h < 100:
        return False, "Image too small: minimum 100×100 pixels required for meaningful reconstruction"
    return True, None

# ============================================================================
# MODEL LOADING
# ============================================================================

print("Loading GLPN model (lightweight)...")
try:
    glpn_processor = GLPNImageProcessor.from_pretrained("vinvino02/glpn-nyu")
    glpn_model     = GLPNForDepthEstimation.from_pretrained("vinvino02/glpn-nyu")
    print("✓ GLPN model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")
    glpn_processor = None
    glpn_model     = None

dpt_model     = None
dpt_processor = None

# ============================================================================
# CORE 3D RECONSTRUCTION
# ============================================================================

def process_image(image, model_choice="GLPN (Recommended)", visualization_type="mesh"):

    def _quality_assessment(metrics):
        notes = []
        pct = (metrics["outliers_removed"] / metrics["initial_points"]) * 100
        notes.append("Very clean depth estimation" if pct < 5 else ("Good depth quality" if pct < 15 else "High noise in depth estimation"))
        if metrics["is_edge_manifold"] and metrics["is_vertex_manifold"]:
            notes.append("Excellent topology")
        elif metrics["is_vertex_manifold"]:
            notes.append("Good local topology")
        else:
            notes.append("Topology issues present")
        notes.append("Watertight mesh: ready for 3D printing!" if metrics["is_watertight"] else "Not watertight —needs repair for 3D printing")
        return "\n".join(f"- {n}" for n in notes)

    if glpn_model is None:
        return None, None, None, "❌ Model failed to load. Please refresh the page.", None

    try:
        # preprocess
        new_height = 480 if image.height > 480 else image.height
        new_height -= new_height % 32
        new_width  = int(new_height * image.width / image.height)
        diff       = new_width % 32
        new_width  = new_width - diff if diff < 16 else new_width + (32 - diff)
        image      = image.resize((new_width, new_height), Image.LANCZOS)

        # model selection
        if model_choice == "GLPN (Recommended)":
            processor, model = glpn_processor, glpn_model
        else:
            global dpt_model, dpt_processor
            if dpt_model is None:
                print("Loading DPT model (first time only)…")
                from transformers import DPTForDepthEstimation, DPTImageProcessor
                dpt_processor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
                dpt_model     = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
                print("✓ DPT model loaded!")
            processor, model = dpt_processor, dpt_model

        inputs = processor(images=image, return_tensors="pt")
        t0 = time.time()
        with torch.no_grad():
            predicted_depth = model(**inputs).predicted_depth
        depth_time = time.time() - t0

        # crop / align
        pad    = 16
        output = predicted_depth.squeeze().cpu().numpy() * 1000.0
        output = output[pad:-pad, pad:-pad]
        image  = image.crop((pad, pad, image.width - pad, image.height - pad))
        dh, dw = output.shape
        iw, ih = image.size
        if dh != ih or dw != iw:
            from scipy import ndimage
            output = ndimage.zoom(output, (ih / dh, iw / dw), order=1)

        # depth visualisation
        fig, ax = plt.subplots(1, 2, figsize=(14, 7))
        ax[0].imshow(image); ax[0].set_title("Original Image", fontsize=14, fontweight="bold"); ax[0].axis("off")
        im = ax[1].imshow(output, cmap="plasma")
        ax[1].set_title("Estimated Depth Map", fontsize=14, fontweight="bold"); ax[1].axis("off")
        plt.colorbar(im, ax=ax[1], fraction=0.046, pad=0.04)
        plt.tight_layout()
        buf = io.BytesIO(); plt.savefig(buf, format="png", dpi=150, bbox_inches="tight"); buf.seek(0)
        depth_viz = Image.open(buf)
        plt.close()

        # point cloud
        w, h = image.size
        if output.shape != (h, w):
            from scipy import ndimage
            output = ndimage.zoom(output, (h / output.shape[0], w / output.shape[1]), order=1)

        depth_img = (output * 255 / np.max(output)).astype(np.uint8)
        img_arr   = np.array(image)
        rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(
            o3d.geometry.Image(img_arr), o3d.geometry.Image(depth_img), convert_rgb_to_intensity=False)
        cam = o3d.camera.PinholeCameraIntrinsic()
        cam.set_intrinsics(w, h, 500, 500, w / 2, h / 2)
        pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd, cam)
        initial_points = len(pcd.points)

        _, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
        pcd    = pcd.select_by_index(ind)
        pcd.estimate_normals()
        pcd.orient_normals_to_align_with_direction()

        # mesh
        t1   = time.time()
        mesh = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd, depth=9, n_threads=1)[0]
        tree = o3d.geometry.KDTreeFlann(pcd)
        mesh.vertex_colors = o3d.utility.Vector3dVector(
            np.array([pcd.colors[tree.search_knn_vector_3d(v, 1)[1][0]] for v in mesh.vertices]))
        mesh.rotate(mesh.get_rotation_matrix_from_xyz((np.pi, 0, 0)), center=(0, 0, 0))
        mesh_time = time.time() - t1
        mesh.compute_vertex_normals()

        # metrics
        metrics = {
            "model_used":               model_choice,
            "depth_estimation_time":    f"{depth_time:.2f}s",
            "mesh_reconstruction_time": f"{mesh_time:.2f}s",
            "total_time":               f"{depth_time + mesh_time:.2f}s",
            "initial_points":           initial_points,
            "outliers_removed":         initial_points - len(pcd.points),
            "final_points":             len(pcd.points),
            "vertices":                 len(mesh.vertices),
            "triangles":                len(mesh.triangles),
            "is_edge_manifold":         mesh.is_edge_manifold(),
            "is_vertex_manifold":       mesh.is_vertex_manifold(),
            "is_watertight":            mesh.is_watertight(),
        }
        try:
            sa = mesh.get_surface_area()
            if sa <= 0:
                verts = np.asarray(mesh.vertices); tris = np.asarray(mesh.triangles)
                cross = np.cross(verts[tris[:,1]] - verts[tris[:,0]], verts[tris[:,2]] - verts[tris[:,0]])
                sa = float(np.sum(0.5 * np.linalg.norm(cross, axis=1)))
            metrics["surface_area"] = float(sa)
        except Exception:
            metrics["surface_area"] = "Unable to compute"
        metrics["volume"] = float(mesh.get_volume()) if mesh.is_watertight() else None

        # 3-D visualisation
        pts  = np.asarray(pcd.points)
        cols = np.asarray(pcd.colors)

        if visualization_type == "point_cloud":
            plotly_fig = go.Figure(data=[go.Scatter3d(
                x=pts[:,0], y=pts[:,1], z=pts[:,2], mode="markers",
                marker=dict(size=2, color=["rgb({},{},{})".format(int(r*255),int(g*255),int(b*255)) for r,g,b in cols]))])
        else:
            verts = np.asarray(mesh.vertices); tris = np.asarray(mesh.triangles)
            vcols = np.asarray(mesh.vertex_colors)
            plotly_fig = go.Figure(data=[go.Mesh3d(
                x=verts[:,0], y=verts[:,1], z=verts[:,2],
                i=tris[:,0], j=tris[:,1], k=tris[:,2],
                vertexcolor=["rgb({},{},{})".format(int(r*255),int(g*255),int(b*255)) for r,g,b in vcols],
                opacity=0.95)])

        plotly_fig.update_layout(
            scene=dict(xaxis=dict(visible=False), yaxis=dict(visible=False), zaxis=dict(visible=False), aspectmode="data"),
            height=700)

        # export
        tmp = Path(tempfile.mkdtemp())
        pcd_path = tmp / "point_cloud.ply";  o3d.io.write_point_cloud(str(pcd_path), pcd)
        mesh_ply = tmp / "mesh.ply";          o3d.io.write_triangle_mesh(str(mesh_ply), mesh)
        mesh_obj = tmp / "mesh.obj";          o3d.io.write_triangle_mesh(str(mesh_obj), mesh)
        mesh_stl = tmp / "mesh.stl";          o3d.io.write_triangle_mesh(str(mesh_stl), mesh)
        met_path = tmp / "metrics.json";      met_path.write_text(json.dumps(metrics, indent=2, default=str))
        zip_path = tmp / "reconstruction_complete.zip"
        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
            for p in [pcd_path, mesh_ply, mesh_obj, mesh_stl, met_path]:
                zf.write(p, p.name)

        report = f"""
## Reconstruction Complete!

### Performance
- **Processing Time**: {metrics['total_time']}
- **Points**: {metrics['final_points']:,}
- **Triangles**: {metrics['triangles']:,}

### Quality
- **Topology**: {'Good' if metrics['is_vertex_manifold'] else 'Issues'}
- **Watertight**: {'Yes' if metrics['is_watertight'] else 'No'}

### Assessment
{_quality_assessment(metrics)}

**Download the complete package below!**
"""
        return depth_viz, plotly_fig, str(zip_path), report, json.dumps(metrics, indent=2, default=str)

    except Exception as e:
        import traceback
        return None, None, None, f"Error: {e}\n\n{traceback.format_exc()}", None


def process_image_with_safeguards(image, model_choice="GLPN (Recommended)", visualization_type="mesh", consent_given=False):
    session_id = generate_session_id()
    if not consent_given:
        return None, None, None, "**You must agree to the Responsible Use Guidelines first.**", None
    if image is None:
        return None, None, None, "Please upload an image first.", None
    _, safety_warning = check_image_safety(image)
    passes, policy_msg = content_policy_check(image)
    if not passes:
        return None, None, None, policy_msg, None

    depth_viz, plotly_fig, zip_path, report, json_metrics = process_image(image, model_choice, visualization_type)
    if safety_warning:
        report = f"**Privacy Notice:**\n{safety_warning}\n\n{report}"
    if json_metrics:
        metrics = json.loads(json_metrics)
        metrics["responsible_ai"] = {"session_id": session_id, "timestamp": datetime.now().isoformat(), "consent_given": True}
        json_metrics = json.dumps(metrics, indent=2)
    return depth_viz, plotly_fig, zip_path, report, json_metrics

# ============================================================================
# GRADIO INTERFACE
# ============================================================================

THEME = gr.themes.Soft(
    primary_hue=gr.themes.colors.amber,
    secondary_hue=gr.themes.colors.blue,
    neutral_hue=gr.themes.colors.stone,
    font=[gr.themes.GoogleFont("IBM Plex Sans"), "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "monospace"],
)

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Fraunces:ital,opsz,wght@0,9..144,600;0,9..144,800;0,9..144,900;1,9..144,500&display=swap');

.gradio-container h1, .gradio-container h2,
.gradio-container h3, .gradio-container h4 {
    font-family: 'Fraunces', Georgia, serif !important; letter-spacing: -.01em;
}
.gradio-container h2 { border-left: 3px solid var(--color-accent); padding-left: .55rem !important; }
.gradio-container h3 { opacity: .85; }

.df-header {
    position: relative; border-radius: 10px;
    padding: 1.8rem 1.7rem 1.4rem; margin: .3rem 0 1rem;
    border-bottom: 2px solid var(--border-color-primary);
}
.df-kicker { font-family: 'IBM Plex Mono', monospace; font-size: .7rem; letter-spacing: .04em; margin-bottom: .3rem; opacity: .6; }
.df-title  { font-family: 'Fraunces', serif; font-weight: 900; font-size: 2.5rem; line-height: 1; margin: 0 0 .3rem; }
.df-title span {
    background: linear-gradient(180deg, transparent 58%, color-mix(in srgb, var(--color-accent) 35%, transparent) 58%);
    padding: 0 .05em;
}
.df-sub   { font-family: 'IBM Plex Mono', monospace; font-size: .74rem; opacity: .6; }
.df-stamp {
    position: absolute; top: 1.2rem; right: 1.2rem; transform: rotate(7deg);
    font-family: 'IBM Plex Mono', monospace; font-size: .62rem; letter-spacing: .08em;
    text-transform: uppercase; border: 1.5px solid; border-radius: 4px; padding: .2rem .45rem; opacity: .75;
    color: var(--error-text-color, #c0392b); border-color: var(--error-text-color, #c0392b);
}
.df-note {
    margin-top: .9rem; font-family: 'IBM Plex Mono', monospace; font-size: .72rem;
    border-left: 3px solid var(--color-accent); padding: .5rem .7rem;
    background: color-mix(in srgb, var(--color-accent) 8%, transparent); border-radius: 0 4px 4px 0;
}
.df-badge {
    display: inline-block; font-family: 'IBM Plex Mono', monospace; font-size: .63rem;
    border: 1px solid var(--border-color-primary); border-radius: 3px;
    padding: .18rem .48rem; margin: .5rem .35rem 0 0; opacity: .75;
}
button.primary {
    text-transform: uppercase !important; letter-spacing: .04em !important;
    font-family: 'IBM Plex Mono', monospace !important;
    box-shadow: 3px 3px 0 color-mix(in srgb, var(--color-accent) 60%, transparent) !important;
    transition: transform .08s, box-shadow .08s !important;
}
button.primary:hover { transform: translate(1px,1px); box-shadow: 2px 2px 0 color-mix(in srgb, var(--color-accent) 60%, transparent) !important; }
.gradio-image, .gradio-plot, [class*="image-container"] {
    border-radius: 8px !important;
    border: 1px solid var(--border-color-primary) !important;
    box-shadow: 3px 3px 0 color-mix(in srgb, var(--body-text-color) 10%, transparent) !important;
}
.gradio-container :not(pre) > code { background: color-mix(in srgb, var(--color-accent) 25%, transparent) !important; border-radius: 3px; padding: .05em .3em; }
.gradio-container pre { border-radius: 8px !important; }
.gradio-container blockquote { border-left: 3px solid var(--color-accent) !important; font-family: 'Fraunces', serif; font-style: italic; }
button.selected { border-bottom: 2px solid var(--color-accent) !important; }
.df-footer {
    text-align: center; padding: 1.5rem; margin-top: 1.4rem;
    border-top: 1px solid var(--border-color-primary);
    font-family: 'IBM Plex Mono', monospace; font-size: .72rem; line-height: 1.7; opacity: .7;
}

/* sample cards */
.sample-card {
    border: 1px solid var(--border-color-primary); border-radius: 8px;
    padding: 1rem 1.1rem;
    background: color-mix(in srgb, var(--color-accent) 4%, transparent);
    margin-bottom: .75rem;
}
.sample-card h4 { font-family: 'Fraunces', Georgia, serif !important; font-size: 1rem; margin: 0 0 .4rem; }
.sample-card p  { font-family: 'IBM Plex Mono', monospace; font-size: .7rem; opacity: .75; margin: 0 0 .6rem; line-height: 1.6; }
.sample-pill {
    display: inline-block; font-family: 'IBM Plex Mono', monospace; font-size: .6rem;
    border-radius: 3px; padding: .12rem .38rem; margin-right: .3rem;
    border: 1px solid; opacity: .85;
}
.sample-pill.good { color: var(--success-text-color, #2d7a2d); border-color: var(--success-text-color, #2d7a2d); }
.sample-pill.warn { color: var(--warning-text-color, #a07000); border-color: var(--warning-text-color, #a07000); }
"""

# ============================================================================
# LAYOUT
# ============================================================================

with gr.Blocks(title="DepthForge") as demo:

    gr.HTML("""
    <div class="df-header">
        <div class="df-stamp">research demo</div>
        <div class="df-kicker">// monocular depth → 3D, a perception experiment</div>
        <div class="df-title">Depth<span>Forge</span></div>
        <div class="df-sub">3D reconstruction from a single photo · GLPN &amp; DPT</div>
        <div style="margin-top:.7rem;">
            <span class="df-badge">single image</span>
            <span class="df-badge">glpn + dpt</span>
            <span class="df-badge">mesh · point cloud</span>
            <span class="df-badge">ply / obj / stl</span>
        </div>
        <div class="df-note">⚠ please skim the Responsible Use tab before you start the consent box on the Reconstruct tab is required.</div>
    </div>
    """)

    with gr.Tabs():

        # ── Responsible Use ─────────────────────────────────────────────────
        with gr.Tab("⚠️ Responsible Use"):
            gr.Markdown(RESPONSIBLE_AI_NOTICE)
            gr.Markdown("""
### Known Limitations & Biases
- Training data is concentrated in Western indoor scenes
- Likely to underperform on scenes outside that distribution
- Scale is relative, not absolute
- Single viewpoint captures only visible surfaces
""")

        # ── Reconstruct ─────────────────────────────────────────────────────
        with gr.Tab("🔬 Reconstruct"):
            consent_checkbox = gr.Checkbox(
                label="I have read and agree to the Responsible Use Guidelines", value=False)

            with gr.Row():
                with gr.Column(scale=1):
                    input_image = gr.Image(type="pil", label="Upload a photo", sources=["upload", "clipboard"])

                    # Sample loaders  click to drop a sample into the uploader
                    # above, right here on this tab (no need to leave for the
                    # Try a Sample page). See that page for what each one shows.
                    gr.Markdown("**No image? Load a sample:**")
                    with gr.Row():
                        load_indoor_btn  = gr.Button("🛋 Indoor sample",  variant="secondary", size="sm")
                        load_outdoor_btn = gr.Button("🛣 Outdoor sample", variant="secondary", size="sm")

                    model_choice = gr.Radio(
                        choices=["GLPN (Recommended)", "DPT (High Quality)"],
                        value="GLPN (Recommended)", label="Depth model",
                        info="GLPN is fast and indoor-trained; DPT handles more varied scenes.")
                    visualization_type = gr.Radio(choices=["mesh", "point_cloud"], value="mesh", label="Show as")
                    reconstruct_btn = gr.Button("▶ Reconstruct", variant="primary", size="lg")

                with gr.Column(scale=2):
                    depth_output = gr.Image(label="Fig. 1 · depth map (the model's estimate)")
                    viewer_3d    = gr.Plot(label="Fig. 2 · interactive 3D model")

            with gr.Row():
                with gr.Column():
                    metrics_output = gr.Markdown(label="Report")
                with gr.Column():
                    json_output = gr.Textbox(label="Metrics (JSON)", lines=8)

            download_output = gr.File(label="Download package (ZIP)")

            # load a sample into the uploader on this same tab
            load_indoor_btn.click(fn=lambda: SAMPLE_INDOOR,  inputs=[], outputs=[input_image])
            load_outdoor_btn.click(fn=lambda: SAMPLE_OUTDOOR, inputs=[], outputs=[input_image])

            reconstruct_btn.click(
                fn=process_image_with_safeguards,
                inputs=[input_image, model_choice, visualization_type, consent_checkbox],
                outputs=[depth_output, viewer_3d, download_output, metrics_output, json_output])

        # ── Try a Sample ─────────────────────────────────────────────────────
        with gr.Tab("🖼️ Try a Sample"):
            gr.Markdown("""
## About the two sample images

To use a sample, head to the **🔬 Reconstruct** tab and click the
**Indoor sample** or **Outdoor sample** button under the uploader  it loads
straight in, no switching back and forth.

This page explains what each sample is for and how to read the results.
The two are a deliberate side-by-side: **Sample 1** shows the model at its
best, **Sample 2** shows where it struggles and when to switch models.
""")

            with gr.Row(equal_height=False):

                # ── Sample 1 ────────────────────────────────────────────────
                with gr.Column():
                    gr.HTML("""
<div class="sample-card">
  <h4>Sample 1 · Indoor scene</h4>
  <p>
    A real 640×480 photo (two cats resting on a couch  the well-known COCO
    image used across Hugging Face vision tutorials). Clear foreground/background
    depth layering with distinct near and far objects makes it a good fit for
    GLPN. Expect a sharp depth map, a dense point cloud, and a clean mesh.
    Start here for your "hello world" reconstruction.
  </p>
  <span class="sample-pill good">✓ real photo</span>
  <span class="sample-pill good">✓ GLPN ideal input</span>
  <span class="sample-pill good">✓ clear depth layering</span>
</div>
""")
                    gr.Image(value=SAMPLE_INDOOR, label="Sample 1 preview", interactive=False)

                # ── Sample 2 ────────────────────────────────────────────────
                with gr.Column():
                    gr.HTML("""
<div class="sample-card">
  <h4>Sample 2 · Outdoor street</h4>
  <p>
    A real outdoor street scene (from Niantic's monodepth2 demo assets). Open
    sky carries no true depth signal and the receding road blends into the
    background, so GLPN tends to produce noisier depth here  outlier % rises
    and the mesh may not be watertight. Run it, then switch to
    <strong>DPT (High Quality)</strong> and compare the two metrics reports to
    see how much the model choice matters.
  </p>
  <span class="sample-pill warn">⚡ real photo</span>
  <span class="sample-pill warn">⚡ tests model limits</span>
  <span class="sample-pill good">✓ switch to DPT here</span>
</div>
""")
                    gr.Image(value=SAMPLE_OUTDOOR, label="Sample 2 preview", interactive=False)

            gr.Markdown("""
---
### What to look for in the results

Depth map colours use the **plasma** colourmap: **bright yellow/white = far**, **dark purple/blue = near**.

| | Sample 1:  indoor | Sample 2: outdoor |
|---|---|---|
| **Depth map** | Sharp near/far contrast; bright areas = far objects | Sky often flat or patchy no true depth signal |
| **Point cloud** | Dense, well-separated depth layers | Sparse / noisy in sky region |
| **Mesh** | Clean topology, dense triangles; not watertight (single view) | More holes/noise at sky/ground boundary |
| **Recommended model** | GLPN (Recommended) | DPT (High Quality) |
| **Expected outlier %** | < 10 % | Can exceed 20 % |

> **Tip:** run Sample 2 once with GLPN, once with DPT, and compare the
> **Metrics (JSON)** output. The `outliers_removed` count and `is_watertight`
> flag are the fastest way to judge which model suited the scene better.
""")

        # ── Theory ──────────────────────────────────────────────────────────
        with gr.Tab("📖 Theory & Background"):
            gr.Markdown("""
## About This Tool

This application converts a single 2D photograph into an interactive 3D model
automatically, using monocular depth estimation  no special equipment, no
multiple cameras, no manual modelling.

## The Models

### GLPN (Global-Local Path Networks)
- **Paper:** Kim et al. (2022): [arXiv:2201.07436](https://arxiv.org/abs/2201.07436), Apache 2.0 licence
- **Training data:** NYU Depth V2: 1,449 densely labelled RGB-D frames across 464 indoor scenes, 480 × 640 px
- **Best for:** Building interiors, objects on tables, architectural corners
- **Speed:** Fast (0.3–2.5 s on CPU)
- **Hosted:** `vinvino02/glpn-nyu` on Hugging Face

### DPT (Dense Prediction Transformer)
- **Paper:** Ranftl et al. (2021),  ICCV 2021, MIT licence
- **Training data:** MIX 6 diverse indoor + outdoor, ~1.4 M images
- **Best for:** Outdoor scenes, varied lighting, inputs outside GLPN's indoor training distribution
- **Speed:** Slower (3–8 s), loads on first use
- **Hosted:** `Intel/dpt-large` on Hugging Face

## Pipeline (10 steps)

1. Image preprocessing: resize to model stride requirements
2. Depth estimation: neural network inference
3. Depth visualisation: plasma colour map
4. Point cloud generation: back-project pixels via pinhole camera model
5. Outlier removal: statistical filtering (nb_neighbors=20, std_ratio=2)
6. Normal estimation:  PCA on local neighbourhoods; oriented toward +Z axis
7. Mesh reconstruction: Poisson surface reconstruction (depth=9)
8. Quality metrics: manifold checks, surface area, volume
9. Interactive 3D visualisation: Plotly Mesh3d / Scatter3d
10. File export: PLY, OBJ, STL, JSON
""")

        # ── Usage Guide ─────────────────────────────────────────────────────
        with gr.Tab("🧭 Usage Guide"):
            gr.Markdown("""
## How to Use This Application

### First time? Use a sample
On the **🔬 Reconstruct** tab, tick the consent box, then click the
**Indoor sample** or **Outdoor sample** button under the uploader. It loads
instantly, hit ▶ Reconstruct and you'll have a result in under a minute, no
image of your own needed. (See the **🖼️ Try a Sample** tab for what each
sample demonstrates.)

### Step 1 · Read the guidelines
The **⚠️ Responsible Use** tab explains privacy rules, prohibited uses,
and model biases. The consent checkbox is required.

### Step 2 · Choose a model

| | GLPN | DPT |
|---|---|---|
| Speed | 0.3–2.5 s | 3–8 s |
| Best for | Indoor, NYU-style | Outdoor, varied scenes |
| Memory | Light (loads at startup) | Heavier (loads on first use) |
| Licence | Apache 2.0 | MIT |

### Step 3 · Reconstruct
Click **▶ Reconstruct**. Total time is usually 10–60 s depending on
image resolution and available hardware.

### Step 4 · Read the depth map
- **Yellow / white** → far from camera
- **Purple / blue**  → close to camera
- Flat uniform patches = depth-ambiguous regions (sky, mirrors, blank walls)

### Step 5 · Explore the 3D viewer
Drag to rotate · Scroll to zoom · Right-click drag to pan · Double-click to reset

### Step 6 · Download
The ZIP contains `point_cloud.ply`, `mesh.ply`, `mesh.obj`, `mesh.stl`,
and `metrics.json`.

**Free viewers:** MeshLab · Blender · CloudCompare · 3dviewer.net

---

## Troubleshooting

| Problem | Fix |
|---|---|
| No output appears | Refresh page; try a smaller image |
| Mesh has holes | Expected for single-view: use MeshLab repair |
| Wrong colours on mesh | Use point cloud mode for accurate colours |
| Very slow | Reduce image resolution; GLPN is faster than DPT |
| Not watertight | Normal for single-view reconstruction, still usable for visualisation |
""")

        # ── Ethics ──────────────────────────────────────────────────────────
        with gr.Tab("⚖️ Ethics & Impact"):
            gr.Markdown("""
## Algorithmic Bias & Fairness

These points are reasoned from the *composition* of the training data
(GLPN's NYU Depth V2 is built from indoor scenes captured in a small number of
US locations) rather than from a formal bias audit. They are framed as open
concerns worth investigating, not measured findings.

**Geographic coverage:** training data is concentrated in North America and
Europe, so performance on scenes from other regions is not well-characterised
and may be less reliable.

**Architectural style:** modern Western interiors are well-represented in the
training set; traditional, vernacular, and indigenous structures appear to be
far less common, so the model's behaviour on them is largely untested.

**Socioeconomic coverage:** because the indoor scenes come from a limited set of
environments, coverage across diverse socioeconomic settings is likely uneven,
this is an inference from the data's origin, not a quantified measurement.

## Potential Harms

- **Privacy:** unauthorised 3D reconstruction of private spaces or individuals
- **Misinformation:** fabricating 3D "evidence" or misleading spatial claims
- **Property rights:** documenting copyrighted designs without permission

## What This Tool Does to Reduce Harm

- Mandatory consent acknowledgment before any processing
- No webcam option; all processing is local to the session
- Clear prohibited-use list in the Responsible Use tab
- Bias and limitation documentation throughout the interface

## User Responsibilities

You are responsible for ensuring lawful use of source images, obtaining
necessary consents, and using outputs ethically and transparently.
""")


    gr.HTML("""
    <div class="df-footer">
      field notes &amp; code by <b>Priyadharshini Ramesh Kumar</b> · 2025–26<br>
      built out of curiosity about how we read depth from a flat image 🔎
    </div>
    """)

if __name__ == "__main__":
    import os

    # On Hugging Face Spaces a public URL is provided automatically and
    # share=True is unsupported, so only enable it for genuinely local runs.
    on_spaces = bool(os.environ.get("SPACE_ID") or os.environ.get("SYSTEM") == "spaces")

    print("=" * 60)
    print("RESPONSIBLE AI 3D RECONSTRUCTION")
    print("=" * 60)
    print("✓ GLPN default, DPT on demand")
    print("✓ Two built-in real sample images (cached at startup)")
    print("✓ No webcam option")
    print("✓ Local processing")
    print("✓ Consent required")
    print(f"✓ Environment: {'Hugging Face Spaces' if on_spaces else 'local'}")
    print("=" * 60)

    demo.launch(
        share=not on_spaces,      # share only makes sense (and works) off-Spaces
        ssr_mode=False,
        theme=THEME,
        css=CSS,
    )