import os
import numpy as np
import torch

class ClassificationData:
	def __init__(self, data_dict):
		self.data_dict = data_dict
		self.pcs = self.find_attribute('pcs')
		self.labels = self.find_attribute('labels')
		self.check_data()

	def find_attribute(self, attribute):
		try:
			attribute_data = self.data_dict[attribute]
		except:
			print("Given data directory has no key attribute \"{}\"".format(attribute))
		return attribute_data

	def check_data(self):
		assert 1 < len(self.pcs.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.pcs.shape)
		assert 0 < len(self.labels.shape) < 3, "Error in dimension of labels! Given data dimension: {}".format(self.labels.shape)
		
		if len(self.pcs.shape)==2: self.pcs = self.pcs.reshape(1, -1, 3)
		if len(self.labels.shape) == 1: self.labels = self.labels.reshape(1, -1)

		assert self.pcs.shape[0] == self.labels.shape[0], "Inconsistency in the number of point clouds and number of ground truth labels!"


	def __len__(self):
		return self.pcs.shape[0]

	def __getitem__(self, index):
		return torch.tensor(self.pcs[index]).float(), torch.from_numpy(self.labels[idx]).type(torch.LongTensor)


class RegistrationData:
	def __init__(self, data_dict):
		self.data_dict = data_dict
		self.template = self.find_attribute('template')
		self.source = self.find_attribute('source')
		self.transformation = self.find_attribute('transformation')
		self.check_data()

	# def find_attribute(self, attribute):
	# 	try:
	# 		attribute_data = self.data[attribute]
	# 	except:
	# 		print("Given data directory has no key attribute \"{}\"".format(attribute))
	# 	return attribute_data
	
	def find_attribute(self, attribute):
		attribute_data = None
		if attribute in self.data_dict:
			attribute_data = self.data_dict[attribute]
		else:
			print("Given data directory has no key attribute \"{}\"".format(attribute))
		return attribute_data

	def check_data(self):
		assert 1 < len(self.template.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.template.shape)
		assert 1 < len(self.source.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.source.shape)
		assert 1 < len(self.transformation.shape) < 4, "Error in dimension of transformations! Given data dimension: {}".format(self.transformation.shape)

		if len(self.template.shape)==2: self.template = self.template.reshape(1, -1, 3)
		if len(self.source.shape)==2: self.source = self.source.reshape(1, -1, 3)
		if len(self.transformation.shape) == 2: self.transformation = self.transformation.reshape(1, 4, 4)

		assert self.template.shape[0] == self.source.shape[0], "Inconsistency in the number of template and source point clouds!"
		assert self.source.shape[0] == self.transformation.shape[0], "Inconsistency in the number of transformation and source point clouds!"

	def __len__(self):
		return self.template.shape[0]

	def __getitem__(self, index):
		return torch.tensor(self.template[index]).float(), torch.tensor(self.source[index]).float(), torch.tensor(self.transformation[index]).float()


class FlowData:
	def __init__(self, data_dict):
		self.data_dict = data_dict
		self.frame1 = self.find_attribute('frame1')
		self.frame2 = self.find_attribute('frame2')
		self.flow = self.find_attribute('flow')
		self.check_data()

	def find_attribute(self, attribute):
		try:
			attribute_data = self.data[attribute]
		except:
			print("Given data directory has no key attribute \"{}\"".format(attribute))
		return attribute_data

	def check_data(self):
		assert 1 < len(self.frame1.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.frame1.shape)
		assert 1 < len(self.frame2.shape) < 4, "Error in dimension of point clouds! Given data dimension: {}".format(self.frame2.shape)
		assert 1 < len(self.flow.shape) < 4, "Error in dimension of flow! Given data dimension: {}".format(self.flow.shape)

		if len(self.frame1.shape)==2: self.frame1 = self.frame1.reshape(1, -1, 3)
		if len(self.frame2.shape)==2: self.frame2 = self.frame2.reshape(1, -1, 3)
		if len(self.flow.shape) == 2: self.flow = self.flow.reshape(1, -1, 3)

		assert self.frame1.shape[0] == self.frame2.shape[0], "Inconsistency in the number of frame1 and frame2 point clouds!"
		assert self.frame2.shape[0] == self.flow.shape[0], "Inconsistency in the number of flow and frame2 point clouds!"

	def __len__(self):
		return self.frame1.shape[0]

	def __getitem__(self, index):
		return torch.tensor(self.frame1[index]).float(), torch.tensor(self.frame2[index]).float(), torch.tensor(self.flow[index]).float()


class UserData:
	def __init__(self, application, data_dict):
		self.application = application

		if self.application == 'classification':
			self.data_class = ClassificationData(data_dict)
		elif self.application == 'registration':
			self.data_class = RegistrationData(data_dict)
		elif self.application == 'flow_estimation':
			self.data_class = FlowData(data_dict)

	def __len__(self):
		return len(self.data_class)

	def __getitem__(self, index):
		return self.data_class[index]