"""3D Scroll: the fit's mesh as a compact point cloud, so the page can roll and unroll the scroll in the browser.

  GET /api/mesh3d.bin?stride=N    the binary below; N is 4, 8 or 16 (default 8) on the in-memory grid
  GET /api/mesh3d.json?stride=N   the same header as JSON, with the byte layout, for tools

The in-memory mesh keeps every XYZ_STRIDE-th node (2) of the 956 by 56,913 tifxyz grid; stride N keeps every N-th of
those, so a node is every 2N-th grid step (0.384 mm times N along the sheet). Only nodes whose x, y and z are all
valid are sent, in row-major order over the sampled grid. Built once per stride with numpy and kept in memory.

Layout of mesh3d.bin, little-endian throughout (HEADER_BYTES = 72, then four arrays, 72 + 13 n bytes in all):

  offset   type          field
  0        char[4]       magic "P4S3"
  4        uint16        version (1)
  6        uint16        stride N
  8        uint32        n, the number of nodes
  12       uint32        rows of the sampled grid (before invalid nodes are dropped)
  16       uint32        cols of the sampled grid
  20       uint32        map width in px (142,283)
  24       uint32        map height in px (2,390)
  28       float32[3]    bbox min x, y, z of the nodes, in 9.6 um scan voxels
  40       float32[3]    bbox max x, y, z
  52       float32       map px per flat unit (2.5: one full mesh grid step)
  56       float32       mm per map px (0.0768)
  60       float32       mm per voxel (0.0096)
  64       uint16        ink pyramid level k
  66       uint16        ink window side in map px (2 ** (k + 1))
  68       uint32        pad, 0
  72       uint16[n][3]  position x, y, z: voxel = min + q * (max - min) / 65535, per axis
  72+6n    uint16[n][2]  flat position: full mesh grid column and row; map x = 2.5 col, map y = 2.5 row
  72+10n   int16[n]      winding at the node (-1 no winding, -2 ambiguous)
  72+12n   uint8[n]      ink: mean of the map pixels in a 2 ** (k + 1) px square centered on the node, from
                         pyramid level k, the level whose 2x2 window is at most the node spacing; 0 where the map
                         has no ink data (exactly 0 there, which the page draws as fine dots, not as low ink)

Scan axes: x and y across the scroll, z along it (up is high z, the top of the map). Map x runs from the outer end
(x = 0) to the inner end. An ink map is model output, not a reading.
"""

from __future__ import annotations

import math
import struct
import threading
import time

import numpy as np
from fastapi import APIRouter, HTTPException, Query
from fastapi.responses import JSONResponse, Response


class _Core:
    """app.py, looked up on use: app imports this module at its end, so a top-level import would be circular."""

    def __getattr__(self, name):
        import app
        return getattr(app, name)


core = _Core()
router = APIRouter()

MAGIC = b"P4S3"
VERSION = 1
STRIDES = (4, 8, 16)
VOXEL_MM = 9.6e-3
HEADER = struct.Struct("<4sHHIIIII3f3ffffHHI")
HEADER_BYTES = HEADER.size  # 72
HEADER_FIELDS = [  # name, struct code, offset: the same layout as the docstring, for tools that read the JSON
    ("magic", "4s", 0), ("version", "uint16", 4), ("stride", "uint16", 6), ("n", "uint32", 8), ("rows", "uint32", 12),
    ("cols", "uint32", 16), ("map_width", "uint32", 20), ("map_height", "uint32", 24),
    ("bbox_min_voxels", "float32[3]", 28), ("bbox_max_voxels", "float32[3]", 40), ("map_px_per_flat_unit", "float32", 52),
    ("mm_per_px", "float32", 56), ("mm_per_voxel", "float32", 60), ("ink_level", "uint16", 64),
    ("ink_window_px", "uint16", 66), ("pad", "uint32", 68)]
CACHE = {"Cache-Control": "public, max-age=86400"}

_cache: dict = {}       # stride -> (source key, bytes, header dict)
_build_lock = threading.Lock()


def _source_key():
    D = core.D
    return id(D["X"]), D["X"].shape, id(D["ink"])


