ahmet / ft_linear_regression

import matplotlib.pyplot as plt
import pandas as pd

DATA_PATH = "data.csv"
THETAS_PATH = "thetas.csv"


def estimate_price(theta0, theta1, mileage):
    """estimatePrice(mileage) = theta0 + (theta1 * mileage)"""
    return theta0 + (theta1 * mileage)


def cal_loss(theta0, theta1, data):
    """Mean squared error of the current hypothesis."""
    return ((data.price - estimate_price(theta0, theta1, data.km)) ** 2).mean()


def cal_gradient(theta0, theta1, data, learning_rate):
    """One gradient descent step, updating both thetas simultaneously."""
    m = len(data)
    error = estimate_price(theta0, theta1, data.km) - data.price
    tmp0 = learning_rate * error.sum() / m
    tmp1 = learning_rate * (error * data.km).sum() / m
    return theta0 - tmp0, theta1 - tmp1


def train(data, learning_rate=0.1, iterations=1000, verbose=True):
    """Both thetas start at 0, as required before training."""
    theta0, theta1 = 0.0, 0.0
    for i in range(iterations):
        theta0, theta1 = cal_gradient(theta0, theta1, data, learning_rate)
        if verbose and i % 100 == 0:
            print(f"  iter {i:>4} | loss = {cal_loss(theta0, theta1, data):.4f}")
    return theta0, theta1


def precision(theta0, theta1, data):
    """R2, RMSE and MAE on the original (unscaled) data."""
    predicted = estimate_price(theta0, theta1, data.km)
    residual = data.price - predicted
    ss_res = (residual**2).sum()
    ss_tot = ((data.price - data.price.mean()) ** 2).sum()
    r2 = 1 - ss_res / ss_tot
    rmse = ((residual**2).mean()) ** 0.5
    mae = residual.abs().mean()
    return r2, rmse, mae


def save_thetas(theta0, theta1):
    thetas = pd.DataFrame({"theta0": [theta0], "theta1": [theta1]})
    thetas.to_csv(THETAS_PATH, index=False)


def graph(data, theta0, theta1):
    plt.scatter(data.km, data.price, label="data")
    x_line = [data.km.min(), data.km.max()]
    y_line = [estimate_price(theta0, theta1, x) for x in x_line]
    plt.plot(x_line, y_line, color="red", label="linear regression")
    plt.title("ft_linear_regression")
    plt.xlabel("Mileage (km)")
    plt.ylabel("Price")
    plt.legend()
    plt.show()


def main():
    try:
        data = pd.read_csv(DATA_PATH)
    except FileNotFoundError:
        print(f"Dataset not found: {DATA_PATH}")
        return

    if "km" not in data.columns or "price" not in data.columns:
        print("Dataset must contain 'km' and 'price' columns.")
        return

    data = data.dropna()
    if len(data) < 2:
        print("Not enough data to train on.")
        return

    km_mean = data.km.mean()
    km_std = data.km.std()
    if km_std == 0:
        print("All mileages are identical, cannot fit a line.")
        return

    scaled = data.copy()
    scaled.km = (data.km - km_mean) / km_std

    print("Training...")
    theta0, theta1 = train(scaled)

    theta1 = theta1 / km_std
    theta0 = theta0 - theta1 * km_mean

    r2, rmse, mae = precision(theta0, theta1, data)
    print(f"\ntheta0 = {theta0:.6f}")
    print(f"theta1 = {theta1:.6f}")
    print(f"price  = {theta0:.2f} + ({theta1:.6f} * km)")
    print(f"\nR2     = {r2:.4f}")
    print(f"RMSE   = {rmse:.2f}")
    print(f"MAE    = {mae:.2f}")

    save_thetas(theta0, theta1)
    print(f"\nThetas saved to {THETAS_PATH}")

    graph(data, theta0, theta1)


if __name__ == "__main__":
    main()