"""Render the training loss curve as an ASCII chart — no plotting dependency.

Reads reports/loss-<config>.csv (written by src/train.py), prints a terminal
chart of train and val loss, and saves the same rendering to
reports/loss-<config>.txt as a committable artifact.

Usage:
    python scripts/plot_loss.py [--config v1]
"""

import argparse
import csv
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent

TRAIN_CH = "."
VAL_CH = "*"


def read_csv(path: Path):
    iters, train, val = [], [], []
    with open(path, newline="") as f:
        for row in csv.DictReader(f):
            iters.append(int(row["iter"]))
            train.append(float(row["train_loss"]))
            val.append(float(row["val_loss"]))
    return iters, train, val


def render(iters, series, width=64, height=18) -> str:
    """series: list of (name, values, char). Returns a multi-line chart string."""
    all_vals = [v for _, vals, _ in series for v in vals]
    lo, hi = min(all_vals), max(all_vals)
    if hi == lo:
        hi = lo + 1.0
    imin, imax = min(iters), max(iters)
    if imax == imin:
        imax = imin + 1

    def col(it: int) -> int:
        return round((it - imin) / (imax - imin) * (width - 1))

    def row(v: float) -> int:
        return round((hi - v) / (hi - lo) * (height - 1))  # row 0 = top = hi

    grid = [[" "] * width for _ in range(height)]
    for _, vals, ch in series:
        for it, v in zip(iters, vals):
            grid[row(v)][col(it)] = ch

    lines = []
    for r in range(height):
        yval = hi - r / (height - 1) * (hi - lo)
        lines.append(f"{yval:6.3f} |" + "".join(grid[r]))
    lines.append(" " * 7 + "+" + "-" * width)
    axis = f"{imin:<{width // 2}}{imax:>{width - width // 2}}"
    lines.append(" " * 8 + axis)
    legend = f"  legend: '{TRAIN_CH}' train   '{VAL_CH}' val   (x: iter, y: loss)"
    lines.append(legend)
    return "\n".join(lines)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="v1")
    args = ap.parse_args()

    csv_path = ROOT / "reports" / f"loss-{args.config}.csv"
    if not csv_path.exists():
        sys.exit(f"no loss log at {csv_path} — run training first")

    iters, train, val = read_csv(csv_path)
    if len(iters) < 2:
        sys.exit(f"only {len(iters)} data point(s) logged so far — need at least 2")

    chart = render(iters, [("train", train, TRAIN_CH), ("val", val, VAL_CH)])
    header = (f"loss curve — {args.config}  "
              f"(final train {train[-1]:.4f}, val {val[-1]:.4f} @ iter {iters[-1]})")
    out = header + "\n" + chart

    print(out)
    txt_path = ROOT / "reports" / f"loss-{args.config}.txt"
    txt_path.write_text(out + "\n")
    print(f"\nsaved {txt_path}")


if __name__ == "__main__":
    main()