import torch
import torch.nn as nn
import numpy as np
from sklearn import datasets
import matplotlib.pyplot as plt

# step 0: prepare data
X_numpy, y_numpy = datasets.make_regression(n_samples = 100, n_features=1, noise=20, random_state=1)
X = torch.from_numpy(X_numpy.astype(np.float32))
y = torch.from_numpy(y_numpy.astype(np.float32))
y = y.view(y.shape[0], 1) # reshape to column vector

n_samples, n_features = X.shape
print('Dataset is ready')
# step 1: model

input_size = n_features
output_size = 1
model = nn.Linear(input_size, output_size)

# step 2: loss and optimizer
learning_rate = 0.01
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

# step 3: training loop
num_epochs = 1000
for epoch in range(num_epochs):
    # forward pass and loss
    y_predicted = model(X)
    loss = criterion(y_predicted, y)

    #backward pass
    loss.backward()

    # update

    optimizer.step()

    optimizer.zero_grad() # never forget this

    if (epoch + 1) % 10 == 0:
        print (f'epoch: {epoch+1}, loss = {loss.item():.4f}')

# plot
predicted = model(X).detach().numpy() #new tensor where graident calculation is false
plt.plot(X_numpy, y_numpy, 'ro')
plt.plot(X_numpy,predicted, 'b')
plt.savefig('plot.png')
print('Training is complete')