"""The plate as a solid — accumulation seen as depth.

Three ways of looking at the same height field.

  relief   raking light across the real surface. This is close to how scribed
           steel actually reads: the eye sees a groove because of the shadow
           in it, not because of any colour.

  iso      the height field as geometry, viewed obliquely. To see it at all
           the depth has to be exaggerated by three orders of magnitude, and
           saying so is part of the work: eight hours of continuous cutting
           takes off eighty micrometres. The accumulation is almost nothing.

  section  a slice through the plate, in micrometres, at true proportion.
           A single pass is three micrometres. Everything deeper than that is
           the line having come back.

Usage:
  python3 render3d.py ../plate/main/depth --mode relief --out ../plate/main/relief.png
"""
from __future__ import annotations
import argparse
import json
import os

import numpy as np
from PIL import Image

PLATE_W, PLATE_H = 420.0, 297.0
STOCK_UM = 1500.0


def load(depth_dir):
    d = np.load(os.path.join(depth_dir, "depth_um.npy"))
    stats = json.load(open(os.path.join(depth_dir, "depth_stats.json")))
    return d, stats


def downsample(a, factor):
    if factor <= 1:
        return a
    ny, nx = a.shape
    ny2, nx2 = ny // factor, nx // factor
    return a[:ny2 * factor, :nx2 * factor].reshape(ny2, factor, nx2, factor).mean(axis=(1, 3))


def blur(a, r=2):
    """Cheap separable box blur, used for the cavity term."""
    if r < 1:
        return a
    k = 2 * r + 1
    x = np.pad(a.astype(np.float32), ((r, r), (0, 0)), mode="edge")
    c = np.concatenate([np.zeros((1, x.shape[1]), np.float32), np.cumsum(x, axis=0)])
    a1 = (c[k:] - c[:-k]) / k
    x = np.pad(a1, ((0, 0), (r, r)), mode="edge")
    c = np.concatenate([np.zeros((x.shape[0], 1), np.float32), np.cumsum(x, axis=1)], axis=1)
    return (c[:, k:] - c[:, :-k]) / k


def render_relief(depth, res_mm, light_deg=32.0, elev_deg=13.0, exagg=260.0):
    """Normal-mapped shading with a low, raking light."""
    h = -depth * exagg * 1e-3          # mm of relief, negative = cut away
    gy, gx = np.gradient(h, res_mm)
    nz = 1.0 / np.sqrt(gx * gx + gy * gy + 1.0)
    nx_, ny_ = -gx * nz, -gy * nz

    a = np.radians(light_deg)
    e = np.radians(elev_deg)
    lx, ly, lz = np.cos(a) * np.cos(e), np.sin(a) * np.cos(e), np.sin(e)

    lam = np.clip(nx_ * lx + ny_ * ly + nz * lz, 0.0, 1.0)
    # specular: brushed steel throws a hard highlight off the groove walls
    hx, hy, hz = lx, ly, lz + 1.0
    n = np.sqrt(hx * hx + hy * hy + hz * hz)
    spec = np.clip(nx_ * hx / n + ny_ * hy / n + nz * hz / n, 0.0, 1.0) ** 48

    # cavity: deep, enclosed places stay dark
    cav = np.clip(1.0 - (blur(depth, 5) - depth) / 30.0, 0.55, 1.0)

    v = 0.40 + 0.44 * lam * cav + 0.95 * spec * cav
    v = np.clip(v, 0.0, 1.0) ** 0.82
    img = (v * 255).astype(np.uint8)
    return np.dstack([img, img, (np.clip(v * 1.03, 0, 1) * 255).astype(np.uint8)])


