import streamlit as st
import os
import cv2
import numpy as np
from PIL import Image
import image_processing
import matplotlib.pyplot as plt

st.set_page_config(page_title="Soil Image Processor", layout="wide")

# Convert PIL image to numpy
def np_image(pil_img):
    return np.array(pil_img.convert("RGB"))

# UI layout
st.title("🧪 Soil Image Processor")

st.subheader("📊 Image Analysis & Prediction")
uploaded_file = st.file_uploader("Upload an image for prediction", type=["png", "jpg", "jpeg"])

if uploaded_file:
    pil_img = Image.open(uploaded_file)
    img_array = np_image(pil_img)

    # Step 1: Show original image
    st.image(pil_img, caption="Uploaded Image", use_container_width=True)

    # Step 2: RGB Histogram
    hist_fig = image_processing.plot_rgb_histogram(img_array)
    st.pyplot(hist_fig)

    # Step 3: Preprocessing
    rgb_img, clahe_img, sharp_img = image_processing.preprocessing(cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB))
    st.image([rgb_img, clahe_img, sharp_img], caption=["RGB Image", "CLAHE Image", "Sharpened Image"], width=250)

    # Step 4: Choose mode: Whole image vs Region-based
    st.markdown("### 🔍 Classification Mode")
    use_segmentation = st.checkbox("Enable multi-region (segmented) classification")

    if use_segmentation:
        k = st.slider("Select number of regions (clusters)", min_value=2, max_value=5)
        segmented_image, region_predictions = image_processing.segment_and_classify_regions(img_array, k_clusters=k)

        st.image(cv2.cvtColor(segmented_image, cv2.COLOR_BGR2RGB), caption="Segmented & Classified", use_column_width=True)

        st.markdown("### 🧩 Region-wise Predictions")
        for idx, region in enumerate(region_predictions):
            st.write(f"**Region {idx + 1}:** `{region['class']}` with confidence `{region['confidence']:.2f}`")
            x, y, w, h = region['bbox']
            cropped = img_array[y:y+h, x:x+w]
            st.image(cropped, caption=f"Region {idx + 1}", width=200)

    else:
        # Step 5: Feature Extraction & Prediction on whole image
        features_df, predicted_class, confidence = image_processing.predict_image_class_with_features(cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB))

        st.subheader("📈 Extracted Features")
        st.dataframe(features_df)

        st.success(f"🔮 Predicted Class: **{predicted_class}** (Confidence: {confidence:.2f})")