"""
Real-time performance graphs for inference latency analysis
Shows when latency spikes occur with different resolutions and devices
"""

from PySide6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QLabel, 
    QGroupBox, QTabWidget, QFrame, QSplitter
)
from PySide6.QtCore import Qt, QTimer, Signal, Slot
from PySide6.QtGui import QPainter, QPen, QBrush, QColor, QFont
import numpy as np
from collections import deque
from typing import Dict, List, Any

class RealTimeGraph(QWidget):
    """Custom widget for drawing real-time graphs"""
    
    def __init__(self, title: str = "Graph", y_label: str = "Value", max_points: int = 300):
        super().__init__()
        self.title = title
        self.y_label = y_label
        self.max_points = max_points
        
        # Data storage
        self.x_data = deque(maxlen=max_points)
        self.y_data = deque(maxlen=max_points)
        self.spike_markers = deque(maxlen=max_points)  # Mark spikes
        self.device_markers = deque(maxlen=max_points)  # Mark device changes
        self.resolution_markers = deque(maxlen=max_points)  # Mark resolution changes
        
        # Graph settings
        self.margin = 40
        self.grid_color = QColor(60, 60, 60)
        self.line_color = QColor(0, 255, 255)  # Cyan
        self.spike_color = QColor(255, 0, 0)   # Red for spikes
        self.cpu_color = QColor(100, 150, 255) # Blue for CPU
        self.gpu_color = QColor(255, 150, 100) # Orange for GPU
        
        # Auto-scaling
        self.y_min = 0
        self.y_max = 100
        self.auto_scale = True
        
        self.setMinimumSize(400, 200)
        
    def add_data_point(self, x: float, y: float, is_spike: bool = False, device: str = "CPU", is_res_change: bool = False):
        """Add a new data point to the graph"""
        self.x_data.append(x)
        self.y_data.append(y)
        self.spike_markers.append(is_spike)
        self.device_markers.append(device)
        self.resolution_markers.append(is_res_change)
        
        # Auto-scale Y axis
        if self.auto_scale and self.y_data:
            data_max = max(self.y_data)
            data_min = min(self.y_data)
            padding = (data_max - data_min) * 0.1
            self.y_max = data_max + padding if data_max > 0 else 100
            self.y_min = max(0, data_min - padding)
        self.update()
        
    def clear_data(self):
        """Clear the graph data"""
        self.x_data.clear()
        self.y_data.clear()
        self.spike_markers.clear()
        self.device_markers.clear()
        self.resolution_markers.clear()
        self.update()
        
    def paintEvent(self, event):
        """Override paint event to draw the graph"""
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        width = self.width()
        height = self.height()
        graph_width = width - 2 * self.margin
        graph_height = height - 2 * self.margin
        
        # Background
        painter.fillRect(self.rect(), QColor(30, 30, 30))
        
        # Title
        painter.setPen(QColor(255, 255, 255))
        painter.setFont(QFont("Arial", 12, QFont.Bold))
        painter.drawText(10, 20, self.title)
        
        # Axes
        painter.setPen(QPen(QColor(200, 200, 200), 2))
        painter.drawLine(self.margin, self.margin, self.margin, height - self.margin)
        painter.drawLine(self.margin, height - self.margin, width - self.margin, height - self.margin)
        
        # Grid
        painter.setPen(QPen(self.grid_color, 1))
        for i in range(5):
            y = self.margin + (graph_height * i / 4)
            painter.drawLine(self.margin, y, width - self.margin, y)
        for i in range(10):
            x = self.margin + (graph_width * i / 9)
            painter.drawLine(x, self.margin, x, height - self.margin)
        
        # Y-axis labels
        painter.setPen(QColor(200, 200, 200))
        painter.setFont(QFont("Arial", 8))
        for i in range(5):
            y_val = self.y_min + (self.y_max - self.y_min) * (4 - i) / 4
            y_pos = self.margin + (graph_height * i / 4)
            painter.drawText(5, y_pos + 5, f"{y_val:.1f}")
        
        # X-axis label
        painter.save()
        painter.translate(15, height // 2)
        painter.rotate(-90)
        painter.drawText(-len(self.y_label) * 3, 0, self.y_label)
        painter.restore()
        
        # Data points
        if len(self.x_data) >= 2 and len(self.y_data) >= 2:
            points = []
            spike_points = []
            device_changes = []
            res_changes = []
            x_min = min(self.x_data) if self.x_data else 0
            x_max = max(self.x_data) if self.x_data else 1
            x_range = x_max - x_min if x_max > x_min else 1
            for i, (x_val, y_val, is_spike, device, is_res_change) in enumerate(zip(
                self.x_data, self.y_data, self.spike_markers, self.device_markers, self.resolution_markers
            )):
                x_screen = self.margin + (x_val - x_min) / x_range * graph_width
                y_screen = height - self.margin - (y_val - self.y_min) / (self.y_max - self.y_min) * graph_height
                points.append((x_screen, y_screen))
                if is_spike:
                    spike_points.append((x_screen, y_screen))
                if i > 0 and device != list(self.device_markers)[i-1]:
                    device_changes.append((x_screen, y_screen, device))
                if is_res_change:
                    res_changes.append((x_screen, y_screen))
            if len(points) >= 2:
                painter.setPen(QPen(self.line_color, 2))
                for i in range(len(points) - 1):
                    x1, y1 = points[i]
                    x2, y2 = points[i + 1]
                    painter.drawLine(x1, y1, x2, y2)
            painter.setPen(QPen(self.spike_color, 3))
            painter.setBrush(QBrush(self.spike_color))
            for x, y in spike_points:
                painter.drawEllipse(x - 3, y - 3, 6, 6)
            for x, y, device in device_changes:
                color = self.gpu_color if device == "GPU" else self.cpu_color
                painter.setPen(QPen(color, 2))
                painter.setBrush(QBrush(color))
                painter.drawRect(x - 2, self.margin, 4, graph_height)
            for x, y in res_changes:
                painter.setPen(QPen(QColor(255, 167, 38), 2))  # Orange for resolution change
                painter.drawLine(x, self.margin, x, height - self.margin)

class PerformanceGraphsWidget(QWidget):
    def __init__(self):
        super().__init__()
        self.setup_ui()
        self.update_timer = QTimer()
        self.update_timer.timeout.connect(self.update_graphs)
        try:
            self.update_timer.start(1000)
        except Exception as e:
            print(f"❌ Error starting performance graph timer: {e}")
        self.start_time = None
        self.latest_data = {}
        self.cpu_usage_history = deque(maxlen=300)
        self.ram_usage_history = deque(maxlen=300)
    def setup_ui(self):
        layout = QVBoxLayout(self)
        title_label = QLabel("🔥 Real-Time Inference Performance & Latency Spike Analysis")
        title_label.setStyleSheet("font-size: 16px; font-weight: bold; color: #FFD700; margin: 10px;")
        layout.addWidget(title_label)
        self.cpu_ram_stats = QLabel("CPU: 0% | RAM: 0%")
        self.cpu_ram_stats.setStyleSheet("color: #FFD700; font-weight: bold; font-size: 14px; margin: 8px;")
        layout.addWidget(self.cpu_ram_stats)
        splitter = QSplitter(Qt.Vertical)
        # Latency graph
        latency_frame = QFrame()
        latency_layout = QVBoxLayout(latency_frame)
        self.latency_graph = RealTimeGraph(
            "Inference Latency Over Time", 
            "Latency (ms)", 
            max_points=300
        )
        latency_layout.addWidget(self.latency_graph)
        latency_info = QHBoxLayout()
        self.latency_stats = QLabel("Avg: 0ms | Max: 0ms | Spikes: 0")
        self.latency_stats.setStyleSheet("color: #00FFFF; font-weight: bold;")
        latency_info.addWidget(self.latency_stats)
        latency_info.addStretch()
        latency_layout.addLayout(latency_info)
        latency_frame.setLayout(latency_layout)
        splitter.addWidget(latency_frame)
        # FPS graph
        fps_frame = QFrame()
        fps_layout = QVBoxLayout(fps_frame)
        self.fps_graph = RealTimeGraph(
            "FPS & Resolution Impact", 
            "FPS", 
            max_points=300
        )
        fps_layout.addWidget(self.fps_graph)
        fps_info = QHBoxLayout()
        self.fps_stats = QLabel("Current FPS: 0 | Resolution: - | Device: -")
        self.fps_stats.setStyleSheet("color: #00FF00; font-weight: bold;")
        fps_info.addWidget(self.fps_stats)
        fps_info.addStretch()
        fps_layout.addLayout(fps_info)
        fps_frame.setLayout(fps_layout)
        splitter.addWidget(fps_frame)
        # Device switching & resolution changes graph
        device_frame = QFrame()
        device_layout = QVBoxLayout(device_frame)
        self.device_graph = RealTimeGraph(
            "Device Switching & Resolution Changes", 
            "-", 
            max_points=300
        )
        device_layout.addWidget(self.device_graph)
        self.device_legend = QLabel("<span style='color:#ff4444;'>CPU Spikes</span>: 0 | <span style='color:#ff5722;'>GPU Spikes</span>: 0 | <span style='color:#2196f3;'>Switches</span>: 0 | <span style='color:#ffa726;'>Res Changes</span>: 0")
        self.device_legend.setStyleSheet("color: #ffb300; font-size: 13px; font-weight: bold; margin: 2px 0 0 8px;")
        device_layout.addWidget(self.device_legend)
        device_frame.setLayout(device_layout)
        splitter.addWidget(device_frame)
        layout.addWidget(splitter)
        self.setLayout(layout)
    def update_graphs(self):
        # Placeholder for updating graphs with new data
        pass
    def update_performance_data(self, analytics_data: Dict[str, Any]):
        """Update graphs with new analytics data, including system metrics"""
        try:
            print(f"[PERF DEBUG] update_performance_data called with: {analytics_data}")
            chart_data = analytics_data.get('real_time_data', {})
            latency_stats = analytics_data.get('latency_statistics', {})
            current_metrics = analytics_data.get('current_metrics', {})
            system_metrics = analytics_data.get('system_metrics', {})
            if not chart_data.get('timestamps'):
                print("[PERF DEBUG] No timestamps in chart_data")
                return
            self.latest_data = {
                'chart_data': chart_data,
                'latency_stats': latency_stats,
                'current_metrics': current_metrics,
                'system_metrics': system_metrics
            }
            self.update_graphs()  # Immediately update graphs on new data
        except Exception as e:
            print(f"❌ Error updating performance data: {e}")
