"""Accumulation as depth — how the plate is actually removed.

A diamond drag stylus does not deposit anything. It takes metal away, and it
takes away almost nothing at a time: a single pass at 25 mm/s leaves a groove
a few micrometres deep. What makes the surface is that the line crosses its
own past. Metal that has been passed once is 3 µm down. Metal that has been
passed forty times is 120 µm down.

So the accumulation is a height field, and the height field is the honest
three-dimensional form of the work: not a drawing on a plate but a plate that
has been eaten into, unevenly, in proportion to how often the line returned.

Depth model
-----------
Per pass the tool removes PASS_UM micrometres at the groove centre, tapering
across a V of width GROOVE mm. Passes accumulate, but not linearly forever —
a groove that is already deep presents less fresh material to the tip, so the
increment decays with existing depth (SATURATION_UM). Nothing is ever added
back: the field is monotonically non-increasing in height, by construction.

Usage:
  python3 depth.py ../plate/main --res 0.1 --out ../plate/main/depth
"""
from __future__ import annotations
import argparse
import json
import os

import numpy as np

PLATE_W, PLATE_H = 420.0, 297.0    # mm
GROOVE = 0.10                       # mm, tip width
PASS_UM = 3.0                       # µm removed by one pass on fresh metal
SATURATION_UM = 140.0               # µm; beyond this a pass removes ~nothing
STOCK_MM = 1.5                      # mm, sheet thickness


def build_depth(pts, res=0.1):
    """Rasterise the stroke into a depth field in micrometres."""
    nx, ny = int(PLATE_W / res), int(PLATE_H / res)
    depth = np.zeros((ny, nx), dtype=np.float32)

    half = max(1, int(round((GROOVE / 2) / res)))
    # V profile across the groove: full removal at the centre, tapering out
    offs = np.arange(-half, half + 1)
    prof = 1.0 - np.abs(offs) / (half + 1.0)

    P = np.asarray(pts, dtype=np.float64)
    # sample every segment at <= res/2 so the groove is continuous
    seg = np.linalg.norm(np.diff(P, axis=0), axis=1)
    steps = np.maximum(1, np.ceil(seg / (res * 0.5)).astype(np.int32))

    total = int(steps.sum())
    xs = np.empty(total, dtype=np.float64)
    ys = np.empty(total, dtype=np.float64)
    k = 0
    for i in range(len(seg)):
        s = steps[i]
        t = (np.arange(s) + 0.5) / s
        xs[k:k + s] = P[i, 0] + (P[i + 1, 0] - P[i, 0]) * t
        ys[k:k + s] = P[i, 1] + (P[i + 1, 1] - P[i, 1]) * t
        k += s

    cx = np.clip((xs / res).astype(np.int32), 0, nx - 1)
    cy = np.clip((ys / res).astype(np.int32), 0, ny - 1)

    # pass count per cell, spread across the groove width in both axes
    passes = np.zeros((ny, nx), dtype=np.float32)
    for dy, wy in zip(offs, prof):
        yy = np.clip(cy + dy, 0, ny - 1)
        for dx, wx in zip(offs, prof):
            xx = np.clip(cx + dx, 0, nx - 1)
            np.add.at(passes, (yy, xx), wy * wx)

    # each pass removes less as the groove deepens
    # d_{n+1} = d_n + PASS_UM * (1 - d_n / SATURATION_UM)
    # closed form for n passes: SAT * (1 - (1 - PASS/SAT)^n)
    r = 1.0 - PASS_UM / SATURATION_UM
    depth = SATURATION_UM * (1.0 - np.power(r, passes, dtype=np.float32))
    return depth, passes


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("plate_dir")
    ap.add_argument("--res", type=float, default=0.1, help="mm per cell")
    ap.add_argument("--out", default=None)
    ap.add_argument("--fraction", type=float, default=1.0,
                    help="use only the first fraction of the stroke")
    args = ap.parse_args()

    d = os.path.abspath(args.plate_dir)
    out = os.path.abspath(args.out or os.path.join(d, "depth"))
    os.makedirs(out, exist_ok=True)

    pts = json.load(open(os.path.join(d, "stroke.json")))["pts"]
    if args.fraction < 1.0:
        pts = pts[:max(2, int(len(pts) * args.fraction))]

    depth, passes = build_depth(pts, args.res)

    cut = depth > 0.05
    stats = {
        "res_mm": args.res,
        "grid": [int(depth.shape[1]), int(depth.shape[0])],
        "fraction_of_stroke": args.fraction,
        "pass_um_fresh": PASS_UM,
        "saturation_um": SATURATION_UM,
        "stock_mm": STOCK_MM,
        "touched_fraction": float(cut.mean()),
        "max_passes": float(passes.max()),
        "mean_passes_where_cut": float(passes[cut].mean()) if cut.any() else 0.0,
        "max_depth_um": float(depth.max()),
        "mean_depth_um_where_cut": float(depth[cut].mean()) if cut.any() else 0.0,
        "removed_mm3": float(depth.sum() * 1e-3 * args.res * args.res),
        "deepest_fraction_of_stock": float(depth.max() / (STOCK_MM * 1000.0)),
    }
    np.save(os.path.join(out, "depth_um.npy"), depth.astype(np.float32))
    with open(os.path.join(out, "depth_stats.json"), "w") as f:
        json.dump(stats, f, ensure_ascii=False, indent=1)
    print(json.dumps(stats, ensure_ascii=False, indent=1))


if __name__ == "__main__":
    main()