import streamlit as st
import cv2
import numpy as np
import plotly.express as px

st.set_page_config(layout="wide")
st.title("2D to 3D Reconstruction Demo")

# SIFT detector
sift = cv2.SIFT_create(nfeatures=5000)

def match_and_reconstruct(img1, img2):

    gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

    kp1, des1 = sift.detectAndCompute(gray1, None)
    kp2, des2 = sift.detectAndCompute(gray2, None)

    if des1 is None or des2 is None:
        return None

    # Brute Force Matcher
    bf = cv2.BFMatcher(cv2.NORM_L2)
    matches = bf.knnMatch(des1, des2, k=2)

    # Lowe Ratio Test
    good = []
    for m, n in matches:
        if m.distance < 0.85 * n.distance:
            good.append(m)

    if len(good) < 12:
        return None

    pts1 = np.float32([kp1[m.queryIdx].pt for m in good])
    pts2 = np.float32([kp2[m.trainIdx].pt for m in good])

    # ---- Approximate Camera Intrinsics ----
    h, w = gray1.shape
    focal = 0.8 * w

    K = np.array([
        [focal, 0, w / 2],
        [0, focal, h / 2],
        [0, 0, 1]
    ])

    # Essential Matrix
    E, mask = cv2.findEssentialMat(
        pts1,
        pts2,
        K,
        method=cv2.RANSAC,
        prob=0.999,
        threshold=3.0
    )

    if E is None:
        return None

    # Recover Pose
    _, R, t, mask_pose = cv2.recoverPose(E, pts1, pts2, K)

    # Projection matrices
    P1 = K @ np.hstack((np.eye(3), np.zeros((3, 1))))
    P2 = K @ np.hstack((R, t))

    # Triangulation
    pts1_norm = pts1.T
    pts2_norm = pts2.T

    points_4d = cv2.triangulatePoints(P1, P2, pts1_norm, pts2_norm)
    points_3d = points_4d[:3] / points_4d[3]

    return points_3d.T


# Upload images
uploaded1 = st.file_uploader("Upload First Image", type=["jpg", "png"])
uploaded2 = st.file_uploader("Upload Second Image", type=["jpg", "png"])

if uploaded1 and uploaded2:

    file_bytes1 = np.asarray(bytearray(uploaded1.read()), dtype=np.uint8)
    file_bytes2 = np.asarray(bytearray(uploaded2.read()), dtype=np.uint8)

    img1 = cv2.imdecode(file_bytes1, 1)
    img2 = cv2.imdecode(file_bytes2, 1)

    st.image(img1, caption="Image 1", use_column_width=True)
    st.image(img2, caption="Image 2", use_column_width=True)

    with st.spinner("Reconstructing 3D points..."):
        points_3d = match_and_reconstruct(img1, img2)

    if points_3d is not None:

        fig = px.scatter_3d(
            x=points_3d[:, 0],
            y=points_3d[:, 1],
            z=points_3d[:, 2],
        )

        fig.update_layout(
            margin=dict(l=0, r=0, b=0, t=0)
        )

        st.plotly_chart(fig, use_container_width=True)

    else:
        st.error("Not enough good matches found. Try images with more texture and slight viewpoint difference.")