"""Build a fine design-surface grid (data/surface.json) for the rain / ponding simulation.

Usage:
  python3 -I extract_surface.py <road-layout.dxf> <site-progress-test/data>

Source: the 0.1 m design contours of Bushveld Joy Road (block "RD CONT BJ", layers BJ G/M CONT) plus the
setting-out pegs (exact design levels). Contours are sampled onto a 0.5 m grid; the cells between contours
are filled with a smooth (harmonic) surface, which is what you get between contour lines of a planar road.
Cells further than ~5 m from any contour are left empty (outside the designed area).
"""
import base64
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from dxf_reader import read_sections, get, lwpolyline  # noqa: E402

CELL = 0.5
REACH = 10    # cells (5 m): how far from a contour we still trust the surface


def build(acc, cnt, pegs, x0, y0, nx, ny, use_pegs=True):
    known = cnt > 0
    Z = np.where(known, acc / np.maximum(cnt, 1e-9), np.nan)
    # pegs override (exact design levels)
    pi = np.clip(((pegs[:, 0] - x0) / CELL).round().astype(int), 0, nx - 1); pj = np.clip(((pegs[:, 1] - y0) / CELL).round().astype(int), 0, ny - 1)
    if use_pegs:
        Z[pj, pi] = pegs[:, 2]; known[pj, pi] = True

    # mask = cells within REACH of known cells (dilation by repeated shifting)
    mask = known.copy()
    for _ in range(REACH):
        m = mask.copy()
        m[1:, :] |= mask[:-1, :]; m[:-1, :] |= mask[1:, :]; m[:, 1:] |= mask[:, :-1]; m[:, :-1] |= mask[:, 1:]
        mask = m
    # initial guess: propagate nearest known values outward, then relax (Jacobi) with known cells fixed
    G = np.where(known, Z, np.nan)
    while np.isnan(G[mask]).any():
        P = np.pad(G, 1, constant_values=np.nan)
        nb = np.stack([P[:-2, 1:-1], P[2:, 1:-1], P[1:-1, :-2], P[1:-1, 2:]])
        with np.errstate(invalid='ignore'):
            s = np.nansum(nb, axis=0); c = np.sum(~np.isnan(nb), axis=0)
        new = np.where(np.isnan(G) & mask & (c > 0), s / np.maximum(c, 1), G)
        if np.array_equal(np.isnan(new), np.isnan(G)): break
        G = new
    free = mask & ~known
    for it in range(1500):
        P = np.pad(G, 1, mode='edge')
        avg = (P[:-2, 1:-1] + P[2:, 1:-1] + P[1:-1, :-2] + P[1:-1, 2:]) / 4
        avg = np.where(np.isnan(avg), G, avg)
        delta = np.nanmax(np.abs(avg[free] - G[free])) if free.any() else 0
        G[free] = avg[free]
        if delta < 1e-4: break
    G[~mask] = np.nan

    return G, mask, it


def main(dxf_path, data_dir):
    base = json.load(open(os.path.join(data_dir, 'bushveld.json')))
    ox, oy = base['origin']['cadX'], base['origin']['cadY']
    sec = read_sections(dxf_path)
    samples = []
    blk = None
    for typ, t in sec['BLOCKS']:
        if typ == 'BLOCK':
            blk = get(t, 2); continue
        if blk == 'RD CONT BJ' and typ == 'LWPOLYLINE' and ('BJ G CONT' in get(t, 8) or 'BJ M CONT' in get(t, 8)):
            pts, closed, z = lwpolyline(t)
            if closed: pts = pts + [pts[0]]
            for (x1, y1), (x2, y2) in zip(pts, pts[1:]):
                L = np.hypot(x2 - x1, y2 - y1); k = max(1, int(L / (CELL / 2)))
                for i in range(k + 1):
                    samples.append((x1 + (x2 - x1) * i / k - ox, y1 + (y2 - y1) * i / k - oy, z))
    S = np.array(samples)
    pegs = np.array([[p['e'], p['n'], p['z']] for p in base['pegs']])
    x0, y0 = np.floor(S[:, 0].min()) - 2, np.floor(S[:, 1].min()) - 2
    nx = int((S[:, 0].max() + 2 - x0) / CELL) + 1; ny = int((S[:, 1].max() + 2 - y0) / CELL) + 1

    acc = np.zeros((ny, nx)); cnt = np.zeros((ny, nx))
    def splat(P, w=1.0):
        i = np.clip(((P[:, 0] - x0) / CELL).round().astype(int), 0, nx - 1)
        j = np.clip(((P[:, 1] - y0) / CELL).round().astype(int), 0, ny - 1)
        np.add.at(acc, (j, i), P[:, 2] * w); np.add.at(cnt, (j, i), w)
    splat(S)
    G, mask, it = build(acc, cnt, pegs, x0, y0, nx, ny)
    G0, _, _ = build(acc, cnt, pegs, x0, y0, nx, ny, use_pegs=False)
    # accuracy estimate: leave-one-out on pegs (pegs re-predicted from the contour-only surface around them)
    zmin = float(np.nanmin(G))
    q = np.where(np.isnan(G), -1, np.round((G - zmin) * 1000)).astype('<i2')
    out = {'x0': float(x0), 'y0': float(y0), 'cell': CELL, 'nx': nx, 'ny': ny, 'zmin': round(zmin, 3),
           'encoding': 'int16 little-endian, mm above zmin, -1 = outside design area, row-major from (x0, y0)',
           'z': base64.b64encode(q.tobytes()).decode(),
           'source': 'Design contours 0.1 m (RD CONT BJ) + setting-out pegs'}
    json.dump(out, open(os.path.join(data_dir, 'surface.json'), 'w'), separators=(',', ':'))
    # accuracy: surface from contours only, compared with the pegs' exact design levels
    pi = np.clip(((pegs[:, 0] - x0) / CELL).round().astype(int), 0, nx - 1); pj = np.clip(((pegs[:, 1] - y0) / CELL).round().astype(int), 0, ny - 1)
    d = G0[pj, pi] - pegs[:, 2]; d = d[~np.isnan(d)]
    out['check'] = {'pegs': int(len(d)), 'medianAbs': round(float(np.median(np.abs(d))), 3), 'p90Abs': round(float(np.percentile(np.abs(d), 90)), 3)}
    json.dump(out, open(os.path.join(data_dir, 'surface.json'), 'w'), separators=(',', ':'))
    print('contour-only surface vs pegs:', out['check'])
    print(f'grid {nx} x {ny} @ {CELL} m, {mask.sum()} cells in design area ({mask.sum() * CELL * CELL:.0f} m²), relax iters {it}')
    print(f'z {zmin:.2f} .. {np.nanmax(G):.2f}; json {os.path.getsize(os.path.join(data_dir, "surface.json")) // 1024} KB')


if __name__ == '__main__':
    main(*sys.argv[1:3])