import trimesh
import numpy as np
import os
from tqdm import tqdm

def compute_smallest_pc(meshdir):
    smallest_pc_list = []
    index_list = []
    # iterate over all meshes in meshdir
    for meshname in tqdm(os.listdir(meshdir)):
        meshpath = os.path.join(meshdir, meshname)
        # extract index from meshname
        index = int(meshname.split("_")[-1].split(".")[0])
        index_list.append(index)
        mesh = trimesh.load(meshpath)
        # sample from mesh
        points = mesh.sample(1000)
        # compute covariance matrix
        cov = np.cov(points.T)
        # compute eigendecomposition
        eigvals, eigvecs = np.linalg.eigh(cov)
        # find smallest principal component
        smallest_pc = eigvecs[:,0]
        smallest_pc_list.append(smallest_pc)
        
    return smallest_pc_list, index_list
        

if __name__ == "__main__":
    meshdir = "/home/ubuntu/orienter-3d/all_upright_airplanes/airplane"
    smallest_pc_list, index_list = compute_smallest_pc(meshdir)
    # Target smallest pc is [0,1,0]
    target_smallest_pc = np.array([0,1,0])
    # Compute absolute cosine similarity of each smallest principal component with the mean
    similarities = []
    for pc in smallest_pc_list:
        similarity = np.abs(np.dot(pc, target_smallest_pc) / (np.linalg.norm(pc) * np.linalg.norm(target_smallest_pc)))
        similarities.append(similarity)
    similarities = np.array(similarities)
    print("Mean abs cos sim wrt target:", np.mean(similarities))
    print("Min abs cos sim wrt target:", np.min(similarities))
    # save the cos sims
    np.save("abs_cos_sims.npy", similarities)
    # save the index list
    np.save("index_list.npy", index_list)