# fixed_animated_pointcloud.py
import gradio as gr
import cv2
import numpy as np
import plotly.graph_objects as go
import plotly.io as pio
import json
import os
from datetime import datetime

def create_animated_point_cloud(
    video,
    resolution: int = 256,
    density: float = 0.25,
    depth: float = 0.5,
    point_size: float = 3.0,
    max_frames: int = 10,
    frame_step: int = 5,
    progress=gr.Progress()
):
    if video is None:
        return None, "Upload a short video first", None, None

    cap = cv2.VideoCapture(video)
    if not cap.isOpened():
        return None, "Cannot open video", None, None

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    max_possible = min(max_frames, (total_frames // frame_step) + 1)

    progress(0, desc="Reading video...")
    all_points = []
    processed = 0

    for i in range(0, total_frames, frame_step):
        if len(all_points) >= max_frames:
            break
        cap.set(cv2.CAP_PROP_POS_FRAMES, i)
        ret, frame = cap.read()
        if not ret:
            break

        small = cv2.resize(frame, (resolution, resolution))
        gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0
        rgb = cv2.cvtColor(small, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0

        mask = np.random.rand(*gray.shape) < density
        ys, xs = np.nonzero(mask)
        zs = (gray[mask] - 0.5) * depth * 8

        xs_norm = (xs / resolution - 0.5) * 12
        ys_norm = (0.5 - ys / resolution) * 12
        colors = rgb[mask].tolist()

        all_points.append({
            'x': xs_norm.tolist(),
            'y': ys_norm.tolist(),
            'z': zs.tolist(),
            'color': colors
        })

        processed += 1
        progress(processed / max_possible, desc=f"Extracted {processed}/{max_possible} frames...")

    cap.release()

    if not all_points:
        return None, "No frames processed", None, None

    # ====================== BUILD FIGURE ======================
    initial = all_points[0]
    initial_colors = [f"rgb({int(255*r)},{int(255*g)},{int(255*b)})" for r,g,b in initial['color']]

    trace = go.Scatter3d(
        x=initial['x'], y=initial['y'], z=initial['z'],
        mode='markers',
        marker=dict(
            size=point_size,
            color=initial_colors,
            opacity=0.85
        )
    )

    fig = go.Figure(data=[trace])

    fig.update_layout(
        scene=dict(
            aspectmode='cube',
            xaxis_title='X',
            yaxis_title='Y',
            zaxis_title='Depth (brightness)',
            camera=dict(eye=dict(x=0, y=0, z=2.8), up=dict(x=0, y=1, z=0), center=dict(x=0, y=0, z=0))
        ),
        title=f"Animated Point Cloud — {len(all_points)} frames (camera-persistent)",
        height=650,
        margin=dict(l=0, r=0, b=0, t=90),
    )

    # Data for JS: convert colors into Plotly color strings
    points_data = []
    for pts in all_points:
        rgb_colors = [f"rgb({int(255*r)},{int(255*g)},{int(255*b)})" for r,g,b in pts['color']]
        points_data.append({'x': pts['x'], 'y': pts['y'], 'z': pts['z'], 'color': rgb_colors})

    # Unique filename
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    html_path = f"animated_pointcloud_{timestamp}.html"

    # Create the HTML using Plotly CDN
    html_str = pio.to_html(
        fig,
        include_plotlyjs='cdn',
        full_html=True,
        default_width="100%",
        default_height="650px"
    )

    # Robust loop + overlay script (now also controls point size)
    overlay_id = f"pc_ctrl_{timestamp}"
    loop_script = f"""
<script>
(function() {{
    const pointsData = {json.dumps(points_data)};
    let loopInterval = null;
    let gd = null;
    let idx = 0;
    let fps = 12;
    const ctrlId = "{overlay_id}";

    function createOverlay() {{
        if (document.getElementById(ctrlId)) return;
        const wrap = document.createElement('div');
        wrap.id = ctrlId;
        wrap.style = "position:fixed; left:12px; top:12px; z-index:9999; background:rgba(0,0,0,0.55); padding:8px 10px; border-radius:8px; color:#fff; font-family:Arial,monospace; display:flex; gap:8px; align-items:center; flex-wrap:wrap;";
        wrap.innerHTML = `
            <button id="{overlay_id}_play" style="padding:6px 8px;border-radius:6px;background:#16a085;color:white;border:none;cursor:pointer;">▶ Play</button>
            <button id="{overlay_id}_pause" style="padding:6px 8px;border-radius:6px;background:#7f8c8d;color:white;border:none;cursor:pointer;">⏸ Pause</button>
            <label style="font-size:12px;margin-left:6px;">FPS</label>
            <input id="{overlay_id}_fps" type="range" min="1" max="60" value="12" style="vertical-align:middle;">
            <span id="{overlay_id}_fps_txt" style="min-width:30px; text-align:center; display:inline-block;">12</span>
            <label style="font-size:12px;margin-left:8px;">Size</label>
            <input id="{overlay_id}_size" type="range" min="1" max="12" step="0.5" value="{point_size}" style="vertical-align:middle;">
            <span id="{overlay_id}_size_txt" style="min-width:30px; text-align:center; display:inline-block;">{point_size}</span>
            <button id="{overlay_id}_resetcam" title="Reset camera" style="padding:6px 8px;border-radius:6px;background:#2d6cdf;color:white;border:none;cursor:pointer;">↺ Cam</button>
        `;
        document.body.appendChild(wrap);
    }}

    function findPlotlyDiv() {{
        const direct = document.querySelector('.js-plotly-plot');
        if (direct) return direct;
        const anyDiv = document.querySelector('[data-plotly]') || document.querySelector('.plotly-graph-div');
        if (anyDiv) return anyDiv;
        return null;
    }}

    function waitForGd(cb) {{
        const existing = findPlotlyDiv();
        if (existing) return cb(existing);

        const mo = new MutationObserver((mut, obs) => {{
            const d = findPlotlyDiv();
            if (d) {{
                obs.disconnect();
                cb(d);
            }}
        }});
        mo.observe(document.body, {{ childList:true, subtree:true }});
        setTimeout(() => {{
            const d = findPlotlyDiv();
            if (!d) {{
                try {{ mo.disconnect(); }} catch(e){{}}
                cb(null);
            }}
        }}, 8000);
    }}

    function updateFrame(i, customSize = null) {{
        if (!gd) return;
        const p = pointsData[i];
        const sizeToUse = customSize !== null ? customSize : {point_size};
        try {{
            Plotly.restyle(gd, {{
                x: [p.x],
                y: [p.y],
                z: [p.z],
                'marker.color': [p.color],
                'marker.size': [sizeToUse]
            }}, [0]);
        }} catch (e) {{
            console.error('restyle failed', e);
        }}
    }}

    function startLoop() {{
        if (loopInterval) return;
        const fpsSlider = document.getElementById('{overlay_id}_fps');
        const sizeSlider = document.getElementById('{overlay_id}_size');
        const val = parseInt(fpsSlider?.value || 12);
        const delay = Math.max(1, Math.round(1000 / val));
        idx = (idx + 1) % pointsData.length;
        updateFrame(idx, parseFloat(sizeSlider?.value || {point_size}));
        loopInterval = setInterval(() => {{
            idx = (idx + 1) % pointsData.length;
            updateFrame(idx, parseFloat(sizeSlider?.value || {point_size}));
        }}, delay);
    }}

    function pauseLoop() {{
        if (loopInterval) {{
            clearInterval(loopInterval);
            loopInterval = null;
        }}
    }}

    function resetCamera() {{
        if (!gd) return;
        try {{
            const layout = gd.layout || {{}};
            if (layout.scene && layout.scene.camera && layout.scene.camera.eye) {{
                Plotly.relayout(gd, {{ 'scene.camera.eye': layout.scene.camera.eye }});
            }} else {{
                Plotly.relayout(gd, {{ 'scene.camera.eye': {{x:0,y:0,z:2.8}} }});
            }}
        }} catch(e) {{
            console.warn('resetCamera failed', e);
        }}
    }}

    function bindOverlay() {{
        createOverlay();
        const playBtn   = document.getElementById('{overlay_id}_play');
        const pauseBtn  = document.getElementById('{overlay_id}_pause');
        const fpsSlider = document.getElementById('{overlay_id}_fps');
        const fpsTxt    = document.getElementById('{overlay_id}_fps_txt');
        const sizeSlider= document.getElementById('{overlay_id}_size');
        const sizeTxt   = document.getElementById('{overlay_id}_size_txt');
        const resetBtn  = document.getElementById('{overlay_id}_resetcam');

        playBtn.onclick  = () => startLoop();
        pauseBtn.onclick = () => pauseLoop();
        resetBtn.onclick = () => resetCamera();

        fpsSlider.oninput = (e) => {{
            const v = e.target.value;
            fpsTxt.innerText = v;
            if (loopInterval) {{
                pauseLoop();
                startLoop();
            }}
        }};

        sizeSlider.oninput = (e) => {{
            const v = parseFloat(e.target.value);
            sizeTxt.innerText = v.toFixed(1);
            // live update even when paused
            updateFrame(idx, v);
        }};
    }}

    function init() {{
        bindOverlay();
        waitForGd(function(div) {{
            if (!div) {{
                console.warn('[pc] Plotly element not found.');
                return;
            }}
            gd = div;
            updateFrame(0);

            try {{
                const ro = new ResizeObserver(() => {{
                    if (gd && gd.layout) Plotly.Plots.resize(gd);
                }});
                ro.observe(gd);
            }} catch(e) {{}}

            window.addEventListener('keydown', (ev) => {{
                if (ev.code === 'Space') {{
                    ev.preventDefault();
                    if (loopInterval) pauseLoop(); else startLoop();
                }}
            }});
        }});
    }}

    window.addEventListener('load', init);
    window.addEventListener('beforeunload', function() {{ pauseLoop(); }});
}})();
</script>
"""

    html_str = html_str.replace('</body>', loop_script + '</body>')

    with open(html_path, 'w', encoding='utf-8') as f:
        f.write(html_str)

    # Iframe src logic (Hugging Face Spaces vs local)
    if os.environ.get("SPACE_ID"):
        space_url = f"https://{os.environ['SPACE_ID']}.hf.space"
        iframe_src = f"{space_url}/file={html_path}"
    else:
        iframe_src = f"/file={html_path}"

    iframe_html = f'''
    <iframe 
        src="{iframe_src}"
        width="100%" 
        height="720" 
        style="border:none; border-radius:8px; box-shadow:0 4px 20px rgba(0,0,0,0.1);"
        allowfullscreen
        sandbox="allow-scripts allow-same-origin allow-popups">
    </iframe>
    '''

    # Static preview (first frame) with chosen point size
    preview_fig = go.Figure(data=[trace])
    preview_fig.update_layout(scene=fig.layout.scene, height=500)

    status = f"✅ Done! {len(all_points)} frames • Point size: {point_size} • Download: {html_path}"

    return preview_fig, status, iframe_html, html_path


# ====================== GRADIO INTERFACE ======================
demo = gr.Interface(
    fn=create_animated_point_cloud,
    inputs=[
        gr.Video(label="Upload Video (short = faster)"),
        gr.Slider(128, 1024, value=256, step=64, label="Resolution"),
        gr.Slider(0.05, 0.5, value=0.25, step=0.05, label="Point Density"),
        gr.Slider(0.2, 1.5, value=0.5, step=0.1, label="Depth Intensity"),
        gr.Slider(0.2, 12, value=3.0, step=0.2, label="Point Size"),
        gr.Slider(4, 65, value=10, step=1, label="Max Frames"),
        gr.Slider(2, 10, value=5, step=1, label="Frame Step (higher = faster)")
    ],
    outputs=[
        gr.Plot(label="Static Preview (first frame)"),
        gr.Textbox(label="Status"),
        gr.HTML(label="🎥 Live Interactive 3D Animation (Top-down + Looping)"),
        gr.File(label="↓ Download HTML (offline use)")
    ],
    title="Video → Interactive Animated 3D Point Cloud (camera-persistent)",
    description="Progress bar shows processing • Full 3D animation appears directly on the page • Live point size & FPS control in animation",
    flagging_mode="never"
)

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