AI Flow: Reverse Engineering

"""What the video may say: the run's own files, read at build time, and nothing else.

`Run` loads the artifacts of a run folder. Every figure a scene shows or speaks goes through `Run.num`, which records
where it was read, and every picture through `Run.used`; both end up in video.json. Board millimetres become picture
pixels by the frame the run itself wrote (scale.json `frame`), 40 px/mm from (100, 100) when it names none.
"""
import json
import os
import re
import numpy as np
import cv2
from PIL import Image, ImageOps

Image.MAX_IMAGE_PIXELS = None


def clean(s):
    """Narration and captions use plain punctuation: no long dashes, no line breaks, one space between words."""
    s = re.sub(r"\s*" + chr(0x2014) + r"\s*", ", ", str(s)).replace(chr(0x2013), " to ").replace(chr(0x2026), "...").replace("\n", " ")
    return re.sub(r"\s+", " ", s).strip()


def say_num(v, digits=1):
    """A number as it is written on screen and read aloud: no trailing zeros, thousands separated."""
    if isinstance(v, bool) or v is None:
        return str(v)
    if isinstance(v, int) or float(v).is_integer():
        return f"{int(v):,}"
    return f"{float(v):,.{digits}f}".rstrip("0").rstrip(".")


def plural(n, one, many=None):
    return f"{say_num(n)} {one if n == 1 else (many or one + 's')}"


class Run:
    def __init__(self, project_dir, facts):
        self.dir = os.path.abspath(project_dir)
        self.flow = os.path.join(self.dir, "reverse-flow")
        self.facts = facts or {}
        self._json = {}
        self.numbers = []     # {scene, label, value, source}
        self.files = {}       # scene -> [paths shown]
        self.notes = []       # things the build wants on record (fallbacks, what was missing)
        self.scene = "run"
        self.out = None       # the video folder, set by the build: drawn stills are kept under it
        fr = (self.j("scale.json") or {}).get("frame") or {}
        self.ppm_frame = float(fr.get("pxPerMm") or 40.0)
        self.origin = tuple(fr.get("origin") or (100.0, 100.0))
        # the board pictures are worked on at most 3000 px wide; k is that reduction (1 for most boards)
        self.k = 1.0
        for side in ("top", "bottom"):
            try:
                with Image.open(self.p(f"board/{side}.jpg")) as im:
                    self.k = min(self.k, 3000.0 / im.width)
            except OSError:
                pass
        self.ppm = self.ppm_frame * self.k

    # ---- files
    def p(self, rel):
        return os.path.join(self.flow, rel)

    def has(self, rel):
        return os.path.isfile(self.p(rel))

    def j(self, rel):
        if rel not in self._json:
            try:
                with open(self.p(rel), encoding="utf-8") as f:
                    self._json[rel] = json.load(f)
            except (OSError, ValueError):
                self._json[rel] = None
        return self._json[rel]

    def rel(self, path):
        return os.path.relpath(path, self.dir)

    def used(self, path):
        """Record a picture or file a scene shows; returns the path."""
        r = self.rel(path) if os.path.isabs(path) else os.path.join("reverse-flow", path)
        self.files.setdefault(self.scene, [])
        if r not in self.files[self.scene]:
            self.files[self.scene].append(r)
        return path

    def img(self, rel_or_abs, record=True):
        """A picture as RGB, or None when it is missing or unreadable."""
        path = rel_or_abs if os.path.isabs(rel_or_abs) else self.p(rel_or_abs)
        if not os.path.isfile(path):
            return None
        try:
            im = ImageOps.exif_transpose(Image.open(path)).convert("RGB")   # a phone photo is turned the way it was held
        except Exception:
            return None
        if record:
            self.used(path)
        return np.asarray(im)

    def keep(self, name, rgb):
        """Save a picture the tool drew over the run's pictures (boxes on parts, blocks, copper), so what the video showed
        stays in the run folder as a still. Recorded with the scene's files."""
        if not self.out:
            return
        d = os.path.join(self.out, "stills", "drawn")
        os.makedirs(d, exist_ok=True)
        path = os.path.join(d, name + ".jpg")
        try:
            Image.fromarray(rgb).save(path, quality=88)
            self.used(path)
        except OSError:
            pass

    # ---- figures
    def num(self, label, value, source):
        """Record a figure with where it was read, and return it."""
        self.numbers.append({"scene": self.scene, "label": label, "value": value, "source": source})
        return value

    def note(self, text):
        self.notes.append({"scene": self.scene, "note": text})

    # ---- the board frame
    def px(self, mm):
        return ((mm[0] - self.origin[0]) * self.ppm, (mm[1] - self.origin[1]) * self.ppm)

    def pts(self, mm_points):
        return np.array([self.px(q) for q in mm_points], np.float32)

    def face(self, side):
        im = self.img(f"board/{side}.jpg")
        if im is not None and self.k < 1.0:
            im = cv2.resize(im, (int(round(im.shape[1] * self.k)), int(round(im.shape[0] * self.k))), interpolation=cv2.INTER_AREA)
        return im

    def sides(self):
        return [s for s in ("top", "bottom") if self.has(f"board/{s}.jpg")]

    # ---- parts
    def parts(self):
        return ((self.j("parts.json") or {}).get("parts")) or []

    def body_corners(self):
        """ref -> four corners in board mm. From the footprints the run was built from (facts.json, written by the
        program); for a part with no such footprint, a box from the sizes its footprint name or hole gives."""
        out = {b["ref"]: b["corners"] for b in self.facts.get("bodies", []) if len(b.get("corners", [])) == 4}
        for p in self.parts():
            ref = p.get("ref")
            if not ref or ref in out or not isinstance(p.get("pos_mm"), list):
                continue
            w, h = _nominal_size(p)
            a = np.radians(float(p.get("rot") or 0))
            c, s = np.cos(a), np.sin(a)
            x, y = p["pos_mm"][:2]
            out[ref] = [[x + lx * c + ly * s, y - lx * s + ly * c] for lx, ly in ((-w / 2, -h / 2), (w / 2, -h / 2), (w / 2, h / 2), (-w / 2, h / 2))]
        return out


