# Tests if all compressors are working properly

from datasets import *
from compressors import *
from util import *
import sys
from copy import deepcopy

make_temp_paths_unique("verify-setup")

# Check if each compressor can compress and decompress a given dataset within a reasonable margin of the target error
def verify_for_dataset(dataset, threads=1, strict=False, repeat=1, selected_compressors=compressors):
	
	print("Start of test with dataset: %s"%dataset.name)
	
	success = True
	for compressor in selected_compressors:
		
		# Set-up for this compressor
		make_temp_paths_unique("verify-setup", parallel_file_system=compressor.needs_parallel_io)
		scratch_dataset = get_scratch_dataset(dataset, parallel_file_system=compressor.needs_parallel_io)
		original = load_data(scratch_dataset.path, scratch_dataset.dtype)
		data_norm = np.linalg.norm(original)
		data_range = compute_data_range(original)
		print_commands = True # For debugging
		target_psnr = 38 #relative_error_to_psnr(1e-1, data_range, dataset.size(), data_norm)
		
		for _ in range(repeat):
			
			target_error = None
			psnr = None
			for _ in range(5):
				
				# Processing
				target_error = compressor.estimate_target_error(target_psnr, data_range=data_range, dataset=scratch_dataset, prev_target_error=target_error, prev_psnr=psnr)
				cleanup()
				results = compress(compressor, scratch_dataset, target_error, print_commands=print_commands, threads=threads)
				decompress(compressor, scratch_dataset, print_commands=print_commands, threads=threads)
				
				# Check results
				decompressed = load_data(get_decompressed_path(), scratch_dataset.dtype)
				psnr = compute_psnr(original, decompressed, data_range)
				passed = (abs(target_psnr - psnr) < 0.01 if strict else 10)
				if passed:
					break
				else:
					print("Compressor %s resulted in %s dB PSNR (relative error: %s, compression factor: %s), attempting to improve target error estimate ..."%(compressor.name, psnr, psnr_to_relative_error(psnr, data_range, dataset.size(), data_norm), results["compression_factor"]))
			
			print("Compressor %s %s test: %s dB PSNR (relative error: %s, compression factor: %s)"%(compressor.name, "passed" if passed else "failed", psnr, psnr_to_relative_error(psnr, data_range, dataset.size(), data_norm), results["compression_factor"]))
			if not passed:
				success = False
			
			# Remove files so that they're not used by the following compressor
			cleanup()
	
	# Summarize result
	if success:
		print("All compressors passed the test!")
	else:
		print("Some compressors did not pass the test!")