def build(stride: int = 8):
    """(bytes, header) for one stride; numpy only, a fraction of a second for the full Paris 4 grid."""
    t0 = time.perf_counter()
    D = core.D
    xyz_stride, per_grid = core.XYZ_STRIDE, core.INK_PER_GRID
    X, Y, Z = (D[k][::stride, ::stride] for k in "XYZ")
    rows, cols = X.shape
    ok = np.isfinite(X) & np.isfinite(Y) & np.isfinite(Z)
    ii, jj = np.nonzero(ok)                      # row-major order
    n = int(ii.size)
    P = np.stack([X[ok], Y[ok], Z[ok]], 1).astype(np.float32)
    if n:
        lo, hi = P.min(0), P.max(0)
    else:
        lo = hi = np.zeros(3, np.float32)
    rng = (hi - lo).astype(np.float32)
    q = np.zeros((n, 3), np.uint16)
    for a in range(3):
        if rng[a] > 0:
            q[:, a] = np.round((P[:, a] - lo[a]) / rng[a] * 65535.0).clip(0, 65535).astype(np.uint16)
    # full mesh grid position of each node, which is also its flat position
    gr = ii.astype(np.int64) * stride * xyz_stride
    gc = jj.astype(np.int64) * stride * xyz_stride
    if n and (gr.max() > 65535 or gc.max() > 65535):
        raise HTTPException(500, "the mesh grid is too large for uint16 flat positions")
    flat = np.stack([gc, gr], 1).astype(np.uint16)
    g = D["wind"]
    wind = g[np.minimum(gr, g.shape[0] - 1), np.minimum(gc, g.shape[1] - 1)].astype(np.int16)
    # ink: a 2x2 window at pyramid level k centered on the node, no wider than the node spacing
    spacing = stride * xyz_stride * per_grid
    pyr = D["pyr"]
    k = int(min(max(0, math.floor(math.log2(spacing)) - 1), len(pyr) - 1))
    a = pyr[k]
    h, w = a.shape
    s = float(2 ** k)
    mx, my = gc * per_grid, gr * per_grid
    x0 = np.clip(np.floor(mx / s - 0.5).astype(np.int64), 0, max(w - 2, 0))
    y0 = np.clip(np.floor(my / s - 0.5).astype(np.int64), 0, max(h - 2, 0))
    x1, y1 = np.minimum(x0 + 1, w - 1), np.minimum(y0 + 1, h - 1)
    tot = (a[y0, x0].astype(np.uint16) + a[y0, x1] + a[y1, x0] + a[y1, x1])
    ink = ((tot + 2) // 4).astype(np.uint8)
    W, H = int(D["W"]), int(D["H"])
    head = HEADER.pack(MAGIC, VERSION, stride, n, rows, cols, W, H, *map(float, lo), *map(float, hi),
                       float(per_grid), float(core.MM_PER_PX), VOXEL_MM, k, 2 ** (k + 1), 0)
    body = b"".join([head, q.astype("<u2").tobytes(), flat.astype("<u2").tobytes(), wind.astype("<i2").tobytes(),
                     ink.tobytes()])
    lo32, hi32 = lo.astype(np.float32), hi.astype(np.float32)
    off = HEADER_BYTES
    layout = [{"name": "header", "dtype": "struct", "shape": [HEADER_BYTES], "offset": 0, "bytes": HEADER_BYTES}]
    for name, dt, shape, size in (("position", "<u2", [n, 3], 6), ("flat", "<u2", [n, 2], 4),
                                  ("winding", "<i2", [n], 2), ("ink", "u1", [n], 1)):
        layout.append({"name": name, "dtype": dt, "shape": shape, "offset": off, "bytes": size * n})
        off += size * n
    info = {
        "format": MAGIC.decode(), "version": VERSION, "stride": stride, "grid_step": stride * xyz_stride,
        "nodes": n, "rows": int(rows), "cols": int(cols), "bytes": len(body), "header_bytes": HEADER_BYTES,
        "map": {"width": W, "height": H, "mm_per_px": core.MM_PER_PX},
        "mm_per_voxel": VOXEL_MM, "node_spacing_mm": round(stride * xyz_stride * core.MM_PER_GRID, 4),
        "bbox_voxels": {"min": [round(float(v), 3) for v in lo32], "max": [round(float(v), 3) for v in hi32]},
        "bbox_mm": {"min": [round(float(v) * VOXEL_MM, 4) for v in lo32],
                    "max": [round(float(v) * VOXEL_MM, 4) for v in hi32]},
        "position": {"dtype": "uint16", "axes": ["x", "y", "z"], "decode": "voxel = min + q * (max - min) / 65535",
                     "quantum_mm": [round(float(r) / 65535 * VOXEL_MM, 6) for r in (hi32 - lo32)]},
        "flat": {"axes": ["grid column", "grid row"], "map_px_per_unit": per_grid,
                 "decode": "map x = 2.5 * column, map y = 2.5 * row"},
        "winding": {"dtype": "int16", "no_winding": -1, "ambiguous": -2},
        "ink": {"dtype": "uint8", "pyramid_level": k, "window_px": 2 ** (k + 1), "no_data": 0,
                "no_data_nodes": int((ink == 0).sum())},
        "header_fields": [{"name": a_, "type": t_, "offset": o_} for a_, t_, o_ in HEADER_FIELDS],
        "layout": layout,
        "axes": "scan x, y across the scroll, z along it; map x from the outer end (0) to the inner end",
        "build_seconds": round(time.perf_counter() - t0, 3),
        "note": "The fit's mesh, sampled; an ink map is model output, not a reading.",
    }
    return body, info


def get(stride: int = 8):
    """Cached (bytes, header); builds on first use under one of the app's heavy-work slots."""
    if stride not in STRIDES:
        raise HTTPException(400, f"stride must be one of {', '.join(map(str, STRIDES))}")
    core.need()
    key = _source_key()
    hit = _cache.get(stride)
    if hit and hit[0] == key:
        return hit[1], hit[2]
    with _build_lock:
        hit = _cache.get(stride)
        if hit and hit[0] == key:
            return hit[1], hit[2]
        with core.busy():
            body, info = build(stride)
        info["url"] = f"{core.SPACE_URL}/api/mesh3d.bin?stride={stride}"
        _cache[stride] = (key, body, info)
        return body, info


@router.get("/api/mesh3d.bin", tags=["map"])
def api_mesh3d_bin(stride: int = Query(8, description="4, 8 or 16 on the in-memory grid")):
    """The fit's mesh as a point cloud: quantized scan position, flat map position, winding and ink per node.
    Layout: see /api/mesh3d.json or the module docstring of scroll3d.py."""
    body, info = get(stride)
    return Response(body, media_type="application/octet-stream",
                    headers=CACHE | {"X-Mesh3d-Nodes": str(info["nodes"]),
                                     "Content-Disposition": f'inline; filename="paris4-mesh3d-s{stride}.bin"'})


@router.get("/api/mesh3d.json", tags=["map"])
def api_mesh3d_json(stride: int = Query(8, description="4, 8 or 16 on the in-memory grid")):
    """The header of /api/mesh3d.bin as JSON: counts, bbox in voxels and mm, scales and the byte layout."""
    _, info = get(stride)
    return JSONResponse(info, headers=CACHE)