import torch
import numpy as np
import os
import cv2
import glob
import trimesh
import pyrender
from PIL import Image

from lib.Renderer import Renderer_pyr


def get_cam_ssns(dataset_fpath):
    cam_ssns = []   # sorted according to their spatial position
    with open(os.path.join(dataset_fpath, 'cam_ssns.txt'), 'r') as fp:
        lns = fp.readlines()
    for ln in lns:
        ln = ln.strip().split(' ')
        if len(ln) > 0:
            cam_ssns.append(ln[0])
    return cam_ssns


class CalibRecorder():
    def __init__(self, data_folder,  param_folder):
        self.data_folder = data_folder
        self.camera_ids = ['01', '02', '04', '05']
        self.param_folder = param_folder
        os.makedirs(self.param_folder, exist_ok=True)
        # self.camera_ids = ['220700191', '221501007', '222200036', '222200037', '222200038', '222200039', '222200040', '222200041',
        #                    '222200042', '222200043', '222200044', '222200045', '222200046', '222200047', '222200048', '222200049']

    def log(self, log_data):
        eids = log_data['eids']
        landmarks_gt = log_data['landmarks_gt']
        landmarks_2d = log_data['landmarks_2d']
        bfm = log_data['bfm']
        intrinsics = log_data['intrinsics']
        extrinsics = log_data['extrinsics']
     
        with torch.no_grad():
            vertices0, landmarks0 = log_data['bfm']()
        # landmarks = torch.cat([landmarks0, vertices0[:, ::100, :]], 1)
        
        for n, eid in enumerate(eids):
            os.makedirs(os.path.join(self.param_folder, eid), exist_ok=True)
            bfm.save('%s/params.npz' % (os.path.join(self.param_folder, eid)), batch_id=n)
            np.save('%s/lmk_3d.npy' % (os.path.join(self.param_folder, eid)), landmarks0[n].cpu().numpy())
            np.save('%s/vertices.npy' % (os.path.join(self.param_folder, eid)), vertices0[n].cpu().numpy())
            # np.save('%s/keypoints_face.npy' % (os.path.join(self.param_folder, eid)), landmarks[n].cpu().numpy())

            
            # for v in range(0, 1):
            #     vertices = vertices0[n].cpu().numpy()
            #     faces = log_data['bfm'].faces.cpu().numpy()
            #     mesh_trimesh = trimesh.Trimesh(vertices=vertices, faces=faces)
            #     mesh = pyrender.Mesh.from_trimesh(mesh_trimesh)

            #     renderer = Renderer_pyr(intrinsic=intrinsics[n, v], extrinsic=extrinsics[n, v])
            #     render_image = renderer.render(mesh)

            #     cv2.imwrite('%s/vis_%d.jpg' % (os.path.join(self.param_folder, eid), v+1), render_image[:,:,::-1])