scripts/run_ablation.py

"""Ablation study → one shareable HTML page.

Runs the REAL pipeline in several configurations over a set of addresses to isolate
what each signal contributes — especially the LiDAR-vs-national-scale question (how
much does dropping LiDAR cost?) and how much the neural segmenter adds over a dumb
color rule. Emits a matrix table (config × address, sqft + % vs the reference) and
a per-address strip of each config's measured-lawn overlay.

    python scripts/run_ablation.py --csv data/evals/addresses/cascade_vs_precascade.csv \
        --out data/outputs/ablation.html

Uses the byte-safe ablation flags (LAWN_FORCE_RGB_ONLY, LAWN_RGB_MODEL_OFF) — all off
in normal runs, so prod is unaffected.
"""
from __future__ import annotations

import argparse
import base64
import csv
import io
import os

import numpy as np
from dotenv import load_dotenv
from PIL import Image, ImageDraw

load_dotenv()
from lawn_estimator.pipeline import run  # noqa: E402

# All flags this study toggles — cleared before each config so runs don't leak.
ALL_FLAGS = ["SAM_RESTRICT", "ROW_TO_CURB", "GREEN_RECLAIM", "LAWN_CASCADE",
             "LAWN_FORCE_RGB_ONLY", "LAWN_RGB_MODEL_OFF", "LAWN_MODEL"]

# The ablation matrix: segmentation {cascade, pre-cascade, color-only} × LiDAR {on, off}.
CONFIGS = [
    ("Cascade + LiDAR", "reference — current prod pipeline",
     {"SAM_RESTRICT": "1", "ROW_TO_CURB": "1", "GREEN_RECLAIM": "1", "LAWN_CASCADE": "1"}),
    ("Cascade · NO LiDAR", "imagery-only — the national-scale question",
     {"SAM_RESTRICT": "1", "ROW_TO_CURB": "1", "GREEN_RECLAIM": "1", "LAWN_CASCADE": "1",
      "LAWN_FORCE_RGB_ONLY": "1"}),
    ("Pre-cascade + LiDAR", "incumbent + green reclaim (v0.2)",
     {"SAM_RESTRICT": "1", "ROW_TO_CURB": "1", "GREEN_RECLAIM": "1"}),
    ("Color baseline + LiDAR", "no neural model — HSV green rule × LiDAR",
     {"SAM_RESTRICT": "1", "ROW_TO_CURB": "1", "LAWN_RGB_MODEL_OFF": "1"}),
    ("Color baseline · NO LiDAR", "the floor: green pixels × pixel area",
     {"ROW_TO_CURB": "1", "LAWN_RGB_MODEL_OFF": "1", "LAWN_FORCE_RGB_ONLY": "1"}),
]


def b64(img: Image.Image, max_w: int = 460) -> str:
    if img.width > max_w:
        img = img.resize((max_w, round(img.height * max_w / img.width)), Image.LANCZOS)
    buf = io.BytesIO()
    img.convert("RGB").save(buf, "JPEG", quality=78)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()


def cell_viz(cap: dict, result: dict) -> str:
    """Compact overlay of the measured lawn for this config."""
    img = cap["image"].convert("RGB")
    est = cap.get("est")
    viz = getattr(est, "viz", {}) if est else {}
    gpx, gpy = viz.get("ground_px"), viz.get("ground_py")
    if gpx is not None:  # LiDAR path — dot the ground points
        d = ImageDraw.Draw(img)
        is_lawn = viz.get("lawn_mask")
        for k in range(0, len(gpx), 2):  # subsample for the thumbnail
            col = (50, 220, 50) if (is_lawn is not None and is_lawn[k]) else (235, 60, 60)
            d.ellipse([gpx[k] - 2, gpy[k] - 2, gpx[k] + 2, gpy[k] + 2], fill=col)
    else:  # RGB-only — fill the lawn mask
        fill = viz.get("lawn_fill", cap.get("lawn_area_in_parcel"))
        if fill is not None:
            arr = np.asarray(img).astype(np.float32)
            arr[fill] = arr[fill] * 0.45 + np.array([50, 220, 50], np.float32) * 0.55
            img = Image.fromarray(arr.astype(np.uint8))
    return b64(img)


