import pandas as pd
import numpy as np
from sklearn.linear_model import Ridge, RidgeCV, LassoCV
from sklearn.preprocessing import StandardScaler
from ISLP import load_data

# === Setup ===
# Load and preprocess data
Hitters = load_data('Hitters').dropna()
Hitters = pd.get_dummies(Hitters, drop_first=True)
y = Hitters['Salary']
X = Hitters.drop(columns=['Salary'])

# Standardize predictors
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# === SLIDE 1: Ridge regression with fixed lambda ===
ridge_fixed = Ridge(alpha=100)
ridge_fixed.fit(X_scaled, y)
ridge_fixed_coeffs = ridge_fixed.coef_
ridge_fixed_preds = ridge_fixed.predict(X_scaled[:5])

# === SLIDE 2: Ridge regression with cross-validation to find best lambda ===
lambdas = 10**np.linspace(10, -2, 100) * 0.5  # Equivalent to R's lambda grid
ridge_cv = RidgeCV(alphas=lambdas, scoring='neg_mean_squared_error', cv=10)
ridge_cv.fit(X_scaled, y)
best_lambda_ridge = ridge_cv.alpha_
ridge_cv_coeffs = ridge_cv.coef_
ridge_cv_preds = ridge_cv.predict(X_scaled[:5])

# === SLIDE 3: Lasso regression with cross-validation ===
lasso_cv = LassoCV(cv=10, max_iter=10000)
lasso_cv.fit(X_scaled, y)
best_lambda_lasso = lasso_cv.alpha_
lasso_cv_coeffs = lasso_cv.coef_
lasso_cv_preds = lasso_cv.predict(X_scaled[:5])

# === Create summary DataFrame ===
summary = pd.DataFrame({
    'Model': ['Ridge (lambda=100)', 'RidgeCV (best lambda)', 'LassoCV (best lambda)'],
    'Best Lambda': [100, best_lambda_ridge, best_lambda_lasso],
    'Non-zero Coefficients': [
        np.sum(ridge_fixed_coeffs != 0),
        np.sum(ridge_cv_coeffs != 0),
        np.sum(lasso_cv_coeffs != 0)
    ]
})

print(summary)
