"""Export traceable viewer-only dent anchors and exact prepared GNN arrays.

Does not alter registration, source files, or training-validity masks.
"""
from pathlib import Path
import json
import numpy as np
from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
WEB = ROOT / 'datas_2/viewer/explorer_data'


def dump(name, data):
    (WEB / name).write_text(json.dumps(data, ensure_ascii=False, allow_nan=False, separators=(',', ':')))


def main():
    meta = json.loads((WEB / 'metadata.json').read_text())
    folders = next(p for p in (ROOT / 'datas_2').glob('1. PLY*') if p.is_dir())
    annotations = {}
    for batch, info in meta['images'].items():
        folder = next(p for p in folders.iterdir() if p.is_dir() and ('16' if batch == '16h' else '18') in p.name)
        path = folder / 'all_merge_pointcloud_full.ply'
        with path.open('rb') as stream:
            header = []
            while True:
                line = stream.readline().strip()
                header.append(line)
                if line == b'end_header':
                    break
                if not line:
                    raise ValueError('Incomplete PLY header')
            offset = stream.tell()
        assert header[1] == b'format binary_little_endian 1.0'
        assert [line for line in header if line.startswith(b'property')] == [
            b'property float x', b'property float y', b'property float z',
            b'property uchar red', b'property uchar green', b'property uchar blue']
        cloud = np.memmap(path, offset=offset, dtype=[('x', '<f4'), ('y', '<f4'), ('z', '<f4'), ('rgb', 'u1', 3)])
        height = np.asarray(Image.open(folder / 'Merged_Height.tif'))
        assert list(height.shape[::-1]) == info['size']
        x0, y0 = float(cloud['x'].min()), float(cloud['y'].max())
        sample = np.stack([cloud[k][::301] for k in ('x', 'y', 'z')], axis=1)
        u = np.rint((sample[:, 0] - x0) / .1).astype(int)
        v = np.rint((y0 - sample[:, 1]) / .1).astype(int)
        assert ((u >= 0) & (u < height.shape[1]) & (v >= 0) & (v < height.shape[0])).all()
        z = height[v, u]
        assert (z > -900).all()
        error = np.abs(z - sample[:, 2])
        assert np.median(error) < .02 and np.quantile(error, .9) < .2
        px, py = info['dent_pixel']
        xyz = np.array([x0 + px * .1, y0 - py * .1, float(height[py, px])])
        assert xyz[2] > -900
        # Snap only to an actual measured point near the reconstructed pixel.
        candidates = np.flatnonzero((abs(cloud['x'] - xyz[0]) < .3) & (abs(cloud['y'] - xyz[1]) < .3))
        assert len(candidates)
        points = np.stack([cloud[k][candidates] for k in ('x', 'y', 'z')], axis=1)
        nearest = np.argmin(np.linalg.norm(points - xyz, axis=1))
        delta = float(np.linalg.norm(points[nearest] - xyz))
        assert delta < .2
        annotations[batch] = {
            'id': 1, 'label': info['label'], 'image_pixel': info['dent_pixel'],
            'raw_xyz_mm': points[nearest].tolist(), 'raw_ply_vertex': int(candidates[nearest]),
            'raster_xyz_mm': xyz.tolist(), 'snap_distance_mm': delta,
            'mapping': {'x_mm': 'min(PLY.x) + 0.1 * pixel_x', 'y_mm': 'max(PLY.y) - 0.1 * pixel_y',
                        'z_mm': 'Merged_Height.tif[pixel_y, pixel_x]', 'origin_raw_xy_mm': [x0, y0]},
            'validation': {'sample_points': len(sample), 'valid_raster_fraction': 1.,
                           'height_abs_error_median_mm': float(np.median(error)),
                           'height_abs_error_p90_mm': float(np.quantile(error, .9))},
            'status': 'Inferred raster-to-PLY mapping, checked against measured heights; CSV top-left pixel origin still unconfirmed. Viewer reference only, not a validated GNN target.'}
        print(batch, annotations[batch]['raw_xyz_mm'], annotations[batch]['validation'], 'snap', delta)
    dump('dent_annotations.json', annotations)

    graph = np.load(WEB / 'cad_graph.npz')
    display = json.loads((WEB / 'design_graph.json').read_text())
    assert np.allclose(graph['pos_mm'].ravel(), display['positions'], atol=1e-3)
    def values(a):
        if a.dtype.kind == 'f':
            return [float(v) if np.isfinite(v) else None for v in a.ravel()]
        return a.tolist()
    data = {'source': 'cad_graph.npz', 'node_count': len(graph['pos_mm']),
            'edge_index': graph['edge_index'].tolist(),
            'edge_length_mm': values(graph['edge_length_mm']),
            'k1_per_mm': values(graph['k1_per_mm']), 'k2_per_mm': values(graph['k2_per_mm']),
            'curvature_valid': values(graph['curvature_valid']), 'observations': {}}
    for pose in meta['variants']:
        obs = np.load(WEB / f'{pose}_graph_observations.npz')
        assert not obs['training_valid'].any()
        data['observations'][pose] = {'source': f'{pose}_graph_observations.npz', **{
            key: values(obs[key]) for key in ('cad_to_scan_distance_mm', 'coverage_within5', 'training_valid')}}
    dump('graph_prepared.json', data)
    print('Exported exact NPZ graph arrays:', data['node_count'], 'nodes,', graph['edge_index'].shape[1], 'directed edges')


if __name__ == '__main__':
    main()