import gradio as gr  # Import Gradio for building the interactive UI
import cv2  # Import OpenCV for video processing and annotation
import os  # Import os for file handling
import numpy as np  # Import NumPy for array operations
from datetime import datetime  # Import datetime for timestamp generation
import matplotlib.pyplot as plt  # Import Matplotlib for plotting trends

# Import custom modules for fault detection, model loading, and settings
from services.detection_service import detect_faults_solar, detect_faults_windmill
from services.anomaly_service import track_faults, predict_fault
from models.solar_model import load_solar_model
from models.windmill_model import load_windmill_model
from config.settings import VIDEO_FOLDER

# Initialize global state to track faults across frames
logs = []  # List to store log entries
fault_counts = []  # List to store fault counts per frame
frame_numbers = []  # List to store frame numbers
total_detected = 0  # Counter for total faults detected

# Custom CSS to style the dashboard, mimicking the screenshot's blue borders and layout
css = """
<style>
.main-header {
    text-align: center;
    font-size: 24px;
    font-weight: bold;
    color: #333;
    margin-bottom: 10px;
}
.status {
    text-align: center;
    font-size: 16px;
    color: #333;
    margin-bottom: 20px;
}
.section-title {
    font-size: 16px;
    font-weight: bold;
    color: #333;
    text-transform: uppercase;
    margin-bottom: 10px;
}
.section-box {
    border: 1px solid #4A90E2;
    padding: 10px;
    border-radius: 5px;
    margin-bottom: 20px;
}
.log-entry {
    font-size: 14px;
    color: #333;
    margin-bottom: 5px;
}
.metrics-text {
    font-size: 14px;
    color: #333;
    margin-bottom: 5px;
}
</style>
"""

# Function to process video frames and detect faults
def process_video(video_path, detection_type):
    global logs, fault_counts, frame_numbers, total_detected
    cap = cv2.VideoCapture(video_path)  # Open the video file
    if not cap.isOpened():
        return "Error: Could not open video file.", None, None, None, None, None

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

    # Clear previous state for a new video session
    logs.clear()
    fault_counts.clear()
    frame_numbers.clear()
    total_detected = 0

    while cap.isOpened():
        ret, frame = cap.read()  # Read each frame
        if not ret:
            break

        frame_count += 1
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # Convert to RGB for display

        # Detect faults using the appropriate model
        faults = detect_faults_solar(model, frame_rgb) if detection_type == "Solar Panel" else detect_faults_windmill(model, frame_rgb)
        num_faults = len(faults)

        # Draw bounding boxes and labels for detected faults
        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)  # Draw blue box
            cv2.putText(frame_rgb, f"{fault['type']}", (x, y-40), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)  # Add fault type label

        # Update state with current frame data
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        log_entry = f"{timestamp} - Frame {frame_count} - Faults: {num_faults}"
        logs.append(log_entry)
        total_detected += num_faults
        fault_counts.append(num_faults)
        frame_numbers.append(frame_count)

        # Limit data to last 100 frames for performance
        if len(frame_numbers) > 100:
            frame_numbers.pop(0)
            fault_counts.pop(0)

        # Prepare outputs for Gradio UI
        video_output = frame_rgb
        metrics = f"faults: {num_faults}<br>total_detected: {total_detected}"
        live_logs = "<br>".join(logs[-20:])  # Display last 20 logs
        last_5_events = "<br>".join(logs[-5:]) if logs else "No events yet"
        prediction = "Potential fault escalation detected!" if predict_fault(fault_counts) else ""

        # Generate fault trends graph
        fig, ax = plt.subplots(figsize=(6, 3))
        ax.plot(frame_numbers, fault_counts, marker='o', color='blue')
        ax.set_title("Faults Over Time", fontsize=10)
        ax.set_xlabel("Frame", fontsize=8)
        ax.set_ylabel("Count", fontsize=8)
        ax.grid(True)
        ax.tick_params(axis='both', which='major', labelsize=6)
        plt.tight_layout()

        return video_output, metrics, live_logs, last_5_events, fig, prediction

# Create Gradio Blocks interface with custom CSS
with gr.Blocks(css=css) as demo:
    gr.Markdown("### THERMAL FAULT DETECTION DASHBOARD")  # Main header
    gr.Markdown("#### 🟢 RUNNING")  # Status indicator

    with gr.Row():
        with gr.Column(scale=3):
            with gr.Column():
                gr.Markdown("**LIVE VIDEO FEED**")  # Section title
                gr.Markdown('<div class="section-box">', unsafe_allow_html=True)
                video_output = gr.Image(label="", interactive=False)  # Display video feed
                gr.Markdown('</div>', unsafe_allow_html=True)
        with gr.Column(scale=1):
            with gr.Column():
                gr.Markdown("**LIVE METRICS**")  # Section title
                gr.Markdown('<div class="section-box">', unsafe_allow_html=True)
                metrics_output = gr.Markdown(label="")  # Display metrics
                prediction_output = gr.Markdown(label="")  # Display prediction warning
                gr.Markdown('</div>', unsafe_allow_html=True)

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Column():
                gr.Markdown("**LIVE LOGS**")  # Section title
                gr.Markdown('<div class="section-box">', unsafe_allow_html=True)
                logs_output = gr.Markdown(label="")  # Display live logs
                gr.Markdown('</div>', unsafe_allow_html=True)
            with gr.Column():
                gr.Markdown("**LAST 5 CAPTURED EVENTS**")  # Section title
                gr.Markdown('<div class="section-box">', unsafe_allow_html=True)
                events_output = gr.Markdown(label="")  # Display last 5 events
                gr.Markdown('</div>', unsafe_allow_html=True)
        with gr.Column(scale=2):
            with gr.Column():
                gr.Markdown("**DETECTION TRENDS**")  # Section title
                gr.Markdown('<div class="section-box">', unsafe_allow_html=True)
                gr.Markdown("**Faults Over Time**")  # Sub-title
                trends_output = gr.Plot(label="")  # Display fault trends graph
                gr.Markdown('</div>', unsafe_allow_html=True)

    # Sidebar for user inputs
    with gr.Row():
        with gr.Column():
            video_files = [f for f in os.listdir(VIDEO_FOLDER) if f.endswith('.mp4')]  # Get video files
            video_input = gr.Dropdown(choices=video_files, label="Select Video")  # Video selection
            detection_type = gr.Dropdown(choices=["Solar Panel", "Windmill"], label="Detection Type")  # Detection type
            submit_btn = gr.Button("Start Processing")  # Trigger button

    # Connect inputs to outputs with event trigger
    submit_btn.click(
        fn=process_video,
        inputs=[video_input, detection_type],
        outputs=[video_output, metrics_output, logs_output, events_output, trends_output, prediction_output],
        _js="() => [document.querySelector('input[type=\"file\"]').value, document.querySelector('select[name=\"detection_type\"]').value]"
    )

# Launch the Gradio app
demo.launch()