"""Spotlight, disc, wall — what the cut metal throws.

The installation. A lamp above and to one side, the disc standing in front of a
wall, the wall catching whatever the disc sends it. Everything is a real length
in millimetres so the result can be checked against a room.

How the projection is computed
------------------------------
Not by tracing rays from the lamp, which would waste almost all of them. Every
piece of the disc is treated as an emitter: for each place on the disc and each
direction its grooves run, work out what leaves and where it lands.

The disc reduces to three numbers per cell, from groove_field.py and
depth_disc.py: how much groove length runs in each direction, and how much of
the cell is still mirror. Then

  mirror area  -> one specular ray, tight lobe. This is the image OF THE LAMP.
  groove length-> the cone from optics.groove_reflect, per direction bin.

Because the groove is prismatic, all of a groove's light lies on a cone about the
groove direction (see optics.py). A cone landing on a flat wall draws a conic
section. So the wall image is a sum of arcs, one family per direction bin, and
the weight of each arc is how many metres of groove point that way. The picture
is literally a histogram of the machine's travel directions, drawn in light.

What is included: Fresnel reflectance for 304 per channel, exact shadowing and
second bounces inside the groove, the lamp's finite angular size, a roughness
lobe that differs for mirror, burnished groove floor and torn ridge, and inverse
square falloff to the wall.

What is not: interreflection between distant parts of the disc, polarisation
tracked through two bounces (the second bounce is charged a scalar), and
diffraction (features are ~100x the wavelength; the blur it adds is a fraction of
a degree, far below the lamp's own angular size).

Usage:
  python3 project.py ../plate/polar80/descent --out ../plate/polar80/descent/light
"""
from __future__ import annotations
import argparse
import json
import math
import os

import numpy as np

import optics as op
import stylus as st

DISC_R = 200.0
R_MIN, R_MAX = 5.0, 195.0


# ---------------------------------------------------------------------------
# the room
# ---------------------------------------------------------------------------
class Install:
    """Where everything is, in millimetres.

    The disc lies in the z = 0 plane, face up, centred on the origin, then is
    tilted by `tilt_deg` about the x axis so that it throws light forward onto a
    wall standing at y = wall_y. Tilting the disc rather than the wall keeps the
    disc's own coordinates the ones the stroke was cut in.
    """

    def __init__(self, lamp=(-700.0, -1500.0, 2200.0), wall_y=2200.0,
                 tilt_deg=20.0, lamp_half_angle_deg=0.6,
                 wall_w=3600.0, wall_h=2600.0, wall_z0=-400.0):
        self.lamp = np.asarray(lamp, float)
        self.wall_y = float(wall_y)
        self.tilt = math.radians(tilt_deg)
        self.lamp_half_angle = math.radians(lamp_half_angle_deg)
        self.wall_w, self.wall_h, self.wall_z0 = wall_w, wall_h, wall_z0

    def disc_to_world(self, x, y):
        """Disc coordinates (mm) to world position and world surface normal."""
        c, s = math.cos(self.tilt), math.sin(self.tilt)
        # rotate about x so the face tips TOWARD the wall:
        # (x, y, 0) -> (x, y c, -y s); normal (0,0,1) -> (0, s, c)
        P = np.stack([x, y * c, -y * s], axis=-1)
        n = np.array([0.0, s, c])
        return P, n

    def basis(self):
        """World directions of the disc's own x and y axes after the tilt."""
        c, s = math.cos(self.tilt), math.sin(self.tilt)
        return np.array([1.0, 0.0, 0.0]), np.array([0.0, c, -s])

    def to_wall(self, P, D):
        """Where rays from P in direction D meet the wall. Returns (u, v, ok, dist)."""
        dy = D[..., 1]
        with np.errstate(divide="ignore", invalid="ignore"):
            tt = (self.wall_y - P[..., 1]) / dy
        ok = np.isfinite(tt) & (tt > 0) & (dy > 1e-9)
        tt = np.where(ok, tt, 0.0)
        X = P[..., 0] + tt * D[..., 0]
        Z = P[..., 2] + tt * D[..., 2]
        dist = tt
        ok &= (np.abs(X) <= self.wall_w * 0.5)
        ok &= (Z >= self.wall_z0) & (Z <= self.wall_z0 + self.wall_h)
        return X, Z, ok, dist

    def describe(self):
        return {
            "lamp_mm": [round(v, 1) for v in self.lamp.tolist()],
            "lamp_distance_to_centre_mm": round(float(np.linalg.norm(self.lamp)), 1),
            "lamp_half_angle_deg": round(math.degrees(self.lamp_half_angle), 2),
            "disc_tilt_deg": round(math.degrees(self.tilt), 1),
            "wall_y_mm": self.wall_y,
            "wall_size_mm": [self.wall_w, self.wall_h],
        }


