# Import the top-level package for core functionality
import torch

# Import neural network functionality
from torch import nn

# Import functional programming tools
import torch.nn.functional as F

# Import optimization functionality
import torch.optim as optim

# Import dataset functions
from torch.utils.data import TensorDataset, DataLoader

# Import evaluation metrics
import torchmetrics

import pandas as pd

# Create tensor from list with tensor()
tnsr = torch.tensor([1, 3, 6, 10])

# Get data type of tensor elements with .dtype
print(tnsr.dtype) # Returns torch.int64

# Get dimensions of tensor with .Size()
print(tnsr.shape) # Returns torch.Size([4])

# Get memory location of tensor with .device
print(tnsr.device) # Returns cpu or gpu

# Create a tensor of zeros with zeros()
tnsr_zrs = torch.zeros(2, 3)

# Create a random tensor with rand()
tnsr_rndm = torch.rand(size=(3, 4)) # Tensor has 3 rows, 4 columns

# 1. Create a sample DataFrame (replace this with your actual data, e.g., pd.read_csv())
data = {
    'feature_1': [1.5, 2.0, 3.1, 4.2, 5.0],
    'feature_2': [10, 20, 30, 40, 50],
    'target': [100, 150, 200, 250, 300]
}
df = pd.DataFrame(data)

# 2. Define the columns and batch size (n)
feature_columns = ['feature_1', 'feature_2']
target_column = 'target'
n = 2  # Batch size

# 3. Extract values and convert to PyTorch tensors
X = df[feature_columns].values
y = df[target_column].values

# 4. Create the dataset and dataloader
dataset = TensorDataset(torch.tensor(X).float(), torch.tensor(y).float())
dataloader = DataLoader(dataset, batch_size=n, shuffle=True)

# 5. Test the DataLoader to verify it works
for batch_X, batch_y in dataloader:
    print("Batch Features (X):\n", batch_X)
    print("Batch Target (y):\n", batch_y)
    break

## Data Preprocessing

# One-hot encode categorical variables with one_hot()
F.one_hot(torch.tensor([0, 1, 2]), num_classes=3) # Returns tensor of 0s and 1s

# Create a sequential model from layers

# Define dummy dimensions and parameters for demonstration
m, n = 10, 5
n_features, i, j, n_classes = 10, 32, 16, 3

# Create a linear layer with m inputs, n outputs with Linear()
lnr = nn.Linear(m, n)

# Get weight of layer with .weight
print("Layer weights:", lnr.weight)

# Get bias of layer with .bias
print("Layer bias:", lnr.bias)

# Create activation and regularization layers
# Create a sigmoid activation layer for binary classification with Sigmoid()
sigmoid_layer = nn.Sigmoid()

# Create a softmax activation layer for multi-class classification with Softmax()
softmax_layer = nn.Softmax(dim=-1)

# Create a rectified linear unit activation layer to avoid saturation with ReLU()
relu_layer = nn.ReLU()

# Create a leaky rectified linear unit activation layer to avoid saturation with LeakyReLU()
leaky_relu_layer = nn.LeakyReLU(negative_slope=0.05)

# Create a dropout layer to regularize and prevent overfitting with Dropout()
dropout_layer = nn.Dropout(p=0.5)

# Create a sequential model from layers (fixed the missing closing parenthesis)
model = nn.Sequential(
    nn.Linear(n_features, i),
    nn.Linear(i, j),       # Input size must match output from previous layer
    nn.Linear(j, n_classes),
    nn.Softmax(dim=-1)     # Activation layer comes last
)

print("\nModel Architecture:\n", model)

# Define dimensions
n_features, i, j, n_classes = 10, 32, 16, 3
batch_size = 5

# Create sample input data
input_data = torch.randn(batch_size, n_features)

# Fit a model to input data with model
prediction = model(input_data).double()

# Get target values
actual = torch.randn(batch_size, n_classes).double()

# Calculate the mean-squared error loss for regression with MSELoss()
mse_loss = nn.MSELoss()(prediction, actual) # Returns tensor(x) 
print("MSE Loss:", mse_loss.item())

# Calculate the L1 loss for robust regression with SmoothL1Loss()
l1_loss = nn.SmoothL1Loss()(prediction, actual) # Returns tensor(x) 
print("Smooth L1 Loss:", l1_loss.item())

# Calculate binary cross-entropy loss for binary classification with BCELoss()
actual_bce = torch.tensor([
    [0.0, 1.0, 0.0],
    [1.0, 1.0, 1.0],
    [0.0, 0.0, 1.0],
    [1.0, 0.0, 0.0],
    [0.0, 1.0, 1.0]
]).double()
bce_loss = nn.BCELoss()(prediction, actual_bce) # Returns tensor(x) 
print("BCE Loss:", bce_loss.item())

# Calculate cross-entropy loss for multi-class classification with CrossEntropyLoss()
ce_loss = nn.CrossEntropyLoss()(prediction, actual) # Returns tensor(x) 
print("Cross-Entropy Loss:", ce_loss.item())

# Calculate the gradients via backprogagation with .backward()
loss = ce_loss
loss.backward()
print("Backpropagation completed successfully!")

# Create a stochastic gradient descent optimizer with SGD(), setting learning rate and momentum
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.95)

# Update neuron parameters with .step()
optimizer.step()

# Set up dummy data and dataloader
n_features, n_classes = 10, 3
batch_size = 5
X_data = torch.randn(batch_size, n_features)
y_data = torch.randn(batch_size, n_classes)
dataset = TensorDataset(X_data, y_data)
dataloader = DataLoader(dataset, shuffle=True)

# Set model to training mode
model.train()
# Set a loss criterion and an optimizer
loss_criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.95)
# Loop over chunks of data in the training set
for data in dataloader:
    # Set the gradients to zero with .zero_grad()
    optimizer.zero_grad()
    # Get features and targets for current chunk of data
    features, targets = data
    # Run a "forward pass" to fit the model to the data
    predictions = model(features)
    # Calculate loss
    loss = loss_criterion(predictions, targets)
    # Calculate gradients using backprogagation
    loss.backward()
    # Update the model parameters
    optimizer.step()
print("Training loop completed successfully!")

# Set model to evaluation mode
model.eval()

# Create accuracy metric with Accuracy()
metric = torchmetrics.Accuracy(task="multiclass", num_classes=3)
# Loop of chunks of data in the validation set
for i, data in enumerate(dataloader, 0):
    # Get features and targets for current chunk of data
    features, targets = data
    # Run a "forward pass" to fit the model to the data
    predictions = model(features)
    # Calculate accuracy over the batch
    accuracy = metric(targets, predictions.argmax(dim=-1))
# Calculate accuracy over all the validation data
accuracy = metric.compute()
print(f"Accuracy on all data: {accuracy}")
# Reset the metric for the next dataset (training or validation)
metric.reset()

# Extract layer 0 (the first Linear layer)
layer = model[0]

# Save a layer of a model to a file with save()
torch.save(layer, 'layer.pth')
print("Layer saved to 'layer.pth'")

# Load a layer of a model from a file with load()
new_layer = torch.load('layer.pth', weights_only=False)
print("Loaded layer architecture:\n", new_layer)

# Freeze the weight for layer 0 with .requires_grad
for name, param in model.named_parameters():
    if name == "0.weight":
        param.requires_grad = False
        print(f"Parameter '{name}' successfully frozen (requires_grad = {param.requires_grad})")