import torch
import matplotlib.pyplot as plt

x = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7.])
print(f'x = {x}')

y = torch.tensor([1.86, 1.31, 0.62, 0.33, 0.09, -0.67, -1.23, -1.37])
print(f'y = {y}')

fig, ax = plt.subplots()
plt.title('Clinical Trial')
plt.xlabel('Drug dosage (mL)')
plt.ylabel('Foregetfulness')
ax.scatter(x, y)

m = torch.tensor([0.9], requires_grad=True)
b = torch.tensor([0.1], requires_grad=True)
print(f'm={m},b={b}')

def regression(my_x, my_m, my_b):
    return my_m * my_x + my_b

def regression_plot(my_x, my_y, my_m, my_b):
    fig, ax = plt.subplots()
    ax.scatter(my_x, my_y)
    x_min, x_max = ax.get_xlim()
    y_min = regression(x_min, my_m, my_b).detach().item()
    y_max = regression(x_max, my_m, my_b).detach().item()
    ax.set_xlim([x_min, x_max])
    ax.plot([x_min, x_max], [y_min, y_max])

regression_plot(x, y, m, b)