import gradio as gr
import torch
import time

def get_gpu_info():
    """Get information about available GPUs."""
    if torch.cuda.is_available():
        gpu_count = torch.cuda.device_count()
        gpu_info = []
        for i in range(gpu_count):
            name = torch.cuda.get_device_name(i)
            mem_total = torch.cuda.get_device_properties(i).total_memory / (1024**3)
            gpu_info.append(f"GPU {i}: {name} ({mem_total:.1f}GB)")
        return "\n".join(gpu_info)
    return "No GPU available"

def test_gpu(gpu_name):
    """Test the selected GPU with a simple computation."""
    if not torch.cuda.is_available():
        return "No GPU available on this system"
    
    # Find the GPU index based on the name
    device_idx = 0
    for i in range(torch.cuda.device_count()):
        if gpu_name in torch.cuda.get_device_name(i):
            device_idx = i
            break
    
    device = torch.device(f"cuda:{device_idx}")
    
    # Simple GPU computation test
    start_time = time.time()
    
    # Create a tensor on GPU
    x = torch.randn(1000, 1000, device=device)
    y = torch.randn(1000, 1000, device=device)
    
    # Perform matrix multiplication
    z = torch.matmul(x, y)
    
    # Move back to CPU
    result = z.cpu()
    
    elapsed = time.time() - start_time
    
    return f"GPU Test Successful!\n\nSelected GPU: {gpu_name}\nComputation Time: {elapsed:.4f} seconds\nResult shape: {result.shape}\nResult sum: {result.sum().item():.2f}"

def get_available_gpus():
    """Get list of available GPU names."""
    if torch.cuda.is_available():
        return [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]
    return ["No GPU available"]

# Get available GPUs
available_gpus = get_available_gpus()
gpu_info = get_gpu_info()

# Create Gradio interface
with gr.Blocks(title="GPU Tester") as demo:
    gr.Markdown("# GPU Selector and Tester")
    gr.Markdown(f"### Available GPU Information:\n```\n{gpu_info}\n```")
    
    with gr.Row():
        with gr.Column():
            gpu_dropdown = gr.Dropdown(
                choices=available_gpus if available_gpus else ["No GPU available"],
                label="Select GPU",
                value=available_gpus[0] if available_gpus else None
            )
            
            test_button = gr.Button("Test GPU", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(label="Test Results", lines=10)
    
    test_button.click(fn=test_gpu, inputs=gpu_dropdown, outputs=output)
    
    gr.Markdown("---")
    gr.Markdown("This app allows you to select a GPU and test its computation capability using PyTorch.")

# Launch the app
demo.launch(server_name="0.0.0.0", server_port=7860)