# Module 2 Regression Project Template
# Module 2 - Practice regression project
# Dataset: /data/wage_sample.csv

# Module 2 Regression Project Template
#
# Use this notebook as a student project starter. Change `y_name` and `x_name` if you add a different dataset later. For now, the template uses `wage_sample.csv`.

# %% Cell 2
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

df = pd.read_csv("wage_sample.csv")

# Student choice: edit these two names for a different project.
y_name = "wage"
x_name = "education"

print("Project question:")
print(f"How is {y_name} associated with {x_name} in this sample?")
print("\nAvailable columns:", list(df.columns))
print(df[[y_name, x_name]].describe())

# %% Cell 3
plt.scatter(df[x_name], df[y_name])
plt.xlabel(x_name)
plt.ylabel(y_name)
plt.title(f"{y_name} and {x_name}")
plt.show()

# %% Cell 4
x = df[x_name]
y = df[y_name]

slope = ((x - x.mean()) * (y - y.mean())).sum() / ((x - x.mean()) ** 2).sum()
intercept = y.mean() - slope * x.mean()
fitted = intercept + slope * x
residuals = y - fitted
r_squared = 1 - (residuals ** 2).sum() / ((y - y.mean()) ** 2).sum()

print("Estimated regression")
print(f"{y_name}_hat = {intercept:.2f} + {slope:.2f} * {x_name}")
print("R-squared:", round(r_squared, 3))

# %% Cell 5
x_grid = np.linspace(x.min(), x.max(), 50)
line = intercept + slope * x_grid

plt.scatter(x, y, label="Actual data")
plt.plot(x_grid, line, label="Fitted line")
plt.xlabel(x_name)
plt.ylabel(y_name)
plt.title("Final project graph")
plt.legend()
plt.show()

# %% Cell 6
interpretation = f"In this sample, one more unit of {x_name} is associated with {slope:.2f} more units of predicted {y_name}."
fit_sentence = f"The R-squared is {r_squared:.3f}, so the model explains about {100*r_squared:.1f}% of the sample variation in {y_name}."
caution = "This is a simple association. A causal claim would require stronger assumptions or a stronger research design."

print(interpretation)
print(fit_sentence)
print(caution)

# Student submission checklist
#
# - State the research question.
# - Identify y and x.
# - Include one scatter plot and one fitted line.
# - Report the fitted equation and R-squared.
# - Interpret the slope in units.
# - Add one caution about causation or omitted variables.