def apply_env(flags: dict) -> None:
    for f in ALL_FLAGS:
        os.environ.pop(f, None)
    for k, v in flags.items():
        os.environ[k] = v


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--imagery", default="google")
    args = ap.parse_args()

    with open(args.csv, newline="", encoding="utf-8") as f:
        addresses = [r[next(iter(r))].strip() for r in csv.DictReader(f) if r[next(iter(r))].strip()]

    # results[config_name][address] = (sqft, viz_b64)
    results: dict[str, dict[str, tuple]] = {}
    for cname, _desc, flags in CONFIGS:
        results[cname] = {}
        for addr in addresses:
            apply_env(flags)
            try:
                cap: dict = {}
                res = run(addr, imagery=args.imagery, capture=cap)
                results[cname][addr] = (res["lawn_sqft"], cell_viz(cap, res))
                print(f"[{cname}] {addr}: {res['lawn_sqft']:,.0f}", flush=True)
            except Exception as exc:
                results[cname][addr] = (None, "")
                print(f"[{cname}] {addr}: FAILED {exc}", flush=True)

    ref = CONFIGS[0][0]
    short = [a.split(",")[0] for a in addresses]

    # matrix table
    head = "".join(f"<th>{s}</th>" for s in short)
    rows = ""
    for cname, desc, _ in CONFIGS:
        cells = ""
        for addr in addresses:
            sqft, _ = results[cname][addr]
            refv = results[ref][addr][0]
            if sqft is None:
                cells += '<td class="na">—</td>'
            else:
                pct = (sqft - refv) / refv * 100 if refv else 0
                cls = "" if cname == ref else ("hi" if pct > 4 else "lo" if pct < -4 else "ok")
                delta = "" if cname == ref else f'<span class="d {cls}">{pct:+.0f}%</span>'
                cells += f'<td>{sqft:,.0f}{delta}</td>'
        rows += f'<tr><th class="cfg">{cname}<span class="ds">{desc}</span></th>{cells}</tr>'

    # per-address viz strips
    strips = ""
    for addr, s in zip(addresses, short, strict=False):
        cardset = ""
        for cname, _d, _f in CONFIGS:
            sqft, viz = results[cname][addr]
            lbl = f"{sqft:,.0f}" if sqft is not None else "—"
            cardset += (f'<figure><img src="{viz}" alt="{cname}">'
                        f'<figcaption>{cname}<b>{lbl}</b></figcaption></figure>')
        strips += f'<section class="strip"><h3>{s}</h3><div class="cards">{cardset}</div></section>'

    html = PAGE.replace("<!--HEAD-->", head).replace("<!--ROWS-->", rows).replace(
        "<!--STRIPS-->", strips).replace("{ref}", ref).replace("{n}", str(len(addresses)))
    with open(args.out, "w", encoding="utf-8") as f:
        f.write(html)
    print(f"\nablation page: {args.out} ({os.path.getsize(args.out)/1e6:.1f} MB)")


PAGE = """<title>Ablation study — what each signal contributes</title>
<meta name="viewport" content="width=device-width, initial-scale=1">
<style>
:root{--bg:#0f130d;--card:#181d14;--ink:#e6ebe0;--soft:#9fb090;--line:#2c3626;--acc:#5db85d}
*{box-sizing:border-box}
body{margin:0;background:var(--bg);color:var(--ink);font:16px/1.5 "Avenir Next",system-ui,sans-serif}
main{max-width:1040px;margin:0 auto;padding:24px 16px 64px}
h1{font-size:clamp(22px,4vw,30px);margin:.2em 0}
.lead{color:var(--soft);max-width:70ch}
.tblwrap{overflow-x:auto;margin:18px 0}
table{border-collapse:collapse;width:100%;font-variant-numeric:tabular-nums}
th,td{border:1px solid var(--line);padding:9px 11px;text-align:right}
thead th{background:#20281a;color:var(--soft);font-weight:600}
th.cfg{text-align:left;background:#20281a}
th.cfg .ds{display:block;color:var(--soft);font-weight:400;font-size:12px}
tr:first-child td{font-weight:600}
.d{display:inline-block;margin-left:6px;font-size:12px;padding:1px 5px;border-radius:5px}
.d.hi{background:#2e7d3222;color:#7ec97e}.d.lo{background:#c6282822;color:#e88}.d.ok{color:var(--soft)}
.na{color:var(--soft)}
.strip{margin:26px 0}
.strip h3{margin:0 0 10px;border-left:3px solid var(--acc);padding-left:8px}
.cards{display:grid;grid-template-columns:repeat(auto-fill,minmax(200px,1fr));gap:10px}
figure{margin:0;background:var(--card);border:1px solid var(--line);border-radius:9px;overflow:hidden}
figure img{width:100%;display:block}
figcaption{font-size:12.5px;padding:6px 8px;display:flex;justify-content:space-between;gap:6px}
figcaption b{color:var(--acc)}
.key{color:var(--soft);font-size:13.5px;margin-top:8px}
</style>
<main>
<h1>Ablation study</h1>
<p class="lead">Each configuration is the REAL pipeline with one signal changed, over {n}
addresses. Reference = <b>{ref}</b>; the % next to each number is the change vs that.
Green dots (below) = LiDAR points counted as lawn, red = removed; solid green = an
imagery-only lawn fill (no LiDAR).</p>
<div class="tblwrap"><table>
<thead><tr><th class="cfg">configuration</th><!--HEAD--></tr></thead>
<tbody><!--ROWS--></tbody>
</table></div>
<p class="key">Read the "· NO LiDAR" rows against the reference: that gap is what dropping
LiDAR costs (the national-scale question). The "Color baseline" rows show what the neural
segmenter adds over a plain green rule.</p>
<!--STRIPS-->
</main>"""


if __name__ == "__main__":
    main()