# ---------------------------------------------------------------------------
# the disc, reduced
# ---------------------------------------------------------------------------
def load_field(plate_dir):
    f = os.path.join(plate_dir, "field", "groove_length.npz")
    if not os.path.exists(f):
        raise SystemExit("run groove_field.py on %s first" % plate_dir)
    z = np.load(f)
    return z["length"], float(z["cell_mm"]), int(z["n_bins"])


# Effective width that a single pass takes out of the mirror. Wider than the
# specified 100 um groove because the ridge either side also destroys the polish,
# but narrower than groove-plus-both-ridges because later passes cut the ridges
# away and the ridge tail is shallower than the 0.05 um detection floor.
# Recalibrated against the windows measured directly from the 5 um height field
# after the groove width was set from the specification; residual about 5
# percentage points, which is worse than before because the random-coverage
# estimate fits less well at the higher coverage a wider groove produces. See
# --self-check, which prints the comparison.
MIRROR_W_UM = 119.0


def mirror_fraction(length, cell_mm, w_um=MIRROR_W_UM):
    """How much of each cell is still the original polished surface.

    Groove length times an effective width is the area the tip has taken out of
    the mirror in that cell. Overlap is handled by the random-coverage estimate:
    after a swept area A over a cell of area C, the untouched fraction is
    exp(-A/C). Exact for a path that crosses itself freely, which this one does.
    """
    swept = length.sum(axis=-1) * (w_um * 1e-3)             # mm^2
    return np.exp(-swept / (cell_mm * cell_mm))


def self_check(plate_dirs):
    """Compare the mirror model against every directly measured window."""
    import glob
    rows = []
    for d in plate_dirs:
        for p in sorted(glob.glob(os.path.join(d, "cut", "cut_crop_*.json"))):
            j = json.load(open(p))
            m = max(0.0, 1.0 - j["cut_area_fraction"] - j["ridge_area_fraction"])
            rows.append((os.path.basename(d), j["window"],
                         j["path_len_in_window_mm"], m))
    if not rows:
        print("no measured windows found; run depth_disc.py first")
        return False
    print("%-10s %-18s %9s %10s %10s %8s"
          % ("run", "window", "path mm", "measured", "model", "error"))
    tot = 0.0
    for run, win, L, m in rows:
        pred = float(np.exp(-L * (MIRROR_W_UM * 1e-3) / 4.0))   # 2x2 mm windows
        tot += abs(pred - m)
        print("%-10s %-18s %9.1f %10.3f %10.3f %+8.3f"
              % (run, win, L, m, pred, pred - m))
    print("mean absolute error %.3f over %d windows (w_eff = %.1f um)"
          % (tot / len(rows), len(rows), MIRROR_W_UM))
    return tot / len(rows) < 0.05