def render_iso(depth, res_mm, relief_px=130.0, tilt=0.42, out_w=1600, slab_px=120):
    """The plate as a solid block, seen obliquely.

    The top face is shaded by the same raking-light model as the relief view,
    so the interior of a groove is genuinely darker than its rim. That shading
    is what tells the eye these are trenches cut down into the metal and not
    ridges standing up from it. The vertical displacement then makes them
    geometry rather than an image of geometry.

    Depth is exaggerated and the factor is printed: at true scale eight hours
    of cutting is eighty micrometres across four hundred and twenty
    millimetres, which is nothing you could see.
    """
    if depth.shape[1] >= out_w:
        f = max(1, int(round(depth.shape[1] / out_w)))
        d = downsample(depth, f)
        step = res_mm * f
    else:
        k = max(1, int(round(out_w / depth.shape[1])))
        d = np.repeat(np.repeat(depth, k, axis=0), k, axis=1)
        d = blur(d, max(2, k))               # the tip has a radius; so should the model
        step = res_mm / k
    ny, nx = d.shape

    dmax = float(d.max()) or 1.0
    h_px = d * (relief_px / dmax)
    exagg = (relief_px * step) / (dmax * 1e-3)

    # the honest surface shading, computed on the real micrometre depths
    shade = render_relief(d, step, light_deg=28.0, elev_deg=15.0, exagg=190.0)

    top_pad = 50
    H = int(ny * tilt + relief_px + slab_px + top_pad + 20)
    img = np.zeros((H, nx, 3), dtype=np.uint8)
    img[:, :] = (10, 11, 13)

    # Deeper metal is FURTHER from the camera, so in this oblique view it
    # projects higher up the image, not lower. Getting this sign wrong is what
    # turns trenches into ridges.
    y_all = np.clip((top_pad + relief_px + np.arange(ny)[:, None] * tilt - h_px).astype(np.int32),
                    0, H - 2)
    cols = np.arange(nx)

    for row in range(ny):                    # far to near; nearer rows overwrite
        y0 = y_all[row]
        img[y0, cols] = shade[row]
        # the metal hanging below this row, down to wherever the next row sits
        y1 = y_all[row + 1] if row + 1 < ny else y0 + slab_px
        drop = np.clip(y1 - y0, 0, slab_px)
        for c in np.where(drop > 0)[0]:
            k = int(drop[c])
            base = int(shade[row, c, 0]) * 0.30
            wall = np.linspace(base, max(5.0, base * 0.18), k).astype(np.uint8)
            img[y0[c] + 1:y0[c] + 1 + k, c, 0] = wall
            img[y0[c] + 1:y0[c] + 1 + k, c, 1] = wall
            img[y0[c] + 1:y0[c] + 1 + k, c, 2] = np.clip(wall.astype(np.int16) + 4, 0, 255).astype(np.uint8)

    y_front = y_all[ny - 1]
    for c in range(nx):
        y = y_front[c] + 1
        k = min(slab_px, H - y)
        if k > 1:
            body = np.linspace(70, 11, k).astype(np.uint8)
            img[y:y + k, c, 0] = body
            img[y:y + k, c, 1] = body
            img[y:y + k, c, 2] = np.clip(body.astype(np.int16) + 4, 0, 255).astype(np.uint8)

    print("  iso: depth exaggerated x%.0f (max %.1f um -> %.0f px)" % (exagg, dmax, relief_px))
    return img


def render_section(depth, res_mm, at_mm=None, span_mm=45.0, out=(1800, 640)):
    """A slice through the plate at true proportion: micrometres down against
    the full thickness of the stock. Almost nothing has been taken away, and
    that is the measurement, not a failure of the drawing."""
    ny, nx = depth.shape
    row = int((at_mm if at_mm is not None else PLATE_H * 0.5) / res_mm)
    row = max(0, min(ny - 1, row))
    x0 = max(0, int((PLATE_W * 0.5 - span_mm / 2) / res_mm))
    prof = depth[row, x0:x0 + max(2, int(span_mm / res_mm))].astype(np.float64)

    W, H = out
    xi = np.linspace(0, len(prof) - 1, W)
    p = np.interp(xi, np.arange(len(prof)), prof)      # µm below the top face

    y = (p / STOCK_UM * (H - 1)).astype(np.int32)      # surface row per column
    rows = np.arange(H)[:, None]
    solid = rows >= y[None, :]                          # metal below the surface

    img = np.full((H, W, 3), 250, dtype=np.uint8)
    img[solid] = (26, 28, 30)
    # a hairline on the original surface level, so the loss can be seen at all
    img[0:2, :] = (150, 156, 160)
    return img, float(p.max()), float(p.mean())


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("depth_dir")
    ap.add_argument("--mode", default="relief", choices=["relief", "iso", "section"])
    ap.add_argument("--out", required=True)
    ap.add_argument("--exagg", type=float, default=None)
    ap.add_argument("--crop", nargs=4, type=float, default=None,
                    metavar=("X", "Y", "W", "H"), help="mm window into the plate")
    args = ap.parse_args()

    depth, stats = load(os.path.abspath(args.depth_dir))
    res = stats["res_mm"]

    if args.crop:
        x, y, w, h = args.crop
        depth = depth[int(y / res):int((y + h) / res), int(x / res):int((x + w) / res)]

    if args.mode == "relief":
        img = render_relief(depth, res, exagg=args.exagg or 260.0)
    elif args.mode == "iso":
        img = render_iso(depth, res, relief_px=args.exagg or 190.0)
    else:
        img, mx, mean = render_section(depth, res)
        print("section: max %.1f um, mean %.1f um, stock %.0f um" % (mx, mean, STOCK_UM))

    Image.fromarray(img).save(os.path.abspath(args.out))
    print("wrote %s  %dx%d" % (args.out, img.shape[1], img.shape[0]))


if __name__ == "__main__":
    main()