def _nominal_size(p):
    """A part's size in mm from what its record states: its hole pad, a metric chip code or an LxW in the footprint name."""
    if isinstance(p.get("hole"), dict):
        d = float(p["hole"].get("pad_mm") or p["hole"].get("drill_mm") or 2.0)
        return d, d
    name = str(p.get("footprint") or "") + " " + str(p.get("package") or "")
    m = re.search(r"_(\d{2})(\d{2})Metric", name)
    if m:
        return int(m.group(1)) / 10.0, int(m.group(2)) / 10.0
    m = re.search(r"(\d+(?:[p.]\d+)?)x(\d+(?:[p.]\d+)?)\s*mm", name)
    if m:
        return float(m.group(1).replace("p", ".")), float(m.group(2).replace("p", "."))
    return 1.0, 0.6


# ---------------------------------------------------------------- drawing on a board picture

def dim(rgb, k=0.5, grey=0.55):
    """The photo dimmed and partly greyed, so colour drawn on it means something."""
    f = rgb.astype(np.float32)
    g = f @ np.array([0.299, 0.587, 0.114], np.float32)
    out = (f * (1 - grey) + g[:, :, None] * grey) * k + 10
    return np.clip(out, 0, 255).astype(np.uint8)


def zone_mask(run, zones, shape):
    """The exact copper of these pours as a mask in the board picture, from each pour's run-length rows (or its outline)."""
    # the run-length rows are in the frame's own pixels: drawn at that size, then brought to the working size
    full = (int(round(shape[0] / run.k)), int(round(shape[1] / run.k))) if run.k < 1.0 else shape[:2]
    m = np.zeros(full, np.uint8)
    h, w = m.shape
    for z in zones:
        rows = z.get("rle_px")
        if rows:
            for r in rows:
                y = int(r[0])
                if not 0 <= y < h:
                    continue
                for k in range(1, len(r) - 1, 2):
                    m[y, max(0, int(r[k])):min(w, int(r[k + 1]))] = 255
        elif z.get("polygon_mm"):
            cv2.fillPoly(m, [np.round(run.pts(z["polygon_mm"]) / run.k).astype(np.int32)], 255)
    if m.shape != tuple(shape[:2]):
        m = cv2.resize(m, (shape[1], shape[0]), interpolation=cv2.INTER_AREA)
    return m


def tint(rgb, mask, colour, alpha):
    a = (mask.astype(np.float32) / 255.0 * alpha)[:, :, None]
    return (rgb * (1 - a) + np.array(colour, np.float32) * a + 0.5).astype(np.uint8)


def draw_tracks(run, rgb, traces, colour, min_px=2):
    for t in traces:
        path = t.get("path_mm") or []
        if len(path) < 2:
            continue
        th = max(min_px, int(round(float(t.get("width_mm") or 0.15) * run.ppm)))
        cv2.polylines(rgb, [np.round(run.pts(path)).astype(np.int32)], False, colour, th, cv2.LINE_AA)


def draw_vias(run, rgb, vias, colour, ring=(10, 12, 17)):
    for v in vias:
        if not isinstance(v.get("pos_mm"), list):
            continue
        x, y = run.px(v["pos_mm"])
        r = max(3, int(round(float(v.get("diameter_mm") or 0.5) * run.ppm / 2)))
        cv2.circle(rgb, (int(round(x)), int(round(y))), r + 2, ring, 2, cv2.LINE_AA)
        cv2.circle(rgb, (int(round(x)), int(round(y))), r, colour, 2, cv2.LINE_AA)


def densest(points_px, shape, win_w, win_h):
    """The window of this size that holds the most points: the place to zoom into. Returns (x0, y0, x1, y1) or None."""
    if len(points_px) < 8:
        return None
    h, w = shape[:2]
    cell = 16
    grid = np.zeros((h // cell + 1, w // cell + 1), np.float32)
    pts = np.asarray(points_px)
    gx = np.clip((pts[:, 0] // cell).astype(int), 0, grid.shape[1] - 1)
    gy = np.clip((pts[:, 1] // cell).astype(int), 0, grid.shape[0] - 1)
    np.add.at(grid, (gy, gx), 1)
    kw, kh = max(1, int(win_w // cell)), max(1, int(win_h // cell))
    box = cv2.boxFilter(grid, -1, (kw, kh), normalize=False, borderType=cv2.BORDER_CONSTANT)
    # keep the window inside the picture
    ys, xs = slice(kh // 2, max(kh // 2 + 1, grid.shape[0] - kh // 2)), slice(kw // 2, max(kw // 2 + 1, grid.shape[1] - kw // 2))
    sub = box[ys, xs]
    if sub.size == 0 or sub.max() <= 0:
        return None
    iy, ix = np.unravel_index(int(sub.argmax()), sub.shape)
    cx, cy = (ix + kw // 2 + 0.5) * cell, (iy + kh // 2 + 0.5) * cell
    return (cx - win_w / 2, cy - win_h / 2, cx + win_w / 2, cy + win_h / 2)