# Experiment 2: Simple Linear Regression
# Requirements: pip install numpy pandas matplotlib seaborn scikit-learn

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score

# ==========================================
# STEP 1: Generate the Data
# ==========================================
# Synthetic data: Years of Experience vs Salary, with random noise.
np.random.seed(42)
n_samples = 100

experience = np.random.uniform(0, 15, n_samples)
noise = np.random.normal(0, 5000, n_samples)
salary = 30000 + 4500 * experience + noise

data = pd.DataFrame({'YearsExperience': experience, 'Salary': salary})
data.to_csv('simple_linear_regression_data.csv', index=False)

# Load the saved CSV
# CHANGE THIS PATH TO YOUR DATASET (and the column names below)
df = pd.read_csv('simple_linear_regression_data.csv')
print("--- Dataset Head ---")
print(df.head())

X = df[['YearsExperience']]   # Feature matrix (2D)
y = df['Salary']              # Target vector

# ==========================================
# STEP 2: Splitting Data into Training and Testing Sets (80/20)
# ==========================================
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)
print(f"\nTraining Set Shape: {X_train.shape}")
print(f"Testing Set Shape:  {X_test.shape}")

# ==========================================
# STEP 3: Training the Simple Linear Regression Model
# ==========================================
model = LinearRegression()
model.fit(X_train, y_train)

print(f"\nIntercept (beta_0): {model.intercept_:.2f}")
print(f"Slope     (beta_1): {model.coef_[0]:.2f}")

# ==========================================
# STEP 4: Making Predictions & Model Evaluation
# ==========================================
y_pred = model.predict(X_test)

mse = mean_squared_error(y_test, y_pred)
rmse = np.sqrt(mse)
r2 = r2_score(y_test, y_pred)

print("\n--- MODEL PERFORMANCE METRICS ---")
print(f"Mean Squared Error (MSE)       : {mse:.2f}")
print(f"Root Mean Squared Error (RMSE) : {rmse:.2f}")
print(f"R^2 Score                      : {r2:.4f}")

# ==========================================
# STEP 5: Visualizing Results
# ==========================================
fig, axes = plt.subplots(1, 2, figsize=(12, 5))

# Training set
sns.scatterplot(x=X_train['YearsExperience'], y=y_train, ax=axes[0], color='blue', label='Actual')
axes[0].plot(X_train['YearsExperience'], model.predict(X_train), color='red', label='Regression Line')
axes[0].set_title('Training Set: Salary vs Experience')
axes[0].set_xlabel('Years of Experience')
axes[0].set_ylabel('Salary')
axes[0].legend()

# Testing set
sns.scatterplot(x=X_test['YearsExperience'], y=y_test, ax=axes[1], color='green', label='Actual')
axes[1].plot(X_train['YearsExperience'], model.predict(X_train), color='red', label='Regression Line')
axes[1].set_title('Testing Set: Salary vs Experience')
axes[1].set_xlabel('Years of Experience')
axes[1].set_ylabel('Salary')
axes[1].legend()

plt.tight_layout()
plt.show()
