"""Diagnostic replay of saved tensors; capture or substitute pre-simplification remesh."""
from pathlib import Path
import json
import numpy as np


def replay(raw_path, layout_path, output, filtered_path=None):
    import torch
    import cumesh
    import o_voxel
    output = Path(output)
    output.mkdir(parents=True, exist_ok=True)
    with np.load(raw_path, allow_pickle=False) as data:
        arrays = {key: data[key].copy() for key in data.files}
    layout = {key: slice(*value) for key, value in json.loads(Path(layout_path).read_text()).items()}
    original = cumesh.remeshing.remesh_narrow_band_dc
    class Captured(Exception):
        pass
    def interception(*args, **kwargs):
        if filtered_path:
            with np.load(filtered_path, allow_pickle=False) as data:
                return (torch.from_numpy(data['vertices'].astype(np.float32)).cuda(),
                        torch.from_numpy(data['faces'].astype(np.int32)).cuda())
        vertices, faces = original(*args, **kwargs)
        np.savez_compressed(output / 'pre-simplify.npz', vertices=vertices.cpu().numpy(), faces=faces.cpu().numpy())
        raise Captured()
    cumesh.remeshing.remesh_narrow_band_dc = interception
    try:
        result = o_voxel.postprocess.to_glb(
            vertices=torch.from_numpy(arrays['vertices']).cuda(), faces=torch.from_numpy(arrays['faces']).cuda(),
            attr_volume=torch.from_numpy(arrays['attrs']).cuda(), coords=torch.from_numpy(arrays['coords']).cuda(),
            attr_layout=layout, grid_size=int(arrays['grid_size']) if arrays['grid_size'].ndim == 0
            else arrays['grid_size'].tolist(), aabb=[[-.5, -.5, -.5], [.5, .5, .5]],
            decimation_target=50000, texture_size=1024, remesh=True, remesh_band=1, remesh_project=0)
        result.export(output / 'filtered-remesh.glb')
    except Captured:
        return str(output / 'pre-simplify.npz')
    finally:
        cumesh.remeshing.remesh_narrow_band_dc = original
    return str(output / 'filtered-remesh.glb')