from pathlib import Path
import json, sys, time
import numpy as np
import open3d as o3d
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

ROOT = Path(__file__).resolve().parents[1]
DATA = ROOT / 'datas_2'
OUT = ROOT / 'analysis' / 'results'
OUT.mkdir(parents=True, exist_ok=True)

def cloud(p):
    c = o3d.geometry.PointCloud()
    c.points = o3d.utility.Vector3dVector(p)
    return c

def main():
    cad = o3d.io.read_triangle_mesh(str(DATA/'viewer/part.stl'))
    cad.remove_duplicated_vertices()
    o3d.utility.random.seed(42)
    cp = np.asarray(cad.sample_points_uniformly(150000).points)
    np.save(OUT/'cad_sample.npy', cp)
    fig, axes = plt.subplots(3, 3, figsize=(16,12))
    records = {}
    for row, (name, pattern) in enumerate([('16h','*16*'), ('18h01','*18*')]):
        folder = next(next(p for p in DATA.glob('1. PLY*') if p.is_dir()).glob(pattern))
        record = {}
        for path in sorted(folder.glob('*.ply')):
            cache = OUT/f'{name}_{path.stem}.npz'
            if cache.exists():
                a = np.load(cache); p=a['points']; n=int(a['n_raw'])
            else:
                c = o3d.io.read_point_cloud(str(path)); n=len(c.points)
                d=c.voxel_down_sample(2.0); p=np.asarray(d.points)
                np.savez_compressed(cache, points=p, n_raw=n)
            record[path.name] = {'n_raw':n, 'n_down':len(p), 'min':p.min(0).tolist(), 'max':p.max(0).tolist()}
            print(name, path.name, record[path.name], flush=True)
            if path.stem == 'all_merge_pointcloud_full':
                for ax,(i,j) in zip(axes[row], [(0,1),(0,2),(1,2)]):
                    ax.scatter(p[:,i],p[:,j],s=.1); ax.set_aspect('equal');ax.set_title(f'{name} raw {"XYZ"[i]} / {"XYZ"[j]}')
        records[name]=record
    for ax,(i,j) in zip(axes[2], [(0,1),(0,2),(1,2)]):
        ax.scatter(cp[:,i],cp[:,j],s=.1,c='gray'); ax.set_aspect('equal');ax.set_title(f'CAD {"XYZ"[i]} / {"XYZ"[j]}')
    fig.tight_layout(); fig.savefig(OUT/'raw_geometry.png',dpi=140)
    (OUT/'inventory.json').write_text(json.dumps(records,indent=2))

if __name__ == '__main__': main()