import numpy as np
from math import sqrt, log10, inf
from subprocess import Popen, PIPE, STDOUT
from multiprocessing import cpu_count
from functools import reduce
import os

import datasets

if "TMPDIR" in os.environ:
	VSC_LOCAL_SCRATCH = os.environ["TMPDIR"] + "/vsc32681-compression-comparison/"
	os.makedirs(VSC_LOCAL_SCRATCH, exist_ok=True)

# Loads data into a flat array
# original_dtype indicates that the file might currently be stored using dtype,
# but should be cast to original_dtype since its original values fall within that domain
def load_data(path, dtype, original_dtype=None):
	data = np.fromfile(path, dtype)
	if original_dtype is None:
		return data
	elif original_dtype not in ("float32", "float64"):
		# Original data type was an integer type, so round first
		data = np.around(data)
	# If a separate original data type was passed, cast to that data type
	return data.astype(original_dtype)

# Short-hand version of load_data
def load_dataset(dataset):
	return load_data(dataset.path, dataset.dtype, original_dtype=dataset.original_dtype)

# Computes the range of an array
def compute_data_range(data):
	return np.amax(data) - np.amin(data)

# Computes RMSE in between two arrays
def compute_rmse(original, decompressed):
	if original.size != decompressed.size:
		raise Exception("Cannot compare two arrays of different size!")
	return np.linalg.norm(original - decompressed)/sqrt(original.size)

# Computes PSNR in between two arrays
def compute_psnr(original, decompressed, data_range):
	if original.size != decompressed.size:
		raise Exception("Cannot compare two arrays of different size!")
	rmse = compute_rmse(original, decompressed)
	return 20*log10(compute_data_range(original)/rmse) if rmse > 0 else inf

def psnr_to_relative_error(psnr, data_range, data_size, data_norm):
	return data_range*sqrt(data_size)/(data_norm*10**(psnr/20))

def relative_error_to_psnr(rel_error, data_range, data_size, data_norm):
	return 20*log10(data_range/((rel_error*data_norm)/sqrt(data_size)))

# Use /usr/bin/time to measure both CPU time, real time and peak memory usage
def measure_time_and_memory_usage(command, env=None, stdout_file_descriptor=STDOUT, measure_time=True, measure_memory=True):
	
	# compressor output will go to stdout
	# time output will go to stderr
	p = Popen(["/usr/bin/time", "-v"] + command, stdout=stdout_file_descriptor, stderr=PIPE, env=env)
	p.wait()
	
	# Parse time output and store in results
	results = {}
	keyword_tuples = []
	if measure_time:
		keyword_tuples.extend([
			["real_time", "Elapsed (wall clock) time (h:mm:ss or m:ss): ", lambda token: reduce(lambda higher_unit, lower_unit: 60*float(higher_unit) + float(lower_unit), token.split(":"))], # Parser reduces h:mm:ss or m:ss to seconds
			["user_time", "User time (seconds): ", float],
			["system_time", "System time (seconds): ", float]
		])
	if measure_memory:
		keyword_tuples.append(["peak_memory", "Maximum resident set size (kbytes): ", lambda token: 1024*int(token)]) # time returns MRSS in kilobytes
	error_output = p.stderr.read().decode()
	for line in error_output.split("\n"):
		for field_name, keyword, parser in keyword_tuples:
			if keyword in line:
				token = line.split(keyword)[1].strip()
				results[field_name] = parser(token)
				break
	
	return results

def get_max_threads():
	return cpu_count()

def on_vsc():
	return "VSC_SCRATCH" in os.environ
