# This file stores the datasets' metadata

import json
from util import *
from copy import deepcopy
import shutil
import os
import time

DATA_PATH = "../data/"

class Dataset:
	
	# path is relative to DATA_PATH
	# sizes are stored in most-significant-index-first order (row-major)
	# dtype is the NumPy string representation of the data type of the stored array
	# original_dtype optionally indicates a (more restrictive) domain of the original data than the domain of dtype
	def __init__(self, name, relative_path, sizes, dtype, original_dtype=None, most_constant_mode=0, spatial_modes=[1, 2]):
		self.name = name
		self.relative_path = relative_path
		self.path = DATA_PATH + relative_path
		self.sizes = sizes
		self.modes = len(sizes)
		self.dtype = dtype
		self.original_dtype=original_dtype
		self.most_constant_mode = most_constant_mode
		self.spatial_modes = spatial_modes
	
	def get_data_norm(self):
		with open(DATA_PATH + "data_properties.json", "r") as f:
			data_properties = json.load(f)
		if self.name not in data_properties.keys():
			raise Exception("No data found for %s in data_properties.json"%self.name)
		elif "data_norm" not in data_properties[self.name].keys():
			raise Exception("No data norm found for %s in data_properties.json"%self.name)
		else:
			return data_properties[self.name]["data_norm"]
	
	def size(self):
		product = 1
		for size in self.sizes:
			product *= size
		return product

def get_scratch_dataset(dataset, parallel_file_system=False):
	# If we're on the VSC, makes sure there's currently a valid copy of the dataset on the parallel file system and use that instead
	# Because multiple processes can start looking for the copy at the same time, there are some synchronization considerations to be made
	# However, we assume that shutil.copy handles locking properly so no two copies can commence at the same time
	if on_vsc():
		if parallel_file_system:
			scratch_path = "../tmp/data/"
		else:
			scratch_path = VSC_LOCAL_SCRATCH + "data/"
			os.makedirs(scratch_path, exist_ok=True)
		scratch_path += dataset.name
		while not os.path.isfile(scratch_path) or os.path.getsize(scratch_path) != os.path.getsize(dataset.path):
			# Dataset is not (fully) copied yet
			# First, check if there is not already an ongoing copy operation
			if not os.path.isfile(scratch_path):
				# No other copy operation, so start a new one
				shutil.copy(dataset.path, scratch_path)
			else:
				# Wait a bit, then check to see if there was any progress on the copy
				prev_size = os.path.getsize(scratch_path)
				time.sleep(1)
				current_size = os.path.getsize(scratch_path)
				if current_size == prev_size and current_size < os.path.getsize(dataset.path):
					# No progress made, probably a partial copy left behind by an aborted program
					shutil.copy(dataset.path, scratch_path)
		scratch_dataset = deepcopy(dataset)
		scratch_dataset.path = scratch_path
		return scratch_dataset
	else:
		return dataset

# Test data sets
test_data_float32_4D = Dataset("test_data_float_4x8x16x32", "test_data_4x8x16x32.raw", [4, 8, 16, 32], "float32")
test_data_float64_4D = Dataset("test_data_double_4x8x16x32", "test_data_double_4x8x16x32.raw", [4, 8, 16, 32], "float64")
test_data_float64_5D = Dataset("test_data_double_5x6x7x8x9", "test_data_double_5x6x7x8x9.raw", [5, 6, 7, 8, 9], "float64")

# Real data sets
moffett_field = Dataset("Moffett-Field", "hyperspectral_images/moffett_field_float32", [512, 1024, 224], "float32", original_dtype="int16", most_constant_mode=2, spatial_modes=[0, 1])
foot = Dataset("Foot", "volumes/foot_float32", [256, 256, 256], "float32", original_dtype="uint8")
engine = Dataset("Engine", "volumes/engine_float32", [256, 256, 128], "float32", original_dtype="uint8")
brain = Dataset("Brain", "dti/preprocessed/brain", [160, 190, 148, 6], "float32", most_constant_mode=2, spatial_modes=[0, 1])
isotropic_fine_pressure = Dataset("Isotropic1024-fine-pressure-512", "turbulence/isotropic1024fine-pressure-512", [512, 512, 512], "float32")
isotropic_fine_pressure_time_64 = Dataset("Isotropic1024-fine-pressure-time-64", "turbulence/isotropic1024fine-pressure-time-64", [64, 64, 64, 64], "float32", most_constant_mode=0)
isotropic_fine_pressure_time = Dataset("Isotropic-PT", "turbulence/isotropic1024fine-pressure-time-128", [100, 128, 128, 128], "float32", most_constant_mode=0)
isotropic_fine_velocity = Dataset("Isotropic-V", "turbulence/isotropic1024fine-velocity-512", [512, 512, 512, 3], "float32")
isotropic_fine_velocity_time = Dataset("Isotropic1024-fine-velocity-time-64", "turbulence/isotropic1024fine-velocity-time-64", [64, 64, 64, 64, 3], "float32", most_constant_mode=0)
deforest_globe_8 = Dataset("Deforest-8", "climate/preprocessed/deforest-globe-8-change", [20, 79, 8, 180, 360], "float32", most_constant_mode=1, spatial_modes=[3, 4])
deforest_globe_33 = Dataset("Deforest-33", "climate/preprocessed/deforest-globe-33-change", [19, 79, 33, 180, 360], "float32", most_constant_mode=1, spatial_modes=[3, 4])
hurricane = Dataset("Hurricane", "weather/preprocessed/hurricane-isabel", [13, 20, 100, 500, 500], "float32", most_constant_mode=1, spatial_modes=[3, 4])
