import streamlit as st
import cv2
import torch
import pandas as pd
import os
from PIL import Image
import numpy as np

# --- CONFIGURATION ---
MODEL_PATH = "best.pt"  # Jo aapki file list mein hai
LOG_FILE = "scan_history.csv"

# Model load karne ka function
@st.cache_resource
def load_my_model():
    # Agar YOLOv8 hai toh ultralytics use karein, v5 hai toh torch.hub
    try:
        model = torch.hub.load('ultralytics/yolov5', 'custom', path=MODEL_PATH)
        return model
    except:
        st.error("Model load nahi ho raha. Check karein ki best.pt sahi jagah hai.")
        return None

# Data auto-save karne ka function
def auto_log_data(result_count):
    new_entry = pd.DataFrame([[pd.Timestamp.now(), result_count]], columns=["Date", "Detections"])
    if not os.path.isfile(LOG_FILE):
        new_entry.to_csv(LOG_FILE, index=False)
    else:
        new_entry.to_csv(LOG_FILE, mode='a', header=False, index=False)

# --- UI SETUP ---
st.set_page_config(page_title="Stroke-IA Detector", layout="wide")
st.title("🧠 Stroke-IA Real-time Analysis")

tab1, tab2 = st.tabs(["🔍 Detection", "📊 Analytics Dashboard"])

model = load_my_model()

with tab1:
    st.subheader("Upload for AI Scanning")
    uploaded_file = st.file_uploader("Image ya Video select karein", type=['jpg', 'jpeg', 'png', 'mp4'])

    if uploaded_file is not None and model is not None:
        # Image Analysis
        if uploaded_file.type.startswith('image'):
            img = Image.open(uploaded_file)
            results = model(img) # Model prediction
            
            # Show Result
            st.image(np.squeeze(results.render()), caption="AI Prediction")
            
            # Auto-Save
            det_count = len(results.pandas().xyxy[0])
            if st.button("Save Result to Dashboard"):
                auto_log_data(det_count)
                st.success(f"Data Saved! {det_count} signs detected.")

with tab2:
    st.subheader("📈 Automatic Analysis History")
    if os.path.exists(LOG_FILE):
        df = pd.read_csv(LOG_FILE)
        st.write("Aapko dobara CSV upload karne ki zaroorat nahi hai. Ye history hai:")
        st.dataframe(df, use_container_width=True)
        st.line_chart(df['Detections'])
    else:
        st.info("Abhi tak koi scan save nahi kiya gaya hai.")