Fit a Least-Squares Line Without NumPy

from __future__ import annotations

import math


def fit_line(points: list[tuple[float, float]]) -> tuple[float, float, float]:
    count = len(points)
    if count < 2:
        raise ValueError("need at least two points")
    mean_x = sum(x for x, _ in points) / count
    mean_y = sum(y for _, y in points) / count
    sxx = sum((x - mean_x) ** 2 for x, _ in points)
    if sxx == 0:
        raise ValueError("x values must not all be equal")
    sxy = sum((x - mean_x) * (y - mean_y) for x, y in points)
    slope = sxy / sxx
    intercept = mean_y - slope * mean_x
    ss_total = sum((y - mean_y) ** 2 for _, y in points)
    ss_residual = sum((y - (slope * x + intercept)) ** 2 for x, y in points)
    r_squared = 1.0 - ss_residual / ss_total if ss_total else 1.0
    return slope, intercept, r_squared


def predict(slope: float, intercept: float, x: float) -> float:
    return slope * x + intercept


def main() -> None:
    points = [(1.0, 2.0), (2.0, 4.1), (3.0, 5.9), (4.0, 8.2), (5.0, 9.8)]
    slope, intercept, r_squared = fit_line(points)
    print(f"slope={slope:.3f} intercept={intercept:.3f} r2={r_squared:.4f}")
    print(f"prediction at x=6: {predict(slope, intercept, 6.0):.3f}")
    assert math.isclose(slope, 1.97, abs_tol=0.02)
    assert 0.99 < r_squared <= 1.0


if __name__ == "__main__":
    main()