import streamlit as st
import cv2
import numpy as np
import tensorflow as tf
from PIL import Image
import time

st.set_page_config(page_title="Lane Segmentation Dashboard", layout="wide")

# MODEL PATH
MODEL_PATH = "/app/model/1"


# LOAD MODEL
@st.cache_resource
def load_model():
    model_layer = tf.keras.layers.TFSMLayer(
        MODEL_PATH,
        call_endpoint="serving_default"
    )

    inputs = tf.keras.Input(shape=(256, 832, 3))
    outputs = model_layer(inputs)
    model = tf.keras.Model(inputs=inputs, outputs=outputs)

    return model

# PROCESSING
def preprocess_image(image):
    image = cv2.resize(image, (832, 256))
    return np.expand_dims(image.astype(np.float32) / 255.0, axis=0)


def process_image(image, model):
    processed = preprocess_image(image)

    prediction = model(processed, training=False)

    if isinstance(prediction, dict):
        prediction = list(prediction.values())[0]

    prediction = prediction.numpy()

    mask = (prediction[0] > 0.5).astype(np.uint8).squeeze()
    mask = cv2.resize(mask, (image.shape[1], image.shape[0]))

    return mask


# SIDEBAR CONTROLS
st.sidebar.header("Controls")

alpha = st.sidebar.slider("Overlay Intensity", 0.0, 1.0, 0.3)
threshold = st.sidebar.slider("Threshold", 0.1, 0.9, 0.5)
show_mask = st.sidebar.checkbox("Show Raw Mask")
show_metrics = st.sidebar.checkbox("Show Metrics", value=True)

# UI HEADER
st.title("Lane Segmentation Dashboard")

model = load_model()

st.success("Model loaded successfully")


# FILE UPLOAD
uploaded_file = st.file_uploader(
    "Upload image",
    type=["jpg", "jpeg", "png"],
    accept_multiple_files=False
)

if uploaded_file is None:
    st.info("Upload an image to begin")
    st.stop()


# LOAD IMAGE
image = np.array(Image.open(uploaded_file).convert("RGB"))

if image.shape[0] > 1500 or image.shape[1] > 1500:
    image = cv2.resize(image, (832, 256))


# RUN BUTTON
run = st.button("Run Lane Detection")

if not run:
    st.image(image, caption="Uploaded Image", width="stretch")
    st.stop()


# INFERENCE
start = time.time()

mask = process_image(image, model)

end = time.time()
inference_time = (end - start) * 1000  # ms


# APPLY THRESHOLD + COLOR
mask = (mask > threshold).astype(np.uint8)

colored_mask = np.zeros_like(image)
colored_mask[mask == 1] = [0, 255, 0]

result = cv2.addWeighted(image, 1 - alpha, colored_mask, alpha, 0)


# METRICS
lane_pixels = np.sum(mask)
total_pixels = mask.size
lane_coverage = (lane_pixels / total_pixels) * 100
fps = 1000 / inference_time if inference_time > 0 else 0


# DISPLAY
col1, col2 = st.columns(2)

with col1:
    st.subheader("Original")
    st.image(image, width="stretch")

with col2:
    st.subheader("Segmentation")
    st.image(result, width="stretch")


# OPTIONAL MASK
if show_mask:
    st.subheader("Raw Mask")
    st.image(mask * 255)


# METRICS DASHBOARD
if show_metrics:
    st.subheader("Performance Metrics")

    m1, m2, m3 = st.columns(3)

    m1.metric("Inference Time (ms)", f"{inference_time:.2f}")
    m2.metric("FPS", f"{fps:.2f}")
    m3.metric("Lane Coverage (%)", f"{lane_coverage:.2f}")