def verify_mode_storage_orders(dataset, selected_compressors=compressors):
	
	print("Start of mode storage order test with dataset: %s"%dataset.name)
	
	success = True
	for compressor in selected_compressors:
	
		# Set-up for this compressor
		make_temp_paths_unique("verify-setup", parallel_file_system=compressor.needs_parallel_io)
		scratch_dataset = get_scratch_dataset(dataset, parallel_file_system=compressor.needs_parallel_io)
		original = load_data(scratch_dataset.path, scratch_dataset.dtype)
		data_norm = np.linalg.norm(original)
		data_range = compute_data_range(original)
		print_commands = True # For debugging
		target_psnr = relative_error_to_psnr(1e-4, data_range, dataset.size(), data_norm)
			
		# Processing
		target_error = compressor.estimate_target_error(target_psnr, data_range=data_range, dataset=scratch_dataset, prev_target_error=None, prev_psnr=None)
		cleanup()
		results = compress(compressor, scratch_dataset, target_error, print_commands=print_commands, threads=threads)
		decompress(compressor, scratch_dataset, print_commands=print_commands, threads=threads)
		
		# Check results
		decompressed = load_data(get_decompressed_path(), scratch_dataset.dtype)
		psnr = compute_psnr(original, decompressed, data_range)
		passed = (abs(target_psnr - psnr) < 10)
		print("Compressor %s %s test: %s dB PSNR (relative error: %s, compression factor: %s)"%(compressor.name, "passed" if passed else "failed", psnr, psnr_to_relative_error(psnr, data_range, dataset.size(), data_norm), results["compression_factor"]))
		if not passed:
			success = False
		
		if dataset.sizes != dataset.sizes[::-1]:
			print("Checking reverse mode storage order")
			original_compression_factor = results["compression_factor"]
			original_psnr = psnr
			transposed_dataset = deepcopy(scratch_dataset)
			transposed_dataset.sizes = transposed_dataset.sizes[::-1]
			cleanup()
			results = compress(compressor, transposed_dataset, target_error, print_commands=print_commands, threads=threads)
			decompress(compressor, transposed_dataset, print_commands=print_commands, threads=threads)
			decompressed = load_data(get_decompressed_path(), scratch_dataset.dtype)
			psnr = compute_psnr(original, decompressed, data_range)
			compression_factor = results["compression_factor"]
			if abs(psnr - original_psnr) < 10 and compression_factor > original_compression_factor:
				success = False
				print("Compressor %s resulted in compression factor %s at %s dB PSNR when reverting mode sizes, but only achieved compression factor %s at %s dB PSNR when using the original sizes! Maybe the mode interpretation for this compressor is reversed?"%(compressor.name, compression_factor, psnr, original_compression_factor, original_psnr))
			else:
				print("Compressor %s passed test, only resulted in compression factor %s at %s dB PSNR when reverting mode sizes, but achieved compression factor %s at %s dB PSNR when using the original sizes"%(compressor.name, compression_factor, psnr, original_compression_factor, original_psnr))
		
		# Remove files so that they're not used by the following compressor
		cleanup()
	
	# Summarize result
	if success:
		print("All compressors passed the test!")
	else:
		print("Some compressors did not pass the test!")

cli_args = deepcopy(sys.argv[1:])
if "--strict" in cli_args:
	strict = True
	cli_args.remove("--strict")
else:
	strict = False
threads = 1
if len(cli_args) > 0:
	threads = int(cli_args[0])
repeat = 1
if len(cli_args) > 1:
	repeat = int(cli_args[1])

no_x265 = [SZ, ZFP, FPZIP, TuckerMPI, TTHRESH, ATC]
tthresh_and_atc_variants = [TTHRESH] + ATC_ALL_OPT_LEVELS

for dataset in [isotropic_fine_pressure_time]:
	verify_for_dataset(dataset, threads=threads, strict=strict, repeat=repeat, selected_compressors=[generate_atc_compressor(rtmss=0.0), TTHRESH])
#verify_for_dataset(brain, threads=threads, strict=strict, selected_compressors=tthresh_and_atc_variants)
#verify_for_dataset(foot, strict=strict, threads=threads, repeat=repeat, selected_compressors=[generate_atc_compressor(optimization_level=2)])
#verify_for_dataset(test_data_float32_4D, threads=threads, selected_compressors=[SZ, ZFP, FPZIP, TTHRESH, ATC])
#verify_for_dataset(test_data_float64_4D, threads=threads, selected_compressors=[SZ, ZFP, FPZIP, TTHRESH, ATC])
#for dataset in [moffett_field, foot, engine, brain, isotropic_fine_pressure, isotropic_fine_pressure_time_64, isotropic_fine_pressure_time, isotropic_fine_velocity, isotropic_fine_velocity_time, deforest_globe_8]: # deforest_globe_33 and hurricane are too big for testing on the login nodes it seems
#	verify_for_dataset(dataset, threads=threads, selected_compressors=[ATC,])

#verify_for_dataset(moffett_field, threads=threads, strict=strict, selected_compressors=tthresh_and_atc_variants)
#verify_for_dataset(foot, threads=threads, selected_compressors=no_x265)
#verify_for_dataset(isotropic_fine_pressure_time, threads=threads, selected_compressors=[TuckerMPI])
#verify_for_dataset(isotropic_fine_velocity)
#verify_for_dataset(isotropic_fine_velocity_time)
