#!/usr/bin/env python3
import torch
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from datetime import datetime
import subprocess
import time
import psutil
import re
from collections import deque
import threading

class GPUBenchmark:
    def __init__(self):
        self.max_temp = 85
        self.temperatures = deque(maxlen=100)
        self.tflops_history = deque(maxlen=100)
        self.times = deque(maxlen=100)
        self.peak_tflops = 0
        self.running = True
        self.stress_size = 8192
        self.fig, (self.ax1, self.ax2) = plt.subplots(2, 1, figsize=(12, 8))
        self.fig.suptitle('Benchmark TFLOPS - Radeon Pro VII', fontsize=16)
        
    def get_gpu_temp(self):
        try:
            result = subprocess.run(['sensors'], capture_output=True, text=True, timeout=1)
            for line in result.stdout.split('\n'):
                if 'edge:' in line.lower():
                    match = re.search(r'([+-]?\d+\.?\d*)\s*°C', line)
                    if match:
                        return float(match.group(1))
        except:
            return 0
        return 0

    def check_system_responsiveness(self):
        try:
            start = time.time()
            _ = psutil.cpu_percent(interval=0.1)
            return (time.time() - start) < 0.5
        except:
            return False

    def calculate_tflops(self, matrix_size, elapsed_time):
        operations = 2 * (matrix_size ** 3)
        return (operations / elapsed_time) / 1e12

    def stress_gpu(self):
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        if device.type == 'cpu':
            print("ERRO: GPU não detectada!")
            self.running = False
            return
        
        print(f"GPU detectada: {torch.cuda.get_device_name(0)}")
        print("Iniciando benchmark...\n")
        
        while self.running:
            temp = self.get_gpu_temp()
            if temp >= self.max_temp:
                print(f"\n⚠ TEMPERATURA LIMITE: {temp}°C")
                self.running = False
                break
            
            if not self.check_system_responsiveness():
                print("\n⚠ SISTEMA TRAVANDO")
                self.running = False
                break
            
            try:
                torch.cuda.synchronize()
                start = time.time()
                a = torch.randn(self.stress_size, self.stress_size, device=device)
                b = torch.randn(self.stress_size, self.stress_size, device=device)
                c = torch.mm(a, b)
                torch.cuda.synchronize()
                elapsed = time.time() - start
                
                tflops = self.calculate_tflops(self.stress_size, elapsed)
                self.temperatures.append(temp)
                self.tflops_history.append(tflops)
                self.times.append(datetime.now())
                
                if tflops > self.peak_tflops:
                    self.peak_tflops = tflops
                
                print(f"TFLOPS: {tflops:.2f} | Temp: {temp}°C | Peak: {self.peak_tflops:.2f}", end='\r')
                
                if temp < 75 and tflops < self.peak_tflops * 0.9:
                    self.stress_size = min(self.stress_size + 256, 16384)
                elif temp > 80:
                    self.stress_size = max(self.stress_size - 256, 4096)
                
                time.sleep(0.1)
            except Exception as e:
                print(f"\n⚠ ERRO: {e}")
                self.running = False
                break

    def update_plot(self, frame):
        if not self.running and len(self.tflops_history) == 0:
            return
        
        self.ax1.clear()
        self.ax2.clear()
        
        if len(self.tflops_history) > 0:
            self.ax1.plot(list(self.tflops_history), 'b-', linewidth=2, label='TFLOPS atual')
            self.ax1.axhline(y=self.peak_tflops, color='g', linestyle='--', 
                           label=f'Peak: {self.peak_tflops:.2f} TFLOPS')
            self.ax1.set_ylabel('TFLOPS', fontsize=12)
            self.ax1.set_title('Desempenho em Tempo Real')
            self.ax1.legend()
            self.ax1.grid(True, alpha=0.3)
        
        if len(self.temperatures) > 0:
            self.ax2.plot(list(self.temperatures), 'r-', linewidth=2, label='Temperatura')
            self.ax2.axhline(y=self.max_temp, color='orange', linestyle='--', 
                           label=f'Limite: {self.max_temp}°C')
            self.ax2.set_ylabel('Temperatura (°C)', fontsize=12)
            self.ax2.set_xlabel('Amostras', fontsize=12)
            self.ax2.legend()
            self.ax2.grid(True, alpha=0.3)
        
        if not self.running:
            self.ax1.text(0.5, 0.5, f'PEAK TFLOPS: {self.peak_tflops:.2f}', 
                         transform=self.ax1.transAxes, fontsize=20, 
                         ha='center', color='green', weight='bold')

    def run(self):
        stress_thread = threading.Thread(target=self.stress_gpu)
        stress_thread.daemon = True
        stress_thread.start()
        
        ani = animation.FuncAnimation(self.fig, self.update_plot, 
                                     interval=500, cache_frame_data=False)
        plt.tight_layout()
        plt.show()
        stress_thread.join(timeout=2)
        
        print(f"\n\n{'='*50}")
        print(f"RESULTADO FINAL")
        print(f"{'='*50}")
        print(f"🏆 PEAK TFLOPS: {self.peak_tflops:.2f}")
        print(f"🌡️  Temp máxima: {max(self.temperatures) if self.temperatures else 0:.1f}°C")
        print(f"{'='*50}\n")

if __name__ == "__main__":
    benchmark = GPUBenchmark()
    benchmark.run()