import pandas as pd
import pickle
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression

# Load dataset
df = pd.read_csv("data.csv")

# Ensure required columns exist
REQUIRED_COLUMNS = {"sqft_living", "price"}
if not REQUIRED_COLUMNS.issubset(df.columns):
    raise ValueError(f"Dataset is missing required columns: {REQUIRED_COLUMNS - set(df.columns)}")

# Features & Target
X = df[["sqft_living"]]  # Input: Square footage
y = df["price"]           # Output: House price

# Split data into training and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Train model
model = LinearRegression()
model.fit(X_train, y_train)

# Save model to 'model.pkl' using built-in pickle
with open("model.pkl", "wb") as file:
    pickle.dump(model, file)

print("✅ Model training complete. Saved as 'model.pkl'.")