# Data setup
import warnings; warnings.filterwarnings("ignore")
import pandas as pd
import seaborn as sns
df = sns.load_dataset("mpg")
print("Dataset shape:", df.shape)
df.head()

# Pick two columns: weight and mpg
df = df[["weight", "mpg"]].dropna()
print("Clean shape:", df.shape)
df.head()

import matplotlib.pyplot as plt
plt.figure(figsize=(8,5))
plt.scatter(df["weight"], df["mpg"], alpha=0.7)
plt.xlabel("Car Weight")
plt.ylabel("Miles Per Gallon (mpg)")
plt.title("Car Weight vs. MPG")
plt.show()

# Fit the linear model
from sklearn.linear_model import LinearRegression
X = df[["weight"]]
y = df["mpg"]
model = LinearRegression()
model.fit(X, y)
slope = model.coef_[0]
intercept = model.intercept_
print("Slope:", slope)
print("Intercept:", intercept)

# Draw the regression line on the scatter plot
plt.figure(figsize=(8,5))
plt.scatter(df["weight"], df["mpg"], alpha=0.5)
plt.plot(df["weight"], model.predict(X), color="red", label="Regression Line")
plt.xlabel("Car Weight")
plt.ylabel("MPG")
plt.title("Regression Line: Weight vs MPG")
plt.legend()
plt.show()

# Predict mpg for a new car weight
test_weight = 2500
pred_mpg = model.predict([[test_weight]])[0]
print("Predicted mpg for 2500 lbs:", pred_mpg)

# Try your own value!
your_weight = float(input("Enter a car weight in pounds (for example, 3400): "))
your_pred = model.predict([[your_weight]])[0]
print("Predicted mpg:", your_pred)

# How well does our model fit?
r2 = model.score(X, y)
print("R-squared:", r2)

# Check for linearity with a residual plot
import numpy as np
residuals = y - model.predict(X)
plt.scatter(df["weight"], residuals, alpha=0.6)
plt.axhline(0, color="red", linestyle="--")
plt.xlabel("Car Weight")
plt.ylabel("Residual (Actual - Predicted)")
plt.title("Residual Plot")
plt.show()

# Try adding another variable: horsepower
X2 = df[["weight", "mpg"]].dropna()
if "horsepower" in sns.load_dataset("mpg").columns:
    df2_full = sns.load_dataset("mpg")
    df2 = df2_full[["weight", "horsepower", "mpg"]].dropna()
    X2 = df2[["weight", "horsepower"]]
    y2 = df2["mpg"]
    model2 = LinearRegression()
    model2.fit(X2, y2)
    r2_2 = model2.score(X2, y2)
    print("Linear regression with weight and horsepower. R-squared:", r2_2)
else:
    print("No horsepower column in dataset.")
    

# Troubleshooting
print("Check your data for missing or weird values.")
if df.isnull().sum().sum() > 0:
    print("Warning: Missing values found.")
if df.duplicated().sum() > 0:
    print("Warning: Duplicate rows found.")
    

# Challenge! Predict mpg for three new car weights.
for i in range(3):
    w = float(input(f"Test weight {i+1}: "))
    print(f"Predicted mpg for {w} lbs: {model.predict([[w]])[0]:.2f}")