import streamlit as st
import cv2
import numpy as np
from PIL import Image
import math

from mediapipe.tasks import python
from mediapipe.tasks.python import vision

# -----------------------------
# CONFIG
# -----------------------------
st.set_page_config(
    page_title="Classroom Dimension Estimator",
    page_icon="📏",
    layout="wide"
)

POSE_MODEL_PATH = "pose_landmarker_lite.task"

# -----------------------------
# UI STYLE
# -----------------------------
st.markdown("""
    <style>
    .main { padding: 2rem; }
    .stAlert { margin-top: 1rem; }
    </style>
""", unsafe_allow_html=True)

# -----------------------------
# UTIL FUNCTIONS
# -----------------------------
def calculate_distance(p1, p2):
    return math.sqrt((p2[0]-p1[0])**2 + (p2[1]-p1[1])**2)

# -----------------------------
# ROOM DETECTION (unchanged)
# -----------------------------
def estimate_room_dimensions(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 50, 150)

    lines = cv2.HoughLinesP(edges, 1, np.pi/180,
                            threshold=100,
                            minLineLength=100,
                            maxLineGap=10)

    if lines is None:
        return None, None, image

    max_w, max_h = 0, 0
    out = image.copy()

    for l in lines:
        x1, y1, x2, y2 = l[0]
        cv2.line(out, (x1, y1), (x2, y2), (0, 255, 0), 2)

        length = calculate_distance((x1,y1), (x2,y2))
        angle = abs(math.degrees(math.atan2(y2-y1, x2-x1)))

        if angle < 45 or angle > 135:
            max_w = max(max_w, length)
        else:
            max_h = max(max_h, length)

    return max_w * 0.01, max_h * 0.01, out

# -----------------------------
# MEDIA PIPE TASKS (NEW)
# -----------------------------
def estimate_with_person_reference(image):

    base_options = python.BaseOptions(model_asset_path=POSE_MODEL_PATH)

    options = vision.PoseLandmarkerOptions(
        base_options=base_options,
        running_mode=vision.RunningMode.IMAGE,
        num_poses=1
    )

    with vision.PoseLandmarker.create_from_options(options) as landmarker:

        rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        mp_image = vision.Image(image_format=vision.ImageFormat.SRGB, data=rgb)

        result = landmarker.detect(mp_image)

        if not result.pose_landmarks:
            return None, image

        lm = result.pose_landmarks[0]

        h, w = image.shape[:2]

        # TASK API = index based landmarks
        nose = lm[0]
        ankle = lm[27]

        head = (int(nose.x * w), int(nose.y * h))
        foot = (int(ankle.x * w), int(ankle.y * h))

        pixel_dist = calculate_distance(head, foot)

        if pixel_dist == 0:
            return None, image

        ratio = 1.7 / pixel_dist

        annotated = image.copy()
        cv2.circle(annotated, head, 6, (0,0,255), -1)
        cv2.circle(annotated, foot, 6, (255,0,0), -1)

        return ratio, annotated

# -----------------------------
# STREAMLIT APP
# -----------------------------
def main():

    st.title("📏 Classroom Dimension Estimator")

    uploaded_file = st.file_uploader(
        "Upload classroom image",
        type=["jpg","jpeg","png"]
    )

    if uploaded_file:

        image = Image.open(uploaded_file)
        image_np = np.array(image)

        tab1, tab2 = st.tabs(["Image", "Results"])

        with tab1:
            st.image(image, use_column_width=True)

        with st.spinner("Processing..."):

            ratio, pose_img = estimate_with_person_reference(image_np)

            if ratio:
                st.success("Person detected → using scale reference")

                w, h, processed = estimate_room_dimensions(image_np)

                if w and h:
                    w *= ratio
                    h *= ratio

                    with tab1:
                        st.image(pose_img, caption="Pose detection")

                    with tab2:
                        st.metric("Width", f"{w:.2f} m")
                        st.metric("Height", f"{h:.2f} m")
                        st.metric("Area", f"{w*h:.2f} m²")

            else:
                st.warning("No person detected → fallback method")

                w, h, processed = estimate_room_dimensions(image_np)

                with tab2:
                    if w and h:
                        st.metric("Width", f"{w:.2f} m")
                        st.metric("Height", f"{h:.2f} m")
                        st.metric("Area", f"{w*h:.2f} m²")

        with tab1:
            st.image(processed, caption="Room detection")

if __name__ == "__main__":
    main()