zonelearn.py

"""Learn occupancy zones from radar target coordinates. No zones to draw.

Feed it a stream of (x, y, t) target positions and it reports where people
actually spend time — sofa, desk, chair — without anyone drawing a box.

    python3 zonelearn.py --csv coords.csv
    python3 zonelearn.py --ha-db ~/.homeassistant/home-assistant_v2.db   # SQLite recorder
    python3 zonelearn.py --ha-recorder sensor.living_target_1            # PostgreSQL recorder

Sensor-agnostic by design: anything that reports 2-D target coordinates works
(HLK-LD2450, LD6002B, Aqara FP2 exports, or your own). Coordinates are
normalized to millimetres before anything else happens.
"""

import argparse
import collections
import csv
import math
import os
import re
import sys

NUMERIC = re.compile(r"^-?\d+(\.\d+)?$")

# --------------------------------------------------------------------------- io


def from_csv(path, scale=1.0):
    """CSV with columns x,y[,t]. `scale` converts your unit into millimetres."""
    out = []
    with open(path) as f:
        for row in csv.DictReader(f):
            try:
                out.append((float(row["x"]) * scale, float(row["y"]) * scale))
            except (KeyError, ValueError):
                continue
    return out


def from_sqlite(db, entity_prefix=None, days=0):
    """Home Assistant's default recorder. No dependencies, no export step.

    `entity_prefix` is optional — with nothing given, the first pair of
    <something>_x / <something>_y target sensors in the database is used.
    """
    import sqlite3

    con = sqlite3.connect(f"file:{os.path.expanduser(db)}?mode=ro", uri=True)
    try:
        meta = con.execute("SELECT name FROM sqlite_master WHERE type='table'"
                           " AND name='states_meta'").fetchone()
        if entity_prefix is None:
            entity_prefix = _guess_prefix(con, meta)
            print(f"reading {entity_prefix}_x and {entity_prefix}_y", file=sys.stderr)

        since = (f"AND s.last_updated_ts > (strftime('%s','now') - {days * 86400})"
                 if days else "")
        if meta:  # recorder 2023.4 and later
            q = ("SELECT s.last_updated_ts, s.state FROM states s"
                 " JOIN states_meta m ON s.metadata_id = m.metadata_id"
                 f" WHERE m.entity_id = ? {since}")
        else:
            q = f"SELECT s.last_updated_ts, s.state FROM states s WHERE s.entity_id = ? {since}"

        series = []
        for suffix in ("_x", "_y"):
            got = {}
            for ts, state in con.execute(q, (entity_prefix + suffix,)):
                if ts is not None and state is not None and NUMERIC.match(str(state)):
                    got[int(float(ts))] = float(state)   # one reading a second is plenty
            series.append(got)
    finally:
        con.close()

    xs, ys = series
    return [(x, ys[t]) for t, x in xs.items() if t in ys]


def _guess_prefix(con, has_meta):
    """Find <prefix>_x / <prefix>_y target sensors so nobody has to look them up."""
    table = "states_meta" if has_meta else "states"
    ids = {r[0] for r in con.execute(
        f"SELECT DISTINCT entity_id FROM {table} WHERE entity_id LIKE '%target%_x'")}
    for i in sorted(ids):
        prefix = i[:-2]
        if con.execute(f"SELECT 1 FROM {table} WHERE entity_id = ? LIMIT 1",
                       (prefix + "_y",)).fetchone():
            return prefix
    sys.exit("no <prefix>_x / <prefix>_y target sensors found — pass --entity")


def from_recorder(entity_prefix, days=0, dsn=None):
    """Home Assistant PostgreSQL recorder: <prefix>_x and <prefix>_y."""
    try:
        import psycopg2
    except ImportError:
        sys.exit("psycopg2 not installed. Use --csv, or: pip install psycopg2-binary")
    since = f"AND s.last_updated_ts > extract(epoch from now()) - {days * 86400}" if days else ""
    q = f"""
    WITH x AS (SELECT floor(s.last_updated_ts) t, s.state::float v FROM states s
               JOIN states_meta sm ON s.metadata_id = sm.metadata_id
               WHERE sm.entity_id = %s AND s.state ~ '^-?[0-9.]+$' {since}),
         y AS (SELECT floor(s.last_updated_ts) t, s.state::float v FROM states s
               JOIN states_meta sm ON s.metadata_id = sm.metadata_id
               WHERE sm.entity_id = %s AND s.state ~ '^-?[0-9.]+$' {since})
    SELECT x.v, y.v FROM x JOIN y ON x.t = y.t;"""
    with psycopg2.connect(dsn) as conn, conn.cursor() as cur:
        cur.execute(q, (f"{entity_prefix}_x", f"{entity_prefix}_y"))
        return [(float(a), float(b)) for a, b in cur.fetchall()]


