import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error

np.random.seed(1)
x_1 = np.absolute(np.random.randn(100, 1) * 10) #2D array with 100 rows, 1 column
x_2 = np.absolute(np.random.randn(100, 1) * 30)
#y = 2*x_1**2 + 3*x_1 + 2 + np.random.randn(100, 1)*20
y = 2*x_1**2 + 3*x_1 + 4*x_2 + np.random.randn(100, 1)*20


fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(10, 4)) #1 line of subplots, 2 rows of subplots
axes[0].scatter(x_1, y)
axes[1].scatter(x_2, y)
axes[0].set_title("x_1 plotted")
axes[1].set_title("x_2 plotted")
plt.show()

#reshape to turn the arrays into 1D arrays with 100 elements each
df = pd.DataFrame({"x_1":x_1.reshape(100,), "x_2":x_2.reshape(100,), "y":y.reshape(100,)}, index=range(0,100)) #index: assign unique identifiers for each df row

#X, y = df[["x_1", "x_2"]], df["y"] could also be written as below
X= df[["x_1", "x_2"]]
y = df["y"]

poly = PolynomialFeatures(degree=2, include_bias=False)
poly_features = poly.fit_transform(X) #transform original features X into polynomial features
X_train, X_test, y_train, y_test = train_test_split(poly_features, y, test_size=0.3, random_state=42)

poly_reg_model = LinearRegression() #create model
poly_reg_model.fit(X_train, y_train) #train model

poly_reg_y_predicted = poly_reg_model.predict(X_test) #make predictions on the test set

poly_reg_rmse = np.sqrt(mean_squared_error(y_test, poly_reg_y_predicted))
print(poly_reg_rmse)