#!/usr/bin/env python3
"""Rasterize upward mesh triangles into a top-down field height map."""
import argparse
import array
import json
import math
import struct
import sys

Q16 = 1 << 16
MASK64 = (1 << 64) - 1


def trunc_div(numerator, denominator):
    return numerator // denominator if numerator >= 0 else -((-numerator) // denominator)


def transition_hash(seed, interface, x, y, z, salt):
    value = 14695981039346656037
    for word in (seed & 0xffffffff, interface & 0xffffffff,
                 x & MASK64, y & MASK64, z & MASK64, salt):
        for byte in range(8):
            value ^= (word >> (byte * 8)) & 0xff
            value = (value * 1099511628211) & MASK64
    value ^= value >> 30
    value = (value * 0xbf58476d1ce4e5b9) & MASK64
    value ^= value >> 27
    value = (value * 0x94d049bb133111eb) & MASK64
    value ^= value >> 31
    return ((value >> 32) >> 15) - Q16


def transition_fade(value):
    value = max(0, min(Q16, value))
    square = value * value // Q16
    return square * (3 * Q16 - 2 * value) // Q16


def gradient_height_um(layout, x_um, y_um):
    upper = next(item for item in layout["strates"] if item["order"] == 0)
    lower = next(item for item in layout["strates"] if item["order"] == 1)
    width = upper["transition_width_um"]
    cell = max(1_000_000, layout["strate_layout_chunk_size_um"] * 4)
    ix = x_um // cell
    iy = y_um // cell
    tx = transition_fade(((x_um - ix * cell) * Q16) // cell)
    ty = transition_fade(((y_um - iy * cell) * Q16) // cell)

    def at(dx, dy):
        noise = transition_hash(0, upper["strate_index"], ix + dx, iy + dy, 0,
                                0x47524144)
        return noise

    def lerp(a, b, t):
        return a + trunc_div((b - a) * t, Q16)

    low = lerp(lerp(at(0, 0), at(1, 0), tx),
               lerp(at(0, 1), at(1, 1), tx), ty)
    amplitude_divisor = 256 if upper.get("generator_kind") == "FloatingIslands" else 8
    displacement = trunc_div(low * (width // amplitude_divisor), Q16)
    seam = upper["bedrock"]["bottom_um"]
    thickness = upper["bedrock"]["top_um"] - seam
    return seam + displacement + thickness


def read_scene(path):
    triangles = []
    with open(path, "rb") as stream:
        if stream.read(8) != b"VFMSH001":
            raise SystemExit(f"{path}: invalid scene header")
        version, voxel_um, count = struct.unpack("<3I", stream.read(12))
        stream.read(40)
        if version != 1 or count > 100000:
            raise SystemExit("unsupported scene")
        for _ in range(count):
            key = struct.unpack("<4i", stream.read(16))
            surface_vertices, skirt_count, vertex_count, index_count, mask_count = struct.unpack(
                "<5I", stream.read(20))
            positions = array.array("f")
            positions.frombytes(stream.read(vertex_count * 12))
            normals = array.array("f")
            normals.frombytes(stream.read(vertex_count * 12))
            indices = array.array("I")
            indices.frombytes(stream.read(index_count * 4))
            masks = stream.read(mask_count)
            if sys.byteorder != "little":
                positions.byteswap()
                normals.byteswap()
                indices.byteswap()
            for index in range(0, index_count, 3):
                if masks and masks[index // 3]:
                    continue
                a, b, c = indices[index:index + 3]
                if max(a, b, c) >= surface_vertices:
                    continue
                nz = (normals[a * 3 + 2] + normals[b * 3 + 2]
                      + normals[c * 3 + 2]) / 3.0
                if nz < 0.35:
                    continue
                vertices = tuple((positions[v * 3], positions[v * 3 + 1],
                                  positions[v * 3 + 2]) for v in (a, b, c))
                triangles.append(vertices)
    return voxel_um, triangles


def rasterize(triangles, bounds, width, height):
    x0, x1, y0, y1 = bounds
    values = [None] * (width * height)
    sx, sy = width / (x1 - x0), height / (y1 - y0)
    for tri in triangles:
        (ax, ay, az), (bx, by, bz), (cx, cy, cz) = tri
        den = (by - cy) * (ax - cx) + (cx - bx) * (ay - cy)
        if abs(den) < 1.0e-12:
            continue
        left = max(0, int(math.floor((min(ax, bx, cx) - x0) * sx)))
        right = min(width - 1, int(math.ceil((max(ax, bx, cx) - x0) * sx)))
        top = max(0, int(math.floor((y1 - max(ay, by, cy)) * sy)))
        bottom = min(height - 1, int(math.ceil((y1 - min(ay, by, cy)) * sy)))
        for py in range(top, bottom + 1):
            y = y1 - (py + 0.5) / sy
            row = py * width
            for px in range(left, right + 1):
                x = x0 + (px + 0.5) / sx
                u = ((by - cy) * (x - cx) + (cx - bx) * (y - cy)) / den
                v = ((cy - ay) * (x - cx) + (ax - cx) * (y - cy)) / den
                w = 1.0 - u - v
                if min(u, v, w) < -1.0e-6:
                    continue
                z = u * az + v * bz + w * cz
                idx = row + px
                if values[idx] is None or z > values[idx]:
                    values[idx] = z
    return values


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("scene")
    parser.add_argument("output_ppm")
    parser.add_argument("--bounds", nargs=4, type=float, required=True,
                        metavar=("XMIN", "XMAX", "YMIN", "YMAX"))
    parser.add_argument("--size", nargs=2, type=int, default=(640, 640))
    parser.add_argument("--range", nargs=2, type=float, default=(-122.0, -118.0),
                        metavar=("ZMIN", "ZMAX"))
    parser.add_argument("--ring-center", nargs=2, type=float,
                        metavar=("X", "Y"))
    parser.add_argument("--probe-max-radius", type=int, default=16)
    parser.add_argument("--layout")
    args = parser.parse_args()
    voxel_um, triangles = read_scene(args.scene)
    width, height = args.size
    sx, sy = width / (args.bounds[1] - args.bounds[0]), height / (args.bounds[3] - args.bounds[2])
    values = rasterize(triangles, args.bounds, width, height)
    finite = sorted(value for value in values if value is not None)
    if not finite:
        raise SystemExit("no upward-facing field samples in bounds")
    low, high = args.range
    if high <= low:
        raise SystemExit("height range must increase")
    stops = ((0.00, (31, 52, 92)), (0.24, (41, 119, 153)),
             (0.50, (85, 177, 143)), (0.73, (226, 207, 105)),
             (1.00, (181, 71, 49)))
    with open(args.output_ppm, "wb") as output:
        output.write(f"P6\n{width} {height}\n255\n".encode())
        for value in values:
            if value is None:
                output.write(bytes((18, 22, 30)))
                continue
            t = max(0.0, min(1.0, (value - low) / (high - low)))
            for index in range(len(stops) - 1):
                if t <= stops[index + 1][0]:
                    ta, ca = stops[index]
                    tb, cb = stops[index + 1]
                    blend = (t - ta) / (tb - ta)
                    rgb = tuple(round(ca[c] + blend * (cb[c] - ca[c]))
                                for c in range(3))
                    break
            else:
                rgb = stops[-1][1]
            output.write(bytes(rgb))
    def quantile(p):
        return finite[round((len(finite) - 1) * p)]
    print(f"[RingFieldRaster] voxel_um={voxel_um} triangles={len(triangles)} "
          f"pixels={len(finite)}/{len(values)} height_min_m={finite[0]:.6f} "
          f"height_p05_m={quantile(.05):.6f} height_median_m={quantile(.5):.6f} "
          f"height_p95_m={quantile(.95):.6f} height_max_m={finite[-1]:.6f} "
          f"height_range_m={finite[-1]-finite[0]:.6f} "
          f"map={args.output_ppm}")
    if args.ring_center:
        cx, cy = args.ring_center
        def at_world(x, y):
            fx = (x - args.bounds[0]) * sx - 0.5
            fy = (args.bounds[3] - y) * sy - 0.5
            ix, iy = math.floor(fx), math.floor(fy)
            tx, ty = fx - ix, fy - iy
            samples = []
            for px, py in ((ix, iy), (ix + 1, iy), (ix, iy + 1), (ix + 1, iy + 1)):
                if not (0 <= px < width and 0 <= py < height):
                    return None
                value = values[py * width + px]
                if value is None:
                    return None
                samples.append(value)
            return ((samples[0] * (1 - tx) + samples[1] * tx) * (1 - ty)
                    + (samples[2] * (1 - tx) + samples[3] * tx) * ty)
        print(f"[RingRadialProbe] center_m=({cx:.3f},{cy:.3f}) "
              "radius_m mean_m angular_sd_m min_m max_m samples")
        for radius in range(1, args.probe_max_radius + 1):
            row = [at_world(cx + radius * math.cos(a * math.tau / 32),
                            cy + radius * math.sin(a * math.tau / 32))
                   for a in range(32)]
            row = [value for value in row if value is not None]
            if row:
                mean = sum(row) / len(row)
                sd = math.sqrt(sum((value - mean) ** 2 for value in row) / len(row))
                print(f"  {radius:>2} {mean:.6f} {sd:.6f} {min(row):.6f} "
                      f"{max(row):.6f} {len(row)}")
    if args.layout:
        layout = json.load(open(args.layout, encoding="utf-8"))
        if next(item for item in layout["strates"] if item["order"] == 0)["transition_mode"] != "Gradient":
            raise SystemExit("--layout must describe a Gradient owner interface")
        errors = []
        for index, actual in enumerate(values):
            if actual is None:
                continue
            px, py = index % width, index // width
            x = args.bounds[0] + (px + 0.5) / sx
            y = args.bounds[3] - (py + 0.5) / sy
            expected = gradient_height_um(layout, round(x * 1_000_000),
                                          round(y * 1_000_000)) / 1_000_000.0
            errors.append(actual - expected)
        absolute = sorted(abs(error) for error in errors)
        rms = math.sqrt(sum(error * error for error in errors) / len(errors))
        print(f"[RingGradientMatch] samples={len(errors)} rms_m={rms:.6f} "
              f"p95_abs_m={absolute[round((len(absolute)-1)*.95)]:.6f} "
              f"max_abs_m={absolute[-1]:.6f} "
              "expected=owner-resolved-2d-hash-value-noise")


if __name__ == "__main__":
    main()