import datetime
print(f"Last Update: {datetime.datetime.now()}")

import torch
import matplotlib.pyplot as plt
from torch import nn

device = "cuda" if torch.cuda.is_available() else "cpu"
device
print(f"Device:: {device}")

# A：create dataset

# create data parameter
weight = 0.3
bias = 0.9

# make X and y using linear regression feature
X = torch.arange(0,1,0.01).unsqueeze(dim = 1)
y = weight * X + bias
print(f"Number of X Sample: {len(X)}")
print(f"Number of y Sample: {len(y)}")
print(f"First 10 X & y Sample:\nX: {X[:10]}\ny: {y[:10]}")

# split data into training and testing
train_split = int(len(X) * 0.8)
X_train = X[:train_split]
y_train = y[:train_split]
X_test = X[train_split:]
y_test = y[train_split:]
len(X_train),len(y_train),len(X_test),len(y_test)

# plot training and testing data
def plot_prediction(train_data = X_train,
                 train_labels = y_train,
                 test_data = X_test,
                 test_labels = y_test,
                 prediction = None):
  plt.figure(figsize = (10,7))
  plt.scatter(train_data,train_labels,c = 'b',s = 4,label = "Training Data")
  plt.scatter(test_data,test_labels,c = 'g',s = 4,label = "Testing Data")

  if prediction is not None:
    plt.scatter(test_data,prediction,c = 'r',s = 4,label = "Prediction")
  plt.legend(prop = {"size" : 14})
plot_prediction()

plt.savefig("VisualizeDataSplittingExercise.svg")

# B：build model
# create PyTorch linear regression model by subclassing nn.Module
## Option 1
class LinearRegressionModel(nn.Module):
  def __init__(self):
    super().__init__()
    self.weight = nn.Parameter(data=torch.randn(1,requires_grad=True,dtype=torch.float))
    self.bias = nn.Parameter(data=torch.randn(1,requires_grad=True,dtype=torch.float))
  def forward(self, x):
    return self.weight * x + self.bias

# ## Option 2
# class LinearRegressionModel(nn.Module):
#   def __init__(self):
#     super().__init__()
#     self.linear_layer = nn.Linear(in_features = 1,out_features = 1)
#   def forward(self,x : torch.Tensor) -> torch.Tensor:
#     return self.linear_layer(x)

torch.manual_seed(42)
model_1 = LinearRegressionModel()
model_1,model_1.state_dict()

next(model_1.parameters()).device

# instantiate model and put it to target device
model_1.to(device)
list(model_1.parameters())

# C：loss function and optimizer
# create loss function and optimizer
loss_fn = nn.L1Loss()
optimizer = torch.optim.SGD(params = model_1.parameters(),lr = 0.01)

# Training loop
# Train model for 300 epoch
torch.manual_seed(42)

epoch = 300

# send data to target device
X_train = X_train.to(device)
X_test = X_test.to(device)
y_train = y_train.to(device)
y_test = y_test.to(device)

for e in range(epoch):
  ### Training

  # Put model in train mode
  model_1.train()

  # 1. Forward pass
  y_pred = model_1(X_train)

  # 2. Calculate loss
  loss = loss_fn(y_pred,y_train)

  # 3. Zero gradients
  optimizer.zero_grad()

  # 4. Backpropagation
  loss.backward()

  # 5. Step the optimizer
  optimizer.step()

  ### Perform testing every 20 epochs
  if epoch % 20 == 0:
    # Put model in evaluation mode and setup inference context
    model_1.eval()
    with torch.inference_mode():
      # 1. Forward pass
      y_preds = model_1(X_test)
      # 2. Calculate test loss
      test_loss = loss_fn(y_preds,y_test)
      # Print out what's happening
      print(f"Epoch: {e} | Train loss: {loss:.3f} | Test loss: {test_loss:.3f}")

# D：Make Prediction

# Make predictions with model
model_1.eval()

with torch.inference_mode():
  y_pred = model_1(X_test)
y_pred

# y_preds.cpu()

# ploy prediction (these may need to be on a specific device)
plot_prediction(prediction = y_pred.cpu())

plt.savefig("VisualizeBuildModelPredictionExercise.svg")

# E:save model
from pathlib import Path

# 1. Create models directory
MODEL_PATH = Path("model")
MODEL_PATH.mkdir(parents = True,exist_ok = True)
# 2. Create model save path
MODEL_NAME = "01_pytorch_model"
MODEL_SAVE_PATH = MODEL_PATH / MODEL_NAME
# 3. Save the model state dict
print(f"Saving model to {MODEL_SAVE_PATH}")
torch.save(obj = model_1.state_dict(),f = MODEL_SAVE_PATH)

# Create new instance of model and load saved
# state dict (make sure to put it on the target device)
loaded_model = LinearRegressionModel()
loaded_model.load_state_dict(torch.load(f = MODEL_SAVE_PATH))
loaded_model.to(device)

# Make predictions with loaded model and compare them to the previous
y_pred_new = loaded_model(X_test)
print(y_pred == y_pred_new)

print(loaded_model.state_dict())