# ---------------------------------------------------------------------------
# the projection
# ---------------------------------------------------------------------------
def render(plate_dir, inst, px=1500, samples=121, bounces=2,
           min_length_mm=0.5, chunk_bins=4000, verbose=True):
    length, cell_mm, n_bins = load_field(plate_dir)
    prof = op.Profile(n=samples)
    gw_um = 2.0 * prof.half_width

    ny, nx, nb = length.shape
    ex, ey = inst.basis()

    # cell centres in disc coordinates
    cx = (np.arange(nx) + 0.5) * cell_mm - DISC_R
    cy = (np.arange(ny) + 0.5) * cell_mm - DISC_R
    CX, CY = np.meshgrid(cx, cy)
    rad = np.hypot(CX, CY)
    on_disc = (rad >= R_MIN) & (rad <= R_MAX)

    mirror = mirror_fraction(length, cell_mm) * on_disc

    H = int(px * inst.wall_h / inst.wall_w)
    img = np.zeros((H, px, 3), dtype=np.float64)
    spec_img = np.zeros((H, px, 3), dtype=np.float64)

    def deposit(target, X, Z, ok, energy, dist):
        if not np.any(ok):
            return 0.0
        u = ((X[ok] + inst.wall_w * 0.5) / inst.wall_w * px).astype(np.int64)
        v = ((inst.wall_z0 + inst.wall_h - Z[ok]) / inst.wall_h * H).astype(np.int64)
        good = (u >= 0) & (u < px) & (v >= 0) & (v < H)
        if not np.any(good):
            return 0.0
        e = energy[ok][good]
        d = np.maximum(dist[ok][good], 1.0)
        e = e / (d * d)[:, None] * 1e6           # inverse square, scaled to taste
        for c in range(3):
            np.add.at(target, (v[good], u[good], c), e[:, c])
        return float(e.sum())

    # ---- the mirror: one specular ray per cell, this is the lamp's image ----
    sel = on_disc & (mirror > 1e-4)
    if sel.any():
        P, nrm = inst.disc_to_world(CX[sel], CY[sel])
        L = inst.lamp[None, :] - P
        dl = np.linalg.norm(L, axis=1, keepdims=True)
        i_vec = -L / dl                          # travelling toward the surface
        cos_in = -(i_vec @ nrm)
        live = cos_in > 1e-3
        if live.any():
            r = i_vec[live] - 2.0 * (i_vec[live] @ nrm)[:, None] * nrm[None, :]
            R = op.steel_rgb(cos_in[live])
            area = (cell_mm * cell_mm) * mirror[sel][live]
            energy = R * (cos_in[live] * area / (dl[live, 0] ** 2) * 1e6)[:, None]
            X, Z, ok, dist = inst.to_wall(P[live], r)
            deposit(spec_img, X, Z, ok, energy, dist)
    if verbose:
        print("  mirror: %d cells still specular, mean mirror fraction %.3f"
              % (int(sel.sum()), float(mirror[on_disc].mean())))

    # ---- the grooves: a cone per direction bin ----
    # Vectorised over bins, looping over the cross-section instead. The escape
    # test is a scalar comparison against each sample's precomputed sky window
    # (Profile.horizon), which is exact because the groove does not vary along
    # its own axis. 500k bins x 121 samples is not something to loop in Python.
    bin_ang = (np.arange(nb) + 0.5) * math.pi / nb
    iy, ix, ib = np.nonzero(length > min_length_mm)
    if verbose:
        print("  grooves: %d (cell, direction) bins carrying %.0f m"
              % (len(iy), length[iy, ix, ib].sum() / 1000.0))
    if len(iy) == 0:
        return img, spec_img, {}

    Lmm = length[iy, ix, ib].astype(np.float64)
    P, nrm = inst.disc_to_world(CX[iy, ix], CY[iy, ix])
    Lv = inst.lamp[None, :] - P
    dl = np.linalg.norm(Lv, axis=1)
    i_vec = -Lv / dl[:, None]

    ang = bin_ang[ib]
    t = np.cos(ang)[:, None] * ex[None, :] + np.sin(ang)[:, None] * ey[None, :]
    uh = np.cross(t, nrm[None, :])
    uh /= np.linalg.norm(uh, axis=1, keepdims=True)

    iu = np.einsum("ij,ij->i", i_vec, uh)
    iz = i_vec @ nrm
    inv_r2 = 1.0 / np.maximum(dl * dl, 1.0)

    # Across-groove roughness rotates the outgoing ray ABOUT the groove axis,
    # which conserves r.t and therefore keeps the ray on its cone: it smears
    # light ALONG the arc rather than off it. Sample that smear here; the small
    # along-groove component, which does take light off the cone, is the image
    # blur applied afterwards.
    n_jit = 7
    jit = np.linspace(-2.2, 2.2, n_jit)
    jw = np.exp(-0.5 * jit * jit)
    jw /= jw.sum()

    for j in range(prof.u.size):
        nu_j, nz_j, ds_j = prof.nu[j], prof.nz[j], prof.ds[j]
        if ds_j <= 0:
            continue
        N = nu_j * uh + nz_j * nrm[None, :]
        cos_in = -np.einsum("ij,ij->i", i_vec, N)
        lit = prof.escapes_at(j, -iu, -iz)
        live = lit & (cos_in > 1e-4)
        if not live.any():
            continue

        k = np.nonzero(live)[0]
        Nk = N[k]
        r = i_vec[k] - 2.0 * np.einsum("ij,ij->i", i_vec[k], Nk)[:, None] * Nk
        ru = np.einsum("ij,ij->i", r, uh[k])
        rz = r @ nrm
        esc = prof.escapes_at(j, ru, rz)

        w = cos_in[k] * ds_j * 1e-3 * Lmm[k]          # mm^2 of projected facet
        if not esc.all():
            # blocked rays hit the opposite facet: its normal is this one with
            # the across-groove component flipped
            blk = ~esc
            N2 = (-nu_j) * uh[k][blk] + nz_j * nrm[None, :]
            rb = r[blk]
            rb = rb - 2.0 * np.einsum("ij,ij->i", rb, N2)[:, None] * N2
            r[blk] = rb
            w[blk] *= 0.55                            # one more reflectance
            rz = r @ nrm

        out = rz > 1e-6
        if not out.any():
            continue
        kk = k[out]
        is_ridge = prof.kind[j] == 2
        scale = op.REFL_SCALE_RIDGE if is_ridge else op.REFL_SCALE_GROOVE
        sigma = 2.0 * (op.SLOPE_RIDGE if is_ridge else op.SLOPE_GROOVE_ACROSS)
        R = op.steel_rgb(cos_in[kk])
        base = R * (w[out] * scale * inv_r2[kk] * 1e6)[:, None]
        r_out, t_out, P_out = r[out], t[kk], P[kk]
        for dj, wj in zip(jit, jw):
            ang = dj * sigma
            if abs(ang) < 1e-9:
                rr = r_out
            else:
                ca, sa = math.cos(ang), math.sin(ang)
                cross = np.cross(t_out, r_out)
                dot = np.einsum("ij,ij->i", t_out, r_out)[:, None]
                rr = r_out * ca + cross * sa + t_out * dot * (1.0 - ca)
            keep = (rr @ nrm) > 1e-6
            if not keep.any():
                continue
            X, Z, ok, dist = inst.to_wall(P_out[keep], rr[keep])
            deposit(img, X, Z, ok, base[keep] * wj, dist)

    return img, spec_img, {
        "cells_on_disc": int(on_disc.sum()),
        "mean_mirror_fraction": round(float(mirror[on_disc].mean()), 4),
        "groove_bins": int(len(iy)),
        "groove_length_m": round(float(length.sum()) / 1000.0, 1),
        "groove_width_um": round(gw_um, 2),
        "wall_px": [px, H],
        "install": inst.describe(),
    }


