import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score

# Create the dataset
X = np.random.rand(100, 1)  # Independent variable
y = 3 * X**2 + np.random.randn(100, 1)  # Dependent variable (quadratic relationship)

# Normalize the data
X_normalized = (X - np.mean(X)) / np.std(X)
y_normalized = (y - np.mean(y)) / np.std(y)

# Split the data into train and test sets
X_train, X_test, y_train, y_test = train_test_split(X_normalized, y_normalized, test_size=0.2, random_state=42)

# Fit the linear regression model
linear_regression = LinearRegression()
linear_regression.fit(X_train, y_train)

# Fit the 4-degree polynomial regression model
poly_features = PolynomialFeatures(degree=4)
X_poly_train = poly_features.fit_transform(X_train)
X_poly_test = poly_features.transform(X_test)
poly_regression = LinearRegression()
poly_regression.fit(X_poly_train, y_train)

# Fit the 15-degree polynomial regression model
poly_features = PolynomialFeatures(degree=15)
X_poly_train = poly_features.fit_transform(X_train)
X_poly_test = poly_features.transform(X_test)
poly_regression_15 = LinearRegression()
poly_regression_15.fit(X_poly_train, y_train)

# Calculate R-squared for train and test sets
linear_train_r2 = linear_regression.score(X_train, y_train)
linear_test_r2 = linear_regression.score(X_test, y_test)

poly_train_r2 = poly_regression.score(X_poly_train, y_train)
poly_test_r2 = poly_regression.score(X_poly_test, y_test)

poly_15_train_r2 = poly_regression_15.score(X_poly_train, y_train)
poly_15_test_r2 = poly_regression_15.score(X_poly_test, y_test)

# Print the results
print("Linear Regression R-squared (Train):", linear_train_r2)
print("Linear Regression R-squared (Test):", linear_test_r2)

print("4-Degree Polynomial Regression R-squared (Train):", poly_train_r2)
print("4-Degree Polynomial Regression R-squared (Test):", poly_test_r2)

print("15-Degree Polynomial Regression R-squared (Train):", poly_15_train_r2)
print("15-Degree Polynomial Regression R-squared (Test):", poly_15_test_r2)