import torch
import torchvision as tv
import numpy as np
import glob
import os
import random
import cv2
from skimage import io
from torch.utils.data import Dataset


class LandmarkDataset(Dataset):
    def __init__(self, data_folder):
        super(LandmarkDataset, self).__init__()

        self.eids = sorted(os.listdir(data_folder))
        self.images_folders = [os.path.join(data_folder, eid) for eid in self.eids]
        self.camera_ids = ['01', '02', '04', '05']

    def get_item(self):
        landmarks = []
        extrinsics = []
        intrinsics = []
        for images_folder in self.images_folders:
            landmarks_ = []
            extrinsics_ = []
            intrinsics_ = []
            for v in range(len(self.camera_ids)):
                if os.path.exists(images_folder + '/landmarks_%s.npy' % self.camera_ids[v]):
                    landmark = np.load(images_folder + '/landmarks_%s.npy' % self.camera_ids[v])
                    landmark = np.vstack([landmark[0:48], landmark[49:54], landmark[55:68]])
                    extrinsic = np.load(images_folder + '/camera_%s.npz' % self.camera_ids[v])['extrinsic']
                    intrinsic = np.load(images_folder + '/camera_%s.npz' % self.camera_ids[v])['intrinsic']
                else:
                    landmark = np.zeros([68, 3], dtype=np.float32)
                    extrinsic = np.ones([3, 4], dtype=np.float32)
                    intrinsic = np.ones([3, 3], dtype=np.float32)
                landmarks_.append(landmark)
                extrinsics_.append(extrinsic)
                intrinsics_.append(intrinsic)
            landmarks_ = np.stack(landmarks_)
            extrinsics_ = np.stack(extrinsics_)
            intrinsics_ = np.stack(intrinsics_)
            landmarks.append(landmarks_)
            extrinsics.append(extrinsics_)
            intrinsics.append(intrinsics_)
        landmarks = np.stack(landmarks)
        extrinsics = np.stack(extrinsics)
        intrinsics = np.stack(intrinsics)

        return landmarks, extrinsics, intrinsics, self.eids
    
    def __len__(self):
        return len(self.images_folders)