# Complex demo: Load California Housing, split with train_test_split and KFold, regress, evaluate, visualize residuals
import pandas as pd
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split, KFold, cross_val_score
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import numpy as np

def advanced_housing_splitting():
    # Load and prep
    housing = fetch_california_housing()
    df = pd.DataFrame(housing.data, columns=housing.feature_names)
    df['Target'] = housing.target
    X = df.drop('Target', axis=1)
    y = df['Target']
    
    # Scale
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    # Train_test_split (no stratification for regression)
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.25, random_state=42)
    
    # Train and predict
    model = LinearRegression()
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    mse = mean_squared_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    print(f"MSE with train_test_split: {mse:.2f}")
    print(f"R² with train_test_split: {r2:.2f}")
    
    # KFold cross-validation
    kf = KFold(n_splits=5, shuffle=True, random_state=42)
    cv_scores = cross_val_score(model, X_scaled, y, cv=kf)
    print(f"Cross-Validation Scores: {cv_scores}")
    print(f"Mean CV Score: {np.mean(cv_scores):.2f}")
    
    # Visualize residuals
    residuals = y_test - y_pred
    plt.scatter(y_test, residuals, c='red', alpha=0.5)
    plt.axhline(0, color='black', linestyle='--')
    plt.xlabel('True Values')
    plt.ylabel('Residuals')
    plt.title('Residuals Plot for Housing Splitting')
    plt.annotate('Low Residual', xy=(y_test.mean(), 0), xytext=(y_test.mean() + 0.5, 0.5), 
                 arrowprops=dict(facecolor='black', shrink=0.05))
    plt.show()

advanced_housing_splitting()