# Lesson 3: train/test split and honest evaluation
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error, r2_score

flats = pd.read_csv("flats.csv")
flats[["age_years", "floor"]] = flats[["age_years", "floor"]].fillna(flats[["age_years", "floor"]].median())
X = pd.get_dummies(flats.drop(columns="price_lakh"), columns=["city"], dtype=int)
y = flats["price_lakh"]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=1)
print("Train rows:", len(X_train), " Test rows:", len(X_test))

model = LinearRegression().fit(X_train, y_train)
pred = model.predict(X_test)
print(f"MAE on unseen flats: {mean_absolute_error(y_test, pred):.1f} lakh")
print(f"R2 on unseen flats:  {r2_score(y_test, pred):.3f}")

# baseline: always guess the average price. A model must beat this.
base = [y_train.mean()] * len(y_test)
print(f"Baseline MAE (always the average): {mean_absolute_error(y_test, base):.1f} lakh")

check = pd.DataFrame({"actual": y_test.values[:5], "predicted": pred[:5].round(1)})
print("\n" + check.to_string(index=False))