# ---------------------------------------------------------------------- geometry


def drop_no_target(points):
    """(0, 0) is what a radar reports when it sees nobody, not a place people sit.

    A CSV exported by hand is usually already clean. A recorder read straight
    from the database is not — empty rooms are most of the day.
    """
    kept = [(x, y) for x, y in points if x or y]
    return kept, len(points) - len(kept)


def clip_to_room(points, x_min=None, x_max=None, y_max=None):
    """Drop points outside the room.

    Radar goes through drywall, so a sensor near a wall reports the neighbours.
    Measure three distances from the sensor and everything else stays automatic:
    how far the wall is to your left (negative x), to your right, and straight
    ahead. Leave any of them out and that side is not clipped.
    """
    if x_min is None and x_max is None and y_max is None:
        return points, 0
    kept, dropped = [], 0
    for x, y in points:
        if (x_min is not None and x < x_min) or (x_max is not None and x > x_max) \
           or (y_max is not None and y > y_max):
            dropped += 1
            continue
        kept.append((x, y))
    return kept, dropped


def find_zones(points, cell=250, frac=0.10):
    """Grid density -> threshold -> 4-neighbour connected components."""
    hist = collections.Counter((int(math.floor(x / cell)), int(math.floor(y / cell)))
                               for x, y in points)
    if not hist:
        return [], hist
    peak = max(hist.values())
    keep = {c for c, n in hist.items() if n >= peak * frac}
    seen, zones = set(), []
    for start in keep:
        if start in seen:
            continue
        comp, stack = [], [start]
        while stack:
            c = stack.pop()
            if c in seen or c not in keep:
                continue
            seen.add(c)
            comp.append(c)
            cx, cy = c
            stack += [(cx + 1, cy), (cx - 1, cy), (cx, cy + 1), (cx, cy - 1)]
        mass = sum(hist[c] for c in comp)
        zones.append({
            "mass": mass,
            "cx": sum(hist[c] * (c[0] + .5) * cell for c in comp) / mass,
            "cy": sum(hist[c] * (c[1] + .5) * cell for c in comp) / mass,
            "cells": len(comp),
            "x0": min(c[0] for c in comp) * cell, "x1": (max(c[0] for c in comp) + 1) * cell,
            "y0": min(c[1] for c in comp) * cell, "y1": (max(c[1] for c in comp) + 1) * cell,
        })
    return sorted(zones, key=lambda z: -z["mass"]), hist


def heatmap(hist, cell):
    """ASCII density map. y increases upward, sensor at the bottom."""
    if not hist:
        return ""
    xs = [c[0] for c in hist]
    ys = [c[1] for c in hist]
    peak = max(hist.values())
    rows = []
    for cy in range(max(ys), min(ys) - 1, -1):
        row = "".join(
            "#" if (r := hist.get((cx, cy), 0) / peak) > .25 else
            "+" if r > .10 else "-" if r > .03 else "." if r > .005 else " "
            for cx in range(min(xs), max(xs) + 1))
        rows.append(f"{cy * cell / 1000:5.1f}m {row}")
    return "\n".join(rows)


# ------------------------------------------------------------------------ emit

ESPHOME_ZONES = 3          # the ld2450 component accepts three, and no more


def as_esphome(zones):
    """Zone entities for an ESPHome ld2450 block, plus the numbers to put in them."""
    out = ["# --- ESPHome: declare the zones under your ld2450: block ---"]
    for i, z in enumerate(zones[:ESPHOME_ZONES], 1):
        out.append(f"  zone_{i}:")
        for k, v in (("x1", z["x0"]), ("y1", z["y0"]), ("x2", z["x1"]), ("y2", z["y1"])):
            out.append(f"    {k}:\n      name: Zone-{i} {k.upper()}   # {v:.0f} mm")
    return "\n".join(out)


