if __name__ == "__main__":
    rng = np.random.default_rng(0)

    # synthetic data: price = 3*sqft/1000 + 2*bedrooms - 1*age + noise
    n = 500
    sqft = rng.uniform(500, 3500, n)
    bedrooms = rng.integers(1, 6, n).astype(float)
    age = rng.uniform(0, 50, n)
    noise = rng.normal(0, 5, n)
    price = 3 * sqft / 1000 + 2 * bedrooms - 1 * age + 50 + noise

    X = np.column_stack([sqft / 1000, bedrooms, age])
    split = int(0.8 * n)
    X_train, X_test = X[:split], X[split:]
    y_train, y_test = price[:split], price[split:]

    w = fit_linear_normal_eq(X_train, y_train, ridge_lambda=0.1)
    preds = predict_linear(X_test, w)
    mse = np.mean((preds - y_test) ** 2)
    print("Learned weights (sqft/1000, bedrooms, age, bias):")
    print(np.round(w, 3))
    print(f"Test MSE: {mse:.2f}")

    # classification demo: is price above the median?
    y_class = (price > np.median(price)).astype(float)
    y_train_c, y_test_c = y_class[:split], y_class[split:]
    w_log = fit_logistic_gd(X_train, y_train_c, lr=0.5, epochs=2000)
    preds_c, probs = predict_logistic(X_test, w_log)
    acc = np.mean(preds_c == y_test_c)
    print(f"Logistic regression test accuracy: {acc:.3f}")
