import streamlit as st
import cv2
import os
import numpy as np
from datetime import datetime
import matplotlib.pyplot as plt
from services.detection_service import detect_faults_solar, detect_faults_windmill
from services.anomaly_service import track_anomalies, predict_anomaly
from models.solar_model import load_solar_model
from models.windmill_model import load_windmill_model
from config.settings import VIDEO_FOLDER

# Initialize session state for logs, metrics, and trends
if 'logs' not in st.session_state:
    st.session_state.logs = []
if 'anomaly_counts' not in st.session_state:
    st.session_state.anomaly_counts = []
if 'frame_numbers' not in st.session_state:
    st.session_state.frame_numbers = []
if 'total_detected' not in st.session_state:
    st.session_state.total_detected = 0

def main():
    st.title("Thermal Anomaly Monitoring Dashboard")
    st.markdown("**Status:** 🟢 Running")

    # Sidebar for video selection and detection type
    st.sidebar.header("Settings")
    video_files = [f for f in os.listdir(VIDEO_FOLDER) if f.endswith('.mp4')]
    if not video_files:
        st.error("No videos found in the 'data' folder. Please add .mp4 files.")
        return
    video_file = st.sidebar.selectbox("Select Video", video_files)
    detection_type = st.sidebar.selectbox("Detection Type", ["Solar Panel", "Windmill"])

    # Load the appropriate model
    model = load_solar_model() if detection_type == "Solar Panel" else load_windmill_model()

    # Layout: Two columns for video feed and metrics
    col1, col2 = st.columns([3, 1])
    
    with col1:
        st.subheader("Live Video Feed")
        video_placeholder = st.empty()
    
    with col2:
        st.subheader("Live Metrics")
        metrics_placeholder = st.empty()

    # Layout: Two columns for logs and trends
    col3, col4 = st.columns([1, 2])
    
    with col3:
        st.subheader("Live Logs")
        logs_placeholder = st.empty()
        st.subheader("Last 5 Captured Events")
        events_placeholder = st.empty()
    
    with col4:
        st.subheader("Detection Trends")
        trends_placeholder = st.empty()

    # Process video
    if video_file:
        video_path = os.path.join(VIDEO_FOLDER, video_file)
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            st.error("Error: Could not open video file.")
            return

        frame_count = 0
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            frame_count += 1
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            # Detect faults
            faults = detect_faults_solar(model, frame_rgb) if detection_type == "Solar Panel" else detect_faults_windmill(model, frame_rgb)
            num_anomalies = len(faults)

            # Draw bounding boxes and labels
            for fault in faults:
                x, y = int(fault['location'][0]), int(fault['location'][1])
                cv2.rectangle(frame_rgb, (x-30, y-30), (x+30, y+30), (255, 0, 0), 2)
                cv2.putText(frame_rgb, f"{fault['type']}", (x, y-40), 
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)

            # Update video feed with timestamp
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            video_placeholder.image(frame_rgb, channels="RGB", caption=f"{timestamp}")

            # Update logs and metrics
            log_entry = f"{timestamp} - Frame {frame_count} - Anomalies: {num_anomalies}"
            st.session_state.logs.append(log_entry)
            st.session_state.total_detected += num_anomalies
            st.session_state.anomaly_counts.append(num_anomalies)
            st.session_state.frame_numbers.append(frame_count)

            # Keep only the last 100 frames for trends to avoid memory issues
            if len(st.session_state.frame_numbers) > 100:
                st.session_state.frame_numbers.pop(0)
                st.session_state.anomaly_counts.pop(0)

            # Update live logs (show all logs, scrollable)
            with logs_placeholder.container():
                for log in st.session_state.logs[::-1]:  # Reverse to show newest first
                    st.write(log)

            # Update last 5 captured events
            with events_placeholder.container():
                for log in st.session_state.logs[-5:][::-1]:  # Last 5, newest first
                    st.write(log)

            # Update live metrics
            metrics_placeholder.write(f"""
            **anomalies:** {num_anomalies}  
            **total_detected:** {st.session_state.total_detected}
            """)

            # Predictive anomaly detection
            prediction = predict_anomaly(st.session_state.anomaly_counts)
            if prediction:
                metrics_placeholder.warning(f"**Prediction:** Potential issue detected - anomaly spike detected!")

            # Update trends graph
            fig, ax = plt.subplots()
            ax.plot(st.session_state.frame_numbers, st.session_state.anomaly_counts, marker='o')
            ax.set_title("Anomalies Over Time")
            ax.set_xlabel("Frame")
            ax.set_ylabel("Count")
            ax.grid(True)
            trends_placeholder.pyplot(fig)
            plt.close(fig)

        cap.release()
        st.success("Video processing completed.")

if __name__ == "__main__":
    main()