# Load in the diabetes dataset from SKLearn, only instance of SKLearn to be used in this notebook
X, y = load_diabetes(return_X_y=True, as_frame=True, scaled=False)
diabetes_data = X
diabetes_data['target'] = y
display(diabetes_data)

import seaborn as sns

sns.relplot(data = diabetes_data, x='bmi', y='target')

# Create a figure with two subplots side by side
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    
# Plot training data
ax1.scatter(bmi_train, y_train, color='blue', label='Training Data')
# Compute regression line for training data
line_x = np.array([bmi_train.min(), bmi_train.max()])
line_y = m * line_x + b
ax1.plot(line_x, line_y, color='orange', label='Regression Line')
    
ax1.set_title('Training Data')
ax1.set_xlabel('Feature [bmi]')
ax1.set_ylabel('Target (disease progression)')
ax1.legend()
    
# Plot test data
ax2.scatter(bmi_test, y_test, color='blue', label='Test Data')
# Compute regression line for test data
line_x = np.array([bmi_test.min(), bmi_test.max()])
line_y = m * line_x + b
ax2.plot(line_x, line_y, color='orange', label='Regression Line')
    
ax2.set_title('Test Data')
ax2.set_xlabel('Feature [bmi]')
ax2.set_ylabel('Target (disease progression)')
ax2.legend()
    
# Adjust layout and show plot
plt.tight_layout()
plt.show()

# TODO add in and understand the equations used
# TODO Add in SE, t, and p for slope and intercept. (The) Math is currently too complicated for me.
def regression_summary(X, y, m, b):
    '''
    Parameters:
    ----------

    Returns:
    --------
    '''
    n = len(X)
    preds = predict(X, m, b)
    residuals = [y_n - y_hat for y_n, y_hat in zip(y, preds)]

    # R-squared
    ss_total = sum((y_n - y.mean())**2 for y_n in y)
    ss_residual = sum(r**2 for r in residuals)
    r_squared = 1 - (ss_residual / ss_total)

    # Adjusted R-squared
    k = 1 #One predictor
    adjusted_r_squared = 1-(1-r_squared)*(n-1)/(n-k-1)
    
    # MSE and RMSE
    mse = ss_residual / n
    rmse = mse**.5
    
    # Print Summary
    print(f"Intercept (b): {b:.4f}")
    print(f"Slope (m): {m:.4f}")
    print(f"MSE: {mse:.4f}, RMSE: {rmse:.4f}")
    print(f"R-squared: {r_squared:.4f}, Adjusted R-squared: {adjusted_r_squared:.4f}\n")
    
regression_summary(bmi_train, y_train, m, b)
regression_summary(bmi_test, y_test, m, b)

# Create a simple dataset
X = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])
y = np.array([3, 7, 11, 15])  # Linearly related to X

# Instantiate and fit your model
model = MultipleLinearRegression(learning_rate=0.01, iterations=1000)
model.fit(X, y)

# Check if weights and bias are reasonable
print(f"Weights: {model.weights}, Bias: {model.bias}")

# Generate some sample data
np.random.seed(42)
X = np.random.randn(100, 3)  # 100 samples, 3 features

#df = pd.DataFrame(X)
#df.to_csv('X.csv')

true_weights = np.array([2.0, -3.5, 1.0])
true_bias = 4.0
y = np.dot(X, true_weights) + true_bias + np.random.randn(100) * 0.5

#df = pd.DataFrame(y)
#df.to_csv('y.csv')
    
# Train the model
model = MultipleLinearRegression(learning_rate=0.01, iterations=1000)
model.fit(X, y)
    
# Check learned parameters
print(f"True weights: {true_weights}, True bias: {true_bias}")
print(f"Learned weights: {model.weights}, Learned bias: {model.bias}")
    
# Plot cost history
model.plot_cost_history()

# Using the Multiple Linear Regression Class
X, y = load_diabetes(return_X_y=True, as_frame=False, scaled=True)
# Create the training and test data
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)

#df = pd.DataFrame(X_train)
#df.to_csv('X_train.csv')

# Standardize, very important!
#X_train = (X_train - X_train.mean(axis=0)) / X_train.std(axis=0)
#X_test = (X_test - X_test.mean(axis=0)) / X_test.std(axis=0)
#y_standard = (y - y.mean(axis=0)) / y.std(axis=0)

#df = pd.DataFrame(y_train)
#df.to_csv('y_train_standard.csv')

#true_weights = np.array([2.0, -3.5, 1.0, .5, 1.0, 1.5, 2.0, -2.0, -1.5, -1.0])
#true_bias = 4.0
#print(X_train.shape)
#y_train = np.dot(X_train, true_weights) + true_bias + np.random.randn(X_train.shape[0]) * 0.5
#print(y_train.shape)

# SOMETHING IS WRONG WITH MY VALUES AND I DON'T KNOW WHAT
# Addendum: I wasn't standardizing my test features. Oops!
model = MultipleLinearRegression(iterations=1000)
model.fit(X_train, y_train)

print(f'Learned weights: {model.weights}, Learned bias: {model.bias}')

# Plot cost history
model.plot_cost_history()

# Let's compare our coef to the ones generated by sklearn!
from sklearn.linear_model import LinearRegression

# Using the Multiple Linear Regression Class
X, y = load_diabetes(return_X_y=True, as_frame=True, scaled=True)
# Create the training and test data
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)

sklearn_model = LinearRegression()
sklearn_model.fit(X_train, y_train)

coefficients = sklearn_model.coef_
intercept = sklearn_model.intercept_

np.set_printoptions(suppress=True)
print(f"Coefficients: {coefficients}") 
print(f"Intercept: {intercept}")

print(model.predict(X_test))
print(sklearn_model.predict(X)[1])

# How does my model compare to SKLearn's? We're going to use MSE and RMSE to compare.
from sklearn.metrics import mean_squared_error

preds_sklearn = sklearn_model.predict(X_test)
mse_sklearn = mean_squared_error(y_test, preds_sklearn)

preds_bespoke = model.predict(X_test)
residuals_bespoke = [y_n - y_hat for y_n, y_hat in zip(y_test, preds_bespoke)]
ss_total = sum((y_n - y_test.mean())**2 for y_n in y_test)
ss_residual_bespoke = sum(r**2 for r in residuals_bespoke)
mse_bespoke = ss_residual_bespoke / y_test.shape[0]

print(mse_sklearn**.5)
print(mse_bespoke**.5)