import os
import sys
import numpy as np
import trimesh as tm
from stl import mesh

# os.environ["CUDA_VISIBLE_DEVICES"] = '1'
# import tensorflow as tf

# gpus = tf.config.experimental.list_physical_devices('GPU')
# if gpus:
#     try:
#         # Currently, memory growth needs to be the same across GPUs
#         for gpu in gpus:
#             tf.config.experimental.set_memory_growth(gpu, True)
#         logical_gpus = tf.config.experimental.list_logical_devices('GPU')
#         print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
#     except RuntimeError as e:
#         # Memory growth must be set before GPUs have been initialized
#         print(e)

# from losses import collision_loss, normal_loss


# from losses import edge_loss_skirt as edge_loss
# from losses import edge_loss


if __name__ == '__main__':
    frame = '04173_193' # '02276_232', '02978_242', '04173_193'
    folder = 'examples/sampling_analysis'
    # original_folder = '/home/workspace/data/dataset/local-cache/train_skirt'
    original_folder = '/home/workspace/local-data/train'
    # folder = 'examples/sampling'

    id = frame.split('_')[0]
    mask = np.load(os.path.join(original_folder, 'mesh_mask', f'{id}.npy'))
    gt = np.load(os.path.join(folder, frame, f'{frame}_gt.npy'))
    faces = np.load(os.path.join(folder, frame, f'{frame}_faces.npy'))
    edges = tm.Trimesh(vertices=gt, faces=faces, validate=False, process=False).edges
    edges_gt = np.linalg.norm(gt[edges[:, 0]] - gt[edges[:, 1]], axis=1)
    area_gt = np.sum(np.linalg.norm(np.cross(gt[faces[:, 0]] - gt[faces[:, 1]], gt[faces[:, 0]] - gt[faces[:, 2]]), axis=1))

    errors = []
    vars = []
    area_diffs = []
    # for i in [1,2,4,10]:
    for i in range(15):
        sample = np.load(os.path.join(folder, frame, f'{frame}_{i}.npy'))
        edges_sample = np.linalg.norm(sample[edges[:, 0]] - sample[edges[:, 1]], axis=1)
        error = np.mean(np.abs(edges_sample - edges_gt))
        var = np.var(np.abs(edges_sample - edges_gt))
        errors.append(error)
        vars.append(var)

        area_sample = np.sum(np.linalg.norm(np.cross(sample[faces[:, 0]] - sample[faces[:, 1]], sample[faces[:, 0]] - sample[faces[:, 2]]), axis=1))
        area_diff = np.abs(area_sample - area_gt)
        area_diffs.append(area_diff)

    print(np.array(errors).mean())
    print(np.array(vars).mean())
    print(np.array(area_diffs).mean())
    print(area_gt)