import gradio as gr
import spaces
import torch
import time

import numpy as np


print(f"torch_version: {torch.__version__}")

# Define two matrices using NumPy arrays
#A = np.array([[1, 2], [3, 4]])
#B = np.array([[5, 6], [7, 8]])

# Define large matrices
A = np.random.rand(10000, 10000)  # Random 10000x10000 matrix
B = np.random.rand(10000, 10000)

# Start the timer
start_time = time.time()

# Perform matrix multiplication
result = np.dot(A, B)

# End the timer
end_time = time.time()

# Calculate and print the time taken
print(f"Time taken for matrix multiplication with NumPy: {end_time - start_time:.6f} seconds")


# Define two matrices
#A = torch.tensor([[1, 2], [3, 4]])
#B = torch.tensor([[5, 6], [7, 8]])

# Define large matrices
A = torch.rand(10000, 10000)  # Random 10000x10000 matrix
B = torch.rand(10000, 10000)

# Start the timer
start_time = time.time()

# Perform matrix multiplication
result = torch.matmul(A, B)

# End the timer
end_time = time.time()

# Calculate and print the time taken
print(f"Time taken for matrix multiplication with PyTorch: {end_time - start_time:.6f} seconds")

@spaces.GPU
def zeroGPU_test(text):
    # Define two matrices
    #A = torch.tensor([[1, 2], [3, 4]])
    #B = torch.tensor([[5, 6], [7, 8]])

    # Define large matrices
    A = torch.rand(10000, 10000).to('cuda')   # Random 10000x10000 matrix
    B = torch.rand(10000, 10000).to('cuda') 
    
    # Start the timer
    start_time = time.time()
    
    # Perform matrix multiplication
    result = torch.matmul(A, B)
    
    # End the timer
    end_time = time.time()
    print(f"Time taken for matrix multiplication with GPU: {end_time - start_time:.6f} seconds")
    return f"Time: {end_time - start_time:.6f} seconds"
     

demo = gr.Interface(fn=zeroGPU_test, inputs=gr.Text(), outputs=gr.Text())
demo.launch()