# Import warning filter to keep our notebook clean
import warnings
warnings.filterwarnings("ignore")

# Data setup: let us use the California Housing dataset for regression
from sklearn.datasets import fetch_california_housing
import pandas as pd

data = fetch_california_housing(as_frame=True)
df = data.frame
print("Shape of the data:", df.shape)
df.head()

# Let us check the target distribution
df["MedHouseVal"].describe()

# Split data: features and target
X = df.drop(columns=["MedHouseVal"])
y = df["MedHouseVal"]

# Split into train and test with 20 percent for testing
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
print("Train shape:", X_train.shape, "Test shape:", X_test.shape)

# Linear regression model
from sklearn.linear_model import LinearRegression
linreg = LinearRegression()
linreg.fit(X_train, y_train)
score_train = linreg.score(X_train, y_train)
score_test = linreg.score(X_test, y_test)
print("Train R^2:", round(score_train, 3), "Test R^2:", round(score_test, 3))

# Ridge regression with default strength
from sklearn.linear_model import Ridge
ridge = Ridge(alpha=1.0)
ridge.fit(X_train, y_train)
ridge_train = ridge.score(X_train, y_train)
ridge_test = ridge.score(X_test, y_test)
print("Ridge Train R^2:", round(ridge_train, 3), "Ridge Test R^2:", round(ridge_test, 3))

# Lasso regression with default strength
from sklearn.linear_model import Lasso
lasso = Lasso(alpha=1.0)
lasso.fit(X_train, y_train)
lasso_train = lasso.score(X_train, y_train)
lasso_test = lasso.score(X_test, y_test)
print("Lasso Train R^2:", round(lasso_train, 3), "Lasso Test R^2:", round(lasso_test, 3))

# ElasticNet regression
from sklearn.linear_model import ElasticNet
elastic = ElasticNet(alpha=1.0, l1_ratio=0.5)
elastic.fit(X_train, y_train)
elastic_train = elastic.score(X_train, y_train)
elastic_test = elastic.score(X_test, y_test)
print("ElasticNet Train R^2:", round(elastic_train, 3), "ElasticNet Test R^2:", round(elastic_test, 3))

# Compare the model coefficients side by side
coef_df = pd.DataFrame({
    "Feature": X.columns,
    "Linear": linreg.coef_,
    "Ridge": ridge.coef_,
    "Lasso": lasso.coef_,
    "ElasticNet": elastic.coef_
})
coef_df

# Practice: Which feature do you think is the most important for predicting house value?
answer = input("Type your guess from the list above: ")
print("Thanks for guessing! Let's look for it in the next cells.")

# Simple test: Try a smaller alpha with Ridge
ridge2 = Ridge(alpha=0.1)
ridge2.fit(X_train, y_train)
print("Ridge alpha=0.1 Test R^2:", round(ridge2.score(X_test, y_test), 3))

# Tip: use cross-validation for the best alpha
from sklearn.model_selection import cross_val_score
scores = cross_val_score(Ridge(alpha=1.0), X_train, y_train, cv=5)
print("Average CV R^2:", round(scores.mean(), 3))

# Mini Project step 1: Fit both Ridge and Lasso, then predict on test set
ridge_full = Ridge(alpha=1.0)
lasso_full = Lasso(alpha=1.0)
ridge_full.fit(X_train, y_train)
lasso_full.fit(X_train, y_train)
ridge_pred = ridge_full.predict(X_test)
lasso_pred = lasso_full.predict(X_test)

# Mini Project step 2: Calculate and compare mean squared error (MSE)
from sklearn.metrics import mean_squared_error
ridge_mse = mean_squared_error(y_test, ridge_pred)
lasso_mse = mean_squared_error(y_test, lasso_pred)
print("Ridge MSE:", round(ridge_mse, 3))
print("Lasso MSE:", round(lasso_mse, 3))