def blur(img, sigma_px):
    """Separable Gaussian, for the lamp's angular size and the roughness lobe."""
    if sigma_px <= 0.3:
        return img
    n = max(3, int(sigma_px * 4) | 1)
    x = np.arange(n) - n // 2
    k = np.exp(-0.5 * (x / sigma_px) ** 2)
    k /= k.sum()
    out = img
    for ax in (0, 1):
        out = np.apply_along_axis(lambda m: np.convolve(m, k, mode="same"), ax, out)
    return out


def tone(img, gain=1.0, gamma=2.2, knee=0.85, expose_pct=99.5):
    """Filmic rolloff, auto-exposed so the picture is actually visible.

    Exposure is set from a high percentile of the non-zero pixels rather than the
    maximum, because the specular image of the lamp is orders of magnitude
    brighter than the scatter and would otherwise set the scale for everything.
    """
    v = img.astype(np.float64)
    nz = v[v > 0]
    ref = np.percentile(nz, expose_pct) if nz.size else 1.0
    v = v * (gain * 0.6 / max(ref, 1e-30))
    v = v / (1.0 + v / knee)
    return np.clip(v, 0.0, 1.0) ** (1.0 / gamma)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("plate_dir")
    ap.add_argument("--out", default=None)
    ap.add_argument("--px", type=int, default=1500)
    ap.add_argument("--samples", type=int, default=121)
    ap.add_argument("--bounces", type=int, default=2)
    ap.add_argument("--tilt", type=float, default=20.0)
    ap.add_argument("--lamp-half-angle", type=float, default=0.6,
                    help="a 60 mm gallery spot at 2.75 m is about 0.6 deg")
    ap.add_argument("--gain", type=float, default=1.0)
    ap.add_argument("--no-specular", action="store_true",
                    help="drop the mirror image of the lamp, to see the grooves alone")
    ap.add_argument("--self-check", action="store_true")
    args = ap.parse_args()

    if args.self_check:
        import glob
        dirs = sorted(glob.glob(os.path.join(
            os.path.dirname(os.path.abspath(args.plate_dir)), "*")))
        raise SystemExit(0 if self_check([d for d in dirs if os.path.isdir(d)]) else 1)

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

    inst = Install(tilt_deg=args.tilt, lamp_half_angle_deg=args.lamp_half_angle)
    print("projecting %s" % os.path.basename(d))
    img, spec, meta = render(d, inst, px=args.px, samples=args.samples,
                             bounces=args.bounces)

    # the lamp's angular size blurs everything; the specular image is blurred by
    # the lamp size and the mirror's own tiny lobe, the grooves by their own
    px_per_rad = args.px / (2.0 * math.atan(inst.wall_w * 0.5 / inst.wall_y))
    lamp_px = px_per_rad * inst.lamp_half_angle
    rough_px = px_per_rad * op.SLOPE_GROOVE_ALONG * 2.0
    print("  blur: lamp %.1f px, groove roughness %.1f px (of %d across the wall)"
          % (lamp_px, rough_px, args.px))
    # a uniform disc source of angular radius rho has sigma = rho/2, not rho
    lamp_px *= 0.5
    sharp = blur(img, lamp_px)                     # lamp size only
    img = blur(img, math.hypot(lamp_px, rough_px))  # and the burnished floor
    spec = blur(spec, math.hypot(lamp_px, px_per_rad * op.SLOPE_MIRROR * 2))

    from PIL import Image
    parts = {"grooves": img, "grooves_sharp": sharp, "specular": spec,
             "both": (img if args.no_specular else img + spec)}
    for name, a in parts.items():
        s = a.sum()
        g = args.gain * (60.0 / max(s / a.size, 1e-12)) ** 0.0 if False else args.gain
        p = os.path.join(out, "wall_%s.png" % name)
        Image.fromarray((tone(a, gain=g) * 255).astype(np.uint8), "RGB").save(p)
        print("  wrote %s  (total energy %.4g)" % (p, float(s)))

    meta["energy"] = {k: float(v.sum()) for k, v in parts.items()}
    with open(os.path.join(out, "projection.json"), "w") as f:
        json.dump(meta, f, ensure_ascii=False, indent=1)
    print(json.dumps(meta, ensure_ascii=False, indent=1))


if __name__ == "__main__":
    main()