import streamlit as st
import cv2
import open3d as o3d
import numpy as np
import tempfile
import os

# Title of the App
st.title("3D Reconstruction Tool from Video 📹 → 🛠️ → 🧊")

# Sidebar: Information
st.sidebar.write("""
## About the App
Upload a video file, extract frames, reconstruct a 3D point cloud using Structure from Motion (SfM), and visualize or download the 3D mesh.
""")

# Step 1: Upload Video File
uploaded_file = st.file_uploader("Upload a Video File (MP4, AVI)", type=["mp4", "avi"])

# Function to extract frames from video
def extract_frames(video_path, frame_rate=10):
    cap = cv2.VideoCapture(video_path)
    frames = []
    count = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        if count % frame_rate == 0:
            frames.append(frame)
        count += 1
    cap.release()
    return frames

# Function to save frames as images
def save_frames_as_images(frames, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    for i, frame in enumerate(frames):
        filename = os.path.join(output_dir, f"frame_{i:04d}.png")
        cv2.imwrite(filename, frame)

# Step 2: Process Uploaded Video
if uploaded_file:
    st.video(uploaded_file)
    st.write("Extracting frames...")

    # Save the uploaded video temporarily
    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_video:
        tmp_video.write(uploaded_file.read())
        video_path = tmp_video.name

    # Extract frames
    frames = extract_frames(video_path, frame_rate=10)
    st.write(f"✅ Extracted {len(frames)} frames from the video.")

    # Save extracted frames
    frames_dir = tempfile.mkdtemp()
    save_frames_as_images(frames, frames_dir)
    st.write(f"Frames saved temporarily at `{frames_dir}`.")

    # Step 3: Structure from Motion (3D Reconstruction)
    st.write("🔄 Reconstructing 3D point cloud using Structure from Motion...")
    
    # Create Open3D Point Cloud
    pcd = o3d.geometry.PointCloud()
    for image_file in sorted(os.listdir(frames_dir)):
        img_path = os.path.join(frames_dir, image_file)
        frame = cv2.imread(img_path)
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        
        # Dummy point cloud generation for simplicity
        height, width = gray.shape
        x, y = np.meshgrid(np.arange(width), np.arange(height))
        z = gray / 255.0  # Use gray intensity as a pseudo depth
        points = np.stack((x.flatten(), y.flatten(), z.flatten()), axis=1)
        pcd.points.extend(o3d.utility.Vector3dVector(points))

    # Step 4: Surface Reconstruction
    st.write("🛠️ Generating mesh using Poisson Reconstruction...")
    pcd.estimate_normals()
    mesh, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd, depth=8)

    # Step 5: Visualization
    st.write("✅ Reconstruction Complete! Visualizing the 3D model:")
    o3d.io.write_triangle_mesh("reconstructed_mesh.stl", mesh)

    # Use Plotly for visualization
    import plotly.graph_objects as go
    vertices = np.asarray(mesh.vertices)
    triangles = np.asarray(mesh.triangles)
    fig = go.Figure(data=[go.Mesh3d(
        x=vertices[:, 0],
        y=vertices[:, 1],
        z=vertices[:, 2],
        i=triangles[:, 0],
        j=triangles[:, 1],
        k=triangles[:, 2],
        color='lightblue',
        opacity=0.50
    )])
    st.plotly_chart(fig)

    # Step 6: Download the Optimized Mesh
    st.write("📥 Download the reconstructed 3D model:")
    with open("reconstructed_mesh.stl", "rb") as f:
        st.download_button("Download 3D Mesh (STL)", f, file_name="reconstructed_3D_Model.stl")