import numpy as np
import trimesh
import pandas as pd
from scipy.stats import entropy
from skimage import color
from sklearn.preprocessing import scale
from nss_functions import * 
import os 

def get_feature_vector(objpath):
  print("Begin loading mesh.")   
  mesh = trimesh.load(objpath)
  
  print("Begin geometry feature extraction.")
  # extract geometry features
  to_origin, extents = trimesh.bounds.oriented_bounds(mesh)
  print("Mean curvature")
  r = max(extents) * 0.003
  #mean_curvature = np.array(trimesh.curvature.discrete_mean_curvature_measure(mesh, mesh.vertices, r))[0]
  mean_curvature = np.array(trimesh.curvature.discrete_gaussian_curvature_measure(mesh, mesh.vertices, r))
  print("Dihedral angle")
  dihedral_angle = mesh.face_adjacency_angles
  print("Face")
  face_area = mesh.area_faces
  face_angle = np.array([i for item in mesh.face_angles for i in item])
  
  print("Begin color feature extraction.")
  # extract color features
  rgb_color = mesh.visual.vertex_colors/255
  rgb_color = np.delete(rgb_color,3,axis=1)
  lab_color = color.rgb2lab(rgb_color)
  l = lab_color[:,0]
  a = lab_color[:,1]
  b = lab_color[:,2]
  
  
  print("Begin NSS parameters estimation.")
  # computer nss parameters
  nss_params = []
  i = 0
  for tmp in [l,a,b,mean_curvature,dihedral_angle,face_angle,face_area]:
      i = i + 1
      print(i)
      params = get_nss_param(tmp)
      #flatten the feature vector
      nss_params = nss_params + [i for item in params for i in item]
  return nss_params

objpath = "Chameleon_SimpMesh_4.obj"
features = get_feature_vector(objpath)
print(features)