"""Where the grooves are and which way they run.

Reflection off a prismatic groove depends on the groove's DIRECTION, not on the
history that put it there. So for the optics the disc reduces to one object: how
much groove length there is, per place, per direction. This builds that.

  length[iy, ix, ibin]   millimetres of groove in cell (ix, iy) running in the
                         direction of bin ibin

Grooves are AXIAL, not vectorial: a groove cut while travelling at 30 degrees is
the same groove as one cut at 210 degrees, so directions are folded onto [0, pi).

The same histogram answers a question the drawings only suggested. On the rim the
rotary axis can only place the tool in 0.38 mm tangential steps while the radial
axis is exact to a micrometre, so lines there should run preferentially RADIALLY.
That is measurable: for each segment take the angle psi between the groove and
the local radial direction, and average cos(2 psi) weighted by length.

  +1  every groove radial
   0  isotropic
  -1  every groove tangential

Usage:
  python3 groove_field.py ../plate/polar80/descent --cell 2.0 --bins 36
"""
from __future__ import annotations
import argparse
import json
import os

import numpy as np

DISC_R = 200.0
R_MIN, R_MAX = 5.0, 195.0


def load_points(d):
    npy = os.path.join(d, "stroke.npy")
    if os.path.exists(npy):
        return np.load(npy).astype(np.float64)
    with open(os.path.join(d, "stroke.json")) as f:
        return np.asarray(json.load(f)["pts"], dtype=np.float64)


def segments(pts, chunk=2_000_000):
    """Midpoints, directions and lengths, yielded in chunks."""
    n = len(pts) - 1
    for c0 in range(0, n, chunk):
        c1 = min(n, c0 + chunk)
        a = pts[c0:c1]
        b = pts[c0 + 1:c1 + 1]
        d = b - a
        L = np.hypot(d[:, 0], d[:, 1])
        keep = L > 1e-9
        if not keep.any():
            continue
        mid = 0.5 * (a[keep] + b[keep])
        yield mid, d[keep], L[keep]


def build(pts, cell_mm=2.0, n_bins=36, n_bands=20):
    nx = ny = int(np.ceil(2.0 * DISC_R / cell_mm))
    length = np.zeros((ny, nx, n_bins), dtype=np.float32)

    # radial-order accumulation, per radius band
    edges = np.linspace(R_MIN, R_MAX, n_bands + 1)
    band_len = np.zeros(n_bands)
    band_cos2 = np.zeros(n_bands)      # sum of L*cos(2 psi)
    band_sin2 = np.zeros(n_bands)      # sum of L*sin(2 psi), for the mean angle
    total = 0.0

    for mid, d, L in segments(pts):
        total += float(L.sum())
        # direction folded to [0, pi)
        phi = np.mod(np.arctan2(d[:, 1], d[:, 0]), np.pi)
        ib = np.minimum((phi / np.pi * n_bins).astype(np.int32), n_bins - 1)
        ix = np.clip(((mid[:, 0] + DISC_R) / cell_mm).astype(np.int32), 0, nx - 1)
        iy = np.clip(((mid[:, 1] + DISC_R) / cell_mm).astype(np.int32), 0, ny - 1)
        np.add.at(length, (iy, ix, ib), L.astype(np.float32))

        r = np.hypot(mid[:, 0], mid[:, 1])
        # psi = groove direction relative to the local radial direction
        rad = np.arctan2(mid[:, 1], mid[:, 0])
        psi = phi - rad
        c2 = np.cos(2.0 * psi)
        s2 = np.sin(2.0 * psi)
        ibnd = np.clip(np.searchsorted(edges, r) - 1, 0, n_bands - 1)
        np.add.at(band_len, ibnd, L)
        np.add.at(band_cos2, ibnd, L * c2)
        np.add.at(band_sin2, ibnd, L * s2)

    with np.errstate(invalid="ignore", divide="ignore"):
        order = np.where(band_len > 0, band_cos2 / band_len, 0.0)
        quad = np.where(band_len > 0, band_sin2 / band_len, 0.0)
    resultant = np.hypot(order, quad)
    mean_psi = 0.5 * np.degrees(np.arctan2(quad, order))

    bands = []
    for i in range(n_bands):
        bands.append({
            "r_lo": round(float(edges[i]), 1),
            "r_hi": round(float(edges[i + 1]), 1),
            "length_m": round(float(band_len[i]) / 1000.0, 2),
            "radial_order": round(float(order[i]), 4),
            "alignment_strength": round(float(resultant[i]), 4),
            "mean_angle_from_radial_deg": round(float(mean_psi[i]), 1),
        })
    return length, {
        "cell_mm": cell_mm,
        "direction_bins": n_bins,
        "total_length_m": round(total / 1000.0, 1),
        "bands": bands,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("plate_dir")
    ap.add_argument("--cell", type=float, default=2.0, help="mm per spatial cell")
    ap.add_argument("--bins", type=int, default=36, help="direction bins over 180 deg")
    ap.add_argument("--bands", type=int, default=20)
    ap.add_argument("--out", default=None)
    args = ap.parse_args()

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

    pts = load_points(d)
    length, stats = build(pts, args.cell, args.bins, args.bands)
    stats["source"] = os.path.basename(d)
    stats["points"] = int(len(pts))

    np.savez_compressed(os.path.join(out, "groove_length.npz"),
                        length=length, cell_mm=args.cell, n_bins=args.bins)
    with open(os.path.join(out, "orientation.json"), "w") as f:
        json.dump(stats, f, ensure_ascii=False, indent=1)

    print("%s: %.1f m of groove, %d points" % (stats["source"], stats["total_length_m"],
                                               stats["points"]))
    print()
    print("%-14s %9s %8s %9s %s" % ("radius band", "length", "radial", "strength", "mean angle"))
    for b in stats["bands"]:
        bar = "#" * int(abs(b["radial_order"]) * 40)
        print("%5.0f-%-5.0f mm %8.1f m %+8.3f %9.3f %6.1f deg  %s"
              % (b["r_lo"], b["r_hi"], b["length_m"], b["radial_order"],
                 b["alignment_strength"], b["mean_angle_from_radial_deg"], bar))
    print()
    print("radial_order: +1 = every groove radial, 0 = isotropic, -1 = tangential")


if __name__ == "__main__":
    main()