import tensorflow as tf
import pandas as pd
from matplotlib import pyplot as plt
tf.random.set_seed(22)

num_obs = 150
x1 = tf.random.uniform([num_obs], minval=0, maxval=10, dtype=tf.float32)
x2 = tf.random.uniform([num_obs], minval=0, maxval=10, dtype=tf.float32)

b, w_1, w_2 = 1.5, 2.0, 5.0
y = b + w_1 * x1 + w_2 * x2

dataset = tf.stack([y, x1, x2], axis=1)
# shape (150, 3)

fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection='3d')
ax.scatter(x1.numpy(), x2.numpy(), y.numpy(), c='red', marker='o')
ax.set_xlabel('x1'); ax.set_ylabel('x2'); ax.set_zlabel('y')
plt.title("3D Scatter: (x1, x2, y)")

dataset_shuffled = tf.random.shuffle(dataset, seed=22)
train_data, test_data = dataset_shuffled[:125], dataset_shuffled[125:]

x_train, y_train = train_data[:, 1:], train_data[:, 0]
x_test, y_test = test_data[:, 1:], test_data[:, 0]
# x_train shape: (125, 2)

class LinearRegression(tf.Module):
  def __init__(self):
    self.built = False

  @tf.function
  def __call__(self, x):
    if not self.built:
      rand_w = tf.random.uniform(shape=[x.shape[-1], 1])
      rand_b = tf.random.uniform(shape=[])
      self.w = tf.Variable(rand_w)
      self.b = tf.Variable(rand_b)
      self.built = True
    y = tf.add(tf.matmul(x, self.w), self.b)
    return tf.squeeze(y, axis=1)

def mse_loss(y_pred, y):
  return tf.reduce_mean(tf.square(y_pred - y))

batch_size = 32
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=x_train.shape[0]).batch(batch_size)

test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_dataset = test_dataset.shuffle(buffer_size=x_test.shape[0]).batch(batch_size)
# Batches: (32,2), (32,2), (32,2), (29,2)

epochs, learning_rate = 100, 0.01
train_losses, test_losses = [], []

for epoch in range(epochs):
  batch_losses_train = []
  for x_batch, y_batch in train_dataset:
    with tf.GradientTape() as tape:
      y_pred_batch = lin_reg(x_batch)
      batch_loss = mse_loss(y_pred_batch, y_batch)
    grads = tape.gradient(batch_loss, lin_reg.variables)
    for g, v in zip(grads, lin_reg.variables):
      v.assign_sub(learning_rate * g)
    batch_losses_train.append(batch_loss)

plt.plot(range(epochs), train_losses, label="Training loss")
plt.plot(range(epochs), test_losses, label="Testing loss")
plt.xlabel("Epoch")
plt.ylabel("Mean squared error loss")
plt.legend()
plt.title("MSE loss vs training iterations")