src/data.py

from __future__ import annotations

from pathlib import Path

import pandas as pd


PROJECT_ROOT = Path(__file__).resolve().parents[1]
DATA_DIR = PROJECT_ROOT / "data"
RAW_TRAIN_PATH = DATA_DIR / "raw" / "train.csv"
SAMPLE_TRAIN_PATH = DATA_DIR / "sample_house_prices.csv"


DEFAULT_FEATURES = [
    "OverallQual",
    "GrLivArea",
    "GarageCars",
    "TotalBsmtSF",
    "FullBath",
    "YearBuilt",
    "Neighborhood",
    "HouseStyle",
]
TARGET_COLUMN = "SalePrice"


def load_training_data(path: str | Path | None = None) -> pd.DataFrame:
    """Load Kaggle training data, falling back to the included sample dataset."""
    if path is not None:
        return pd.read_csv(path)

    if RAW_TRAIN_PATH.exists():
        return pd.read_csv(RAW_TRAIN_PATH)

    return pd.read_csv(SAMPLE_TRAIN_PATH)


def select_model_frame(frame: pd.DataFrame) -> pd.DataFrame:
    required = [*DEFAULT_FEATURES, TARGET_COLUMN]
    missing = [column for column in required if column not in frame.columns]
    if missing:
        raise ValueError(f"Training data is missing required columns: {', '.join(missing)}")
    return frame[required].copy()