#!/usr/bin/env python3
"""
Test GPU Client (Antigravity Version)
=====================================
Questo script usa la libreria 'antigravity' per verificare la GPU remota.
"""

from antigravity_sdk import RemoteGPU

# Codice di verifica da eseguire sul server
REMOTE_CHECK_CODE = r'''
import torch
import sys
import time

print("=" * 60)
print("🔍 REMOTE GPU DEEP CHECK")
print("=" * 60)

# 1. Verifica disponibilità CUDA
print(f"   PyTorch version: {torch.__version__}")
print(f"   CUDA disponibile: {torch.cuda.is_available()}")

if torch.cuda.is_available():
    print(f"   CUDA version: {torch.version.cuda}")
    print(f"   Numero GPU: {torch.cuda.device_count()}")
    
    # 2. Informazioni sulla GPU
    print("\n🎮 Dettagli GPU:")
    for i in range(torch.cuda.device_count()):
        props = torch.cuda.get_device_properties(i)
        print(f"   GPU {i}: {props.name}")
        print(f"      - Memoria Totale: {props.total_memory / 1024**3:.2f} GB")
        print(f"      - Multi-Processors: {props.multi_processor_count}")
    
    # 3. Test calcolo tensoriale
    print("\n⚡ Test Calcolo Tensoriale:")
    size = 4000
    print(f"   Matrice {size}x{size} (Float32)...")
    
    device = torch.device("cuda")
    try:
        a = torch.randn(size, size, device=device)
        b = torch.randn(size, size, device=device)
        
        # Warmup
        torch.matmul(a[:100,:100], b[:100,:100])
        torch.cuda.synchronize()
        
        start = time.time()
        c = torch.matmul(a, b)
        torch.cuda.synchronize()
        elapsed = time.time() - start
        
        print(f"   ✅ Calcolo completato in {elapsed:.4f} secondi")
        print(f"   🚀 Speed: { (2*size**3) / (elapsed * 1e12):.2f} TFLOPS (approx)")
        
    except Exception as e:
        print(f"   ❌ Errore Calcolo: {e}")
    
    # 4. Verifica memoria
    print("\n💾 Utilizzo Memoria GPU:")
    allocated = torch.cuda.memory_allocated() / 1024**3
    print(f"   Allocata: {allocated:.2f} GB")
    
    print("\n" + "=" * 60)
    print("✅ GPU REMOTA PIENAMENTE OPERATIVA")
    print("=" * 60)
    
    # 5. Genera Plot di prova (così il client ha qualcosa da scaricare)
    print("\n🎨 Generazione Grafico di Test...")
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    
    plt.figure(figsize=(8, 4))
    plt.bar(['MatMul GPU'], [elapsed], color='green')
    plt.title(f'GPU Computation Time: {elapsed:.4f}s')
    plt.ylabel('Seconds')
    plt.savefig('gpu_test_plot.png')
    print("📦 Grafico salvato: gpu_test_plot.png")

else:
    print("\n❌ CUDA NON DISPONIBILE - Running on CPU")
    sys.exit(1)

'''

def main():
    print("🚀 Connecting using 'antigravity' library...")
    
    try:
        # Inizializza client (URL default)
        gpu = RemoteGPU()
        
        # Esegue codice
        result = gpu.run(REMOTE_CHECK_CODE)
        
        # Il risultato viene già stampato se verbose=True (default),
        # ma per sicurezza controlliamo l'output
        if "GPU REMOTA PIENAMENTE OPERATIVA" in result.output:
            print("\n🎉 GREAT SUCCESS! La libreria funziona e la GPU risponde.")
        else:
            print("\n⚠️  Il test non sembra aver confermato la GPU (controlla output sopra).")
            
    except Exception as e:
        print(f"❌ Errore Libreria: {e}")

if __name__ == "__main__":
    main()