import streamlit as st
import torch
import io
import sys

# Function to execute the input code and capture print statements
def execute_code(code):
    # Redirect stdout to capture print statements
    old_stdout = sys.stdout
    sys.stdout = mystdout = io.StringIO()
    
    global_vars = {"torch": torch}
    local_vars = {}
    try:
        exec(code, global_vars, local_vars)
        output = mystdout.getvalue()
    except Exception as e:
        output = str(e)
    finally:
        # Reset redirect.
        sys.stdout = old_stdout
    
    return output, local_vars

# Dictionary with exercise details
exercises = {
    "Exercise 1: Create and Manipulate Tensors": {
        "description": "Tensors are the core data structure in PyTorch, similar to arrays in NumPy but with additional capabilities for GPU acceleration. This exercise introduces how to create tensors from various data sources such as lists and NumPy arrays. It also covers basic tensor operations like addition, subtraction, and element-wise multiplication, which are fundamental for manipulating data in PyTorch.",
        "code": '''import torch
import numpy as np

# Creating tensors from Python lists
# This creates a 1D tensor from the list [1, 2, 3]
tensor_from_list = torch.tensor([1, 2, 3])
print("Tensor from list:", tensor_from_list)

# Creating tensors from NumPy arrays
# This converts a NumPy array to a tensor
numpy_array = np.array([4, 5, 6])
tensor_from_numpy = torch.tensor(numpy_array)
print("Tensor from NumPy array:", tensor_from_numpy)

# Performing basic tensor operations
tensor1 = torch.tensor([1, 2, 3])
tensor2 = torch.tensor([4, 5, 6])

# Addition
addition = tensor1 + tensor2
print("Addition:", addition)

# Subtraction
subtraction = tensor1 - tensor2
print("Subtraction:", subtraction)

# Element-wise multiplication
elementwise_multiplication = tensor1 * tensor2
print("Element-wise Multiplication:", elementwise_multiplication)
'''
    },
    "Exercise 2: Tensor Indexing and Slicing": {
        "description": "Indexing and slicing allow you to access and manipulate specific elements and sub-tensors. This is crucial for tasks such as data preprocessing and manipulation in machine learning workflows. This exercise demonstrates how to index and slice tensors to extract and modify elements efficiently.",
        "code": '''import torch

# Creating a 2D tensor (matrix)
tensor = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# Indexing elements
# Accessing the element at the 2nd row and 3rd column (indexing starts at 0)
element = tensor[1, 2]
print("Element at index [1, 2]:", element)

# Slicing sub-tensors
# Extracting the entire second row
row = tensor[1, :]
print("Second row:", row)

# Extracting the entire third column
column = tensor[:, 2]
print("Third column:", column)

# Modifying elements
# Changing the first element of the tensor to 10
tensor[0, 0] = 10
print("Modified tensor:", tensor)
'''
    },
    "Exercise 3: Tensor Reshaping and Transposing": {
        "description": "Reshaping and transposing tensors are common operations when preparing data for neural networks. Reshaping allows you to change the layout of tensor data without altering its data. Transposing changes the orientation of the data, which is useful for operations like matrix multiplication. This exercise covers reshaping using `view()` and `reshape()`, and transposing using `transpose()` and `permute()`.",
        "code": '''import torch

# Creating a 2D tensor
tensor = torch.tensor([[1, 2, 3], [4, 5, 6]])

# Reshaping the tensor
# Changing the shape of the tensor to (3, 2)
reshaped_tensor = tensor.view(3, 2)
print("Reshaped tensor:", reshaped_tensor)

# Another way to reshape using reshape()
reshaped_tensor2 = tensor.reshape(-1)  # Flattening the tensor
print("Reshaped tensor (flattened):", reshaped_tensor2)

# Transposing the tensor
# Swapping the dimensions of the tensor (transpose rows and columns)
transposed_tensor = tensor.t()
print("Transposed tensor:", transposed_tensor)

# Using permute for higher-dimensional tensors
tensor_3d = torch.randn(2, 3, 4)  # Creating a random 3D tensor
# Permuting dimensions
permuted_tensor = tensor_3d.permute(2, 0, 1)
print("Permuted tensor shape:", permuted_tensor.shape)
'''
    },
    "Exercise 4: Tensor Operations and Broadcasting": {
        "description": "Broadcasting in PyTorch allows you to perform arithmetic operations on tensors of different shapes. This can simplify code and reduce the need for explicit reshaping. This exercise demonstrates matrix multiplication, the concept of broadcasting, and the application of element-wise functions like sine and exponential.",
        "code": '''import torch

# Matrix multiplication
tensor1 = torch.tensor([[1, 2], [3, 4]])
tensor2 = torch.tensor([[5, 6], [7, 8]])

# Performing matrix multiplication using matmul
matmul_result = torch.matmul(tensor1, tensor2)
print("Matrix multiplication result:", matmul_result)

# Broadcasting example
# Broadcasting allows tensor3 and tensor4 to be added despite their different shapes
tensor3 = torch.tensor([1, 2, 3])
tensor4 = torch.tensor([[1], [2], [3]])
broadcast_result = tensor3 + tensor4
print("Broadcasting result:", broadcast_result)

# Element-wise functions
# Applying sine function element-wise
sin_result = torch.sin(tensor3)
print("Sine of tensor3:", sin_result)

# Applying exponential function element-wise
exp_result = torch.exp(tensor3)
print("Exponential of tensor3:", exp_result)

# Applying square root function element-wise
# Note: sqrt requires the tensor to be of floating-point type
sqrt_result = torch.sqrt(tensor3.float())
print("Square root of tensor3:", sqrt_result)
'''
    },
    "Exercise 5: Tensor Initialization": {
        "description": "Proper initialization of tensors is crucial for neural network training and other machine learning tasks. This exercise explores various methods to initialize tensors, such as filling them with zeros, ones, random values, and specific distributions. Understanding these methods helps in setting up model parameters and test data.",
        "code": '''import torch

# Creating tensors filled with zeros and ones
zeros_tensor = torch.zeros(3, 3)
print("Zeros tensor:", zeros_tensor)

ones_tensor = torch.ones(3, 3)
print("Ones tensor:", ones_tensor)

# Randomly initialized tensors
# Uniform distribution in the range [0, 1)
rand_tensor = torch.rand(3, 3)
print("Uniform random tensor:", rand_tensor)

# Normal distribution with mean 0 and variance 1
randn_tensor = torch.randn(3, 3)
print("Normal random tensor:", randn_tensor)

# Random integers in the range [0, 10)
randint_tensor = torch.randint(0, 10, (3, 3))
print("Random integer tensor:", randint_tensor)

# Initializing tensors with specific distributions
# Normal distribution with custom mean and standard deviation
normal_tensor = torch.normal(mean=0, std=1, size=(3, 3))
print("Normal distribution tensor:", normal_tensor)

# Uniform distribution in a custom range [0, 1)
uniform_tensor = torch.empty(3, 3).uniform_(0, 1)
print("Uniform distribution tensor:", uniform_tensor)
'''
    },
    "Exercise 6: Tensor Arithmetic Operations": {
        "description": "Arithmetic operations on tensors are essential for numerous tasks in machine learning, including data preprocessing and neural network computations. This exercise covers basic arithmetic with tensors, in-place operations, and aggregation functions that compute summaries over tensor elements.",
        "code": '''import torch

# Creating a 1D tensor
tensor = torch.tensor([1.0, 2.0, 3.0])

# Scalar-tensor operations
# Adding a scalar to each element of the tensor
add_scalar = tensor + 5
print("Addition with scalar:", add_scalar)

# Multiplying each element by a scalar
mul_scalar = tensor * 2
print("Multiplication with scalar:", mul_scalar)

# In-place operations (modify the original tensor)
# Adding 5 to each element of the tensor
tensor.add_(5)
print("In-place addition:", tensor)

# Tensor aggregation operations
tensor = torch.tensor([[1, 2], [3, 4]])

# Summing all elements
sum_all = tensor.sum()
print("Sum of all elements:", sum_all)

# Computing the mean of all elements
mean_all = tensor.mean()
print("Mean of all elements:", mean_all)

# Finding the maximum element
max_all = tensor.max()
print("Maximum element:", max_all)
'''
    },
    "Exercise 7: Tensor Comparison and Logical Operations": {
        "description": "Comparison and logical operations enable you to make conditional selections and perform logical checks on tensors. This is crucial for tasks such as filtering data and implementing decision-based logic in machine learning workflows. This exercise demonstrates the use of comparison operators and logical functions, which are often used in neural network training and data analysis.",
        "code": '''import torch

# Creating two tensors for comparison
tensor1 = torch.tensor([1, 2, 3])
tensor2 = torch.tensor([3, 2, 1])

# Comparison operations
# Checking if elements in tensor1 are greater than corresponding elements in tensor2
greater_than = tensor1 > tensor2
print("Greater than comparison:", greater_than)

# Checking if elements in tensor1 are equal to corresponding elements in tensor2
equal = tensor1 == tensor2
print("Equality comparison:", equal)

# Logical operations
# Logical AND operation between comparison results
logical_and = torch.logical_and(tensor1 > 1, tensor2 < 3)
print("Logical AND:", logical_and)

# Logical OR operation between comparison results
logical_or = torch.logical_or(tensor1 < 3, tensor2 > 1)
print("Logical OR:", logical_or)

# Logical NOT operation
logical_not = torch.logical_not(tensor1 < 3)
print("Logical NOT:", logical_not)

# Conditional selection using torch.where
# Selecting elements from tensor1 if condition is True, otherwise from tensor2
selected_tensor = torch.where(tensor1 > 2, tensor1, tensor2)
print("Conditional selection:", selected_tensor)
'''
    },
    "Exercise 8: Tensor Reduction Operations": {
        "description": "Reduction operations aggregate tensor elements along specified dimensions, providing summarized results like sums, means, and indices of maximum or minimum values. This exercise covers essential reduction operations, which are commonly used in loss functions and performance metrics in machine learning.",
        "code": '''import torch

# Creating a 2D tensor (matrix)
tensor = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# Reduction operations along dimensions
# Summing elements along dimension 0 (rows)
sum_dim0 = tensor.sum(dim=0)
print("Sum along dimension 0:", sum_dim0)

# Computing the mean along dimension 1 (columns)
mean_dim1 = tensor.mean(dim=1)
print("Mean along dimension 1:", mean_dim1)

# Computing the product of all elements in the tensor
prod_all = tensor.prod()
print("Product of all elements:", prod_all)

# Advanced reduction operations
# Finding the index of the maximum element in the flattened tensor
argmax_all = tensor.argmax()
print("Index of maximum element:", argmax_all)

# Finding the indices of the minimum elements along dimension 0
argmin_dim0 = tensor.argmin(dim=0)
print("Indices of minimum elements along dimension 0:", argmin_dim0)
'''
    },
    "Exercise 9: Tensor Cloning and Detachment": {
        "description": "Cloning and detaching tensors are important when working with gradients in neural networks. Cloning creates a copy of a tensor, while detaching removes it from the computation graph to prevent gradient tracking. This exercise demonstrates these operations, which are crucial for managing tensor operations in complex neural networks.",
        "code": '''import torch

# Creating a tensor with gradient tracking enabled
tensor = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# Cloning the tensor
# This creates a new tensor with the same data but no gradient history
cloned_tensor = tensor.clone()
print("Original tensor:", tensor)
print("Cloned tensor:", cloned_tensor)

# Detaching the tensor from the computation graph
# This removes the tensor from the gradient computation graph
detached_tensor = tensor.detach()
print("Detached tensor:", detached_tensor)

# Comparing original and detached tensors
# Modifying the original tensor (note that in-place operations modify the tensor in place)
tensor.add_(1)
print("Modified original tensor:", tensor)
print("Detached tensor remains unchanged:", detached_tensor)
'''
    },
    "Exercise 10: GPU Operations with Tensors": {
        "description": "Utilizing GPUs can significantly speed up tensor operations, which is crucial for training large neural networks. This exercise covers checking for GPU availability, moving tensors to GPU, and performing operations on GPU to compare performance with CPU operations.",
        "code": '''import torch
import time

# Check for GPU availability
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

# Creating a tensor and moving it to GPU if available
tensor = torch.tensor([1.0, 2.0, 3.0])
tensor = tensor.to(device)
print("Tensor on GPU:", tensor)

# Perform operations on GPU
# Adding a scalar to each element of the tensor on GPU
tensor_add = tensor + 5
print("Tensor after addition on GPU:", tensor_add)

# Compare speed of operations on CPU vs GPU
# Creating a large tensor for performance comparison
cpu_tensor = torch.randn(10000, 10000)

# CPU operation
start_time = time.time()
cpu_result = cpu_tensor @ cpu_tensor.T  # Matrix multiplication
end_time = time.time()
print("CPU operation time:", end_time - start_time, "seconds")

# GPU operation (if available)
if torch.cuda.is_available():
    gpu_tensor = cpu_tensor.to(device)
    start_time = time.time()
    gpu_result = gpu_tensor @ gpu_tensor.T  # Matrix multiplication
    torch.cuda.synchronize()  # Wait for GPU operation to finish
    end_time = time.time()
    print("GPU operation time:", end_time - start_time, "seconds")
'''
    }
}

st.title('PyTorch Code Runner')

# Side menu for exercises
exercise_choice = st.sidebar.radio("Choose an exercise", list(exercises.keys()))

# Display the chosen exercise description
st.subheader(exercise_choice)
st.write(exercises[exercise_choice]["description"])

# Text area for inputting the PyTorch code
code_input = st.text_area("Enter your PyTorch code here", height=300, value=exercises[exercise_choice]["code"])

# Button to execute the code
if st.button("Run Code"):
    # Prepend the import statement
    code_to_run = "import torch\n" + code_input
    
    # Execute the code and capture the output
    output, variables = execute_code(code_to_run)
    
    # Display the output
    st.subheader('Output')
    st.text(output)
    
    # Display returned variables
    if variables:
        st.subheader('Variables')
        for key, value in variables.items():
            st.text(f"{key}: {value}")