#!/usr/bin/python3
# -*- encoding: utf-8 -*-
"""
Reconstruct a mesh from a fused point cloud with normals using TSDF and Marching Cubes.
Uses a vectorized splatting approach for efficient SDF generation.

Install:
  pip install numpy scikit-image plyfile tqdm argparse

Example usage:
  python3 MvsPointCloud2TSDF.py -i input_cloud.ply [-o output_mesh.ply] [--voxel_size VOXEL_SIZE] [--truncation_mult TRUNCATION_MULT]
"""

import numpy as np
import argparse
import os
from plyfile import PlyData, PlyElement
from skimage.measure import marching_cubes
from tqdm import tqdm
from scipy.spatial import cKDTree

def load_ply(path):
    """
    Load points and normals from a PLY file.
    """
    plydata = PlyData.read(path)
    vertex = plydata['vertex']
    
    points = np.stack([vertex['x'], vertex['y'], vertex['z']], axis=-1)
    
    if 'nx' in vertex.data.dtype.names and 'ny' in vertex.data.dtype.names and 'nz' in vertex.data.dtype.names:
        normals = np.stack([vertex['nx'], vertex['ny'], vertex['nz']], axis=-1)
    else:
        raise ValueError("PLY file must contain normals (nx, ny, nz).")
        
    return points, normals

def estimate_voxel_size(points, num_samples=3000):
    """
    Estimate voxel size based on nearest neighbor distances of a subset of points.
    Using a brute-force approach on a small sample to avoid scipy dependency if possible,
    but simplest is just to take 1% of bounding box diagonal or similar heuristic
    if we want to avoid KDTree/scipy completely. 
    However, decent estimation requires spatial awareness.
    
    Let's use a simple heuristic for now: 
    Average distance to nearest neighbor in a small random subset.
    """
    print("Estimating voxel size...")
    if len(points) > num_samples:
        idx = np.random.choice(len(points), num_samples, replace=False)
        sample = points[idx]
    else:
        sample = points

    # Brute force NN for estimation (fast enough for 1000 points)
    # dist matrix: (N, N)
    dists = np.sqrt(np.sum((sample[:, None, :] - sample[None, :, :]) ** 2, axis=-1))
    np.fill_diagonal(dists, np.inf)
    min_dists = np.min(dists, axis=1)
    
    return np.median(min_dists)

def splat_points_to_tsdf(points, normals, voxel_size, truncation_mult=4.0, chunk_size=10000):
    """
    Splat points into a TSDF volume using sparse voxel hashing concept (dictionary).
      - voxel_size: size of each voxel in scene units (0 for auto-estimation)
      - truncation_mult: voxel size multiplier to set the truncation value for signed distance function
      - chunk_size: to vectorize, we can process points in chunks, adjust based on memory
    """
    if voxel_size  0.
            sdf_vals = np.sum(diff * nrms_chunk, axis=1)
            
            # Check truncation
            valid_mask = np.abs(sdf_vals) < truncation
            
            if not np.any(valid_mask):
                continue
                
            valid_indices = cand_vox_indices[valid_mask]
            valid_sdfs = sdf_vals[valid_mask]
            
            # Update TSDF
            # We can't easily vector-update a simple dict. 
            # But we can create a local hash map/list and merge?
            # Or just loop for the valid ones (should be smaller subset)
            
            # Optimization: Weighting
            # weight = 1.0 (Simple)
            weights = np.ones_like(valid_sdfs)
            
            # We need to aggregate. 
            # Since we are in Python, dict access is slow in a tight loop.
            # Faster approach: Store all updates in a list/arrays and aggregate later.
            # But memory might be an issue.
            
            # Let's try to aggregate locally in chunk then update global dict?
            # Or use a flat array of 'hashed' indices if domain is known?
            # Since we don't know the full domain size perfectly without allocating dense grid,
            # let's stick to dictionary but maybe use a flat key?
            
            # Flatten keys for dictionary
            keys = tuple(map(tuple, valid_indices))
            
            for k, sdf, w in zip(keys, valid_sdfs, weights):
                if k in tsdf_vol:
                    tsdf_vol[k][0] += w
                    tsdf_vol[k][1] += sdf * w
                else:
                    tsdf_vol[k] = [w, sdf * w] # [weight, weighted_sdf]

    return tsdf_vol, min_bound, voxel_size

def compute_tsdf_voxel(points, normals, voxel_size, truncation_mult=4.0, chunk_size=1000000):
    """
    Compute TSDF by creating a dense grid and querying nearest neighbors using KDTree.
    This is the "per-voxel" iteration method.
    """
    if voxel_size