plot.py

#!/usr/bin/env python3
"""Bonus: visualise the data and the regression line.

Uses matplotlib when it is installed (saves regression.png, or pops a window
with --show); otherwise falls back to a hand-written, dependency-free SVG. Use
--svg to force the SVG path. Either way it also prints a quick ASCII scatter to
the terminal for instant feedback.

Usage:
    python3 plot.py [data.csv] [--svg] [--show]
"""

from __future__ import annotations

import sys

from linear_regression import THETA_FILE, Model, estimate_price, load_dataset, load_model

WIDTH = 820
HEIGHT = 600
PAD = 70


def _scale(v: float, lo: float, hi: float, a: float, b: float) -> float:
    if hi == lo:
        return (a + b) / 2
    return a + (v - lo) / (hi - lo) * (b - a)


def write_svg(xs: list[float], ys: list[float], model: Model, path: str = "regression.svg") -> None:
    """Render the scatter plus the fitted line as a standalone SVG file."""
    xmin, xmax = min(xs), max(xs)
    line = [
        estimate_price(xmin, model.theta0, model.theta1),
        estimate_price(xmax, model.theta0, model.theta1),
    ]
    ymin = min(min(ys), *line)
    ymax = max(max(ys), *line)

    def px(x: float) -> float:
        return _scale(x, xmin, xmax, PAD, WIDTH - PAD)

    def py(y: float) -> float:
        return _scale(y, ymin, ymax, HEIGHT - PAD, PAD)

    out: list[str] = []
    out.append(
        f'<svg xmlns="http://www.w3.org/2000/svg" width="{WIDTH}" '
        f'height="{HEIGHT}" font-family="sans-serif">'
    )
    out.append(f'<rect width="{WIDTH}" height="{HEIGHT}" fill="white"/>')
    out.append(
        f'<line x1="{PAD}" y1="{HEIGHT - PAD}" x2="{WIDTH - PAD}" '
        f'y2="{HEIGHT - PAD}" stroke="#333"/>'
    )
    out.append(f'<line x1="{PAD}" y1="{PAD}" x2="{PAD}" y2="{HEIGHT - PAD}" stroke="#333"/>')
    for i in range(6):
        gx = xmin + (xmax - xmin) * i / 5
        gy = ymin + (ymax - ymin) * i / 5
        out.append(
            f'<text x="{px(gx):.0f}" y="{HEIGHT - PAD + 18}" font-size="11" '
            f'text-anchor="middle">{gx:,.0f}</text>'
        )
        out.append(
            f'<text x="{PAD - 8}" y="{py(gy) + 4:.0f}" font-size="11" '
            f'text-anchor="end">{gy:,.0f}</text>'
        )
    out.append(
        f'<text x="{WIDTH / 2:.0f}" y="{HEIGHT - 22}" font-size="13" '
        f'text-anchor="middle">mileage (km)</text>'
    )
    out.append(
        f'<line x1="{px(xmin):.1f}" y1="{py(line[0]):.1f}" x2="{px(xmax):.1f}" '
        f'y2="{py(line[1]):.1f}" stroke="#e4572e" stroke-width="2"/>'
    )
    for x, y in zip(xs, ys, strict=True):
        out.append(
            f'<circle cx="{px(x):.1f}" cy="{py(y):.1f}" r="4" fill="#2e86de" fill-opacity="0.85"/>'
        )
    out.append("</svg>")
    with open(path, "w") as f:
        f.write("\n".join(out))


def ascii_plot(
    xs: list[float], ys: list[float], model: Model, width: int = 64, height: int = 18
) -> None:
    """Print a compact terminal scatter (``o``) with the fitted line (``.``)."""
    xmin, xmax = min(xs), max(xs)
    ymin, ymax = min(ys), max(ys)
    grid = [[" "] * width for _ in range(height)]

    def col(x: float) -> int:
        return min(width - 1, max(0, round((x - xmin) / (xmax - xmin) * (width - 1))))

    def row(y: float) -> int:
        return min(height - 1, max(0, round((ymax - y) / (ymax - ymin) * (height - 1))))

    for c in range(width):
        x = xmin + (xmax - xmin) * c / (width - 1)
        y = estimate_price(x, model.theta0, model.theta1)
        if ymin <= y <= ymax:
            grid[row(y)][c] = "."
    for x, y in zip(xs, ys, strict=True):
        grid[row(y)][col(x)] = "o"
    print(f"price {ymax:,.0f}")
    for r in grid:
        print("  |" + "".join(r))
    print("  +" + "-" * width)
    print(f"   {xmin:,.0f}{' ' * (width - 14)}{xmax:,.0f} km")


def has_matplotlib() -> bool:
    """True if matplotlib is importable (optional, richer rendering)."""
    import importlib.util

    return importlib.util.find_spec("matplotlib") is not None


def write_matplotlib(
    xs: list[float],
    ys: list[float],
    model: Model,
    path: str = "regression.png",
    show: bool = False,
) -> None:
    """Render with matplotlib: save a PNG (and optionally pop a window)."""
    import matplotlib

    if not show:
        matplotlib.use("Agg")  # headless: only write the file
    import matplotlib.pyplot as plt

    xmin, xmax = min(xs), max(xs)
    line_y = [
        estimate_price(xmin, model.theta0, model.theta1),
        estimate_price(xmax, model.theta0, model.theta1),
    ]
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.scatter(xs, ys, color="#2e86de", label="data", zorder=3)
    ax.plot([xmin, xmax], line_y, color="#e4572e", linewidth=2, label="regression")
    ax.set_xlabel("mileage (km)")
    ax.set_ylabel("price")
    ax.set_title("ft_linear_regression")
    ax.legend()
    ax.grid(visible=True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(path, dpi=120)
    if show:
        plt.show()
    plt.close(fig)


def main(argv: list[str] | None = None) -> int:
    argv = sys.argv if argv is None else argv
    flags = {a for a in argv[1:] if a.startswith("-")}
    positional = [a for a in argv[1:] if not a.startswith("-")]
    data = positional[0] if positional else "data.csv"
    force_svg = "--svg" in flags
    show = "--show" in flags
    try:
        xs, ys = load_dataset(data)
    except (OSError, ValueError) as exc:
        print(f"plot: cannot read {data}: {exc}", file=sys.stderr)
        return 1
    model = load_model(THETA_FILE)
    ascii_plot(xs, ys, model)
    if not force_svg and has_matplotlib():
        write_matplotlib(xs, ys, model, "regression.png", show=show)
        print("\nWrote regression.png (matplotlib).")
    else:
        if show and not force_svg:
            print("plot: matplotlib not installed; writing SVG instead.", file=sys.stderr)
        write_svg(xs, ys, model, "regression.svg")
        print("\nWrote regression.svg (open in a browser for the full graph).")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())