def as_ha_actions(zones, prefix="ld2450"):
    """Paste into Developer Tools > Actions to set the zones without clicking."""
    out = ["# --- Home Assistant: Developer Tools > Actions, YAML mode ---"]
    for i, z in enumerate(zones[:ESPHOME_ZONES], 1):
        for k, v in (("x1", z["x0"]), ("y1", z["y0"]), ("x2", z["x1"]), ("y2", z["y1"])):
            out.append(f"- action: number.set_value\n  target:\n"
                       f"    entity_id: number.{prefix}_zone_{i}_{k}\n"
                       f"  data:\n    value: {v:.0f}")
    return "\n".join(out)


# -------------------------------------------------------------------------- cli


def main():
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    src = p.add_mutually_exclusive_group(required=True)
    src.add_argument("--csv", help="CSV with x,y columns")
    src.add_argument("--ha-db", metavar="PATH",
                     help="Home Assistant SQLite recorder, e.g. ~/.homeassistant/home-assistant_v2.db")
    src.add_argument("--ha-recorder", metavar="ENTITY_PREFIX",
                     help="e.g. sensor.living_target_1 (reads _x and _y)")
    p.add_argument("--entity", metavar="PREFIX",
                   help="entity prefix for --ha-db; auto-detected when omitted")
    p.add_argument("--dsn", help="PostgreSQL DSN for --ha-recorder")
    p.add_argument("--keep-origin", action="store_true",
                   help="keep (0,0) readings instead of treating them as 'nobody there'")
    p.add_argument("--emit", choices=("esphome", "ha", "both"),
                   help="also print zone config to paste into ESPHome or Home Assistant")
    p.add_argument("--days", type=float, default=0, help="only the last N days")
    p.add_argument("--scale", type=float, default=1.0, help="multiply input by this to get mm")
    p.add_argument("--cell", type=int, default=250, help="grid size in mm (default 250)")
    p.add_argument("--frac", type=float, default=0.10, help="keep cells above this share of the peak")
    p.add_argument("--x-min", type=float, metavar="MM",
                   help="wall distance to the left of the sensor, negative mm (e.g. -3000)")
    p.add_argument("--x-max", type=float, metavar="MM",
                   help="wall distance to the right of the sensor, mm (e.g. 500)")
    p.add_argument("--y-max", type=float, metavar="MM",
                   help="wall distance straight ahead, mm")
    a = p.parse_args()

    if a.csv:
        pts = from_csv(a.csv, a.scale)
    elif a.ha_db:
        pts = from_sqlite(a.ha_db, a.entity, a.days)
    else:
        pts = from_recorder(a.ha_recorder, a.days, a.dsn)
    if not pts:
        sys.exit("no coordinates found")
    raw = len(pts)
    empty = 0
    if not a.keep_origin:
        pts, empty = drop_no_target(pts)
    pts, dropped = clip_to_room(pts, a.x_min, a.x_max, a.y_max)
    if not pts:
        sys.exit("every reading was (0,0) or outside the room — nothing left to cluster")
    zones, hist = find_zones(pts, a.cell, a.frac)

    print(f"{raw:,} coordinate pairs", end="")
    if empty:
        print(f" — {empty:,} ({empty / raw * 100:.1f}%) with no target, dropped", end="")
    if dropped:
        print(f" — {dropped:,} ({dropped / raw * 100:.1f}%) outside the room, dropped", end="")
    print(f"\ngrid {a.cell} mm, keeping cells above {a.frac:.0%} of the peak\n")
    print(f"{'zone':>4} {'share':>7} {'centre x':>9} {'centre y':>9} {'cells':>6}")
    for i, z in enumerate(zones, 1):
        print(f"{i:>4} {z['mass'] / len(pts) * 100:6.1f}% "
              f"{z['cx'] / 1000:+9.2f} {z['cy'] / 1000:9.2f} {z['cells']:>6}")
    print("\n" + heatmap(hist, a.cell))
    print("\nsensor at x=0, y=0. Nobody drew these zones.")

    if a.emit:
        if len(zones) > ESPHOME_ZONES:
            print(f"\n{len(zones)} zones found, ESPHome takes {ESPHOME_ZONES}. "
                  f"Writing the {ESPHOME_ZONES} busiest — raise --frac to merge the rest.")
        print()
        if a.emit in ("esphome", "both"):
            print(as_esphome(zones))
        if a.emit in ("ha", "both"):
            print(("\n" if a.emit == "both" else "") + as_ha_actions(zones))


if __name__ == "__main__":
    main()