import os
import sys
import json
from math import inf, ceil, floor
import numpy as np
import gc

from compressors import *
from datasets import *

# Constants
LOGS_PATH = "../logs/"
RESULTS_PATH = "../results/"

def safe_json_file_update(path, data):
	temp_path = path + ".tmp"
	with open(temp_path, "w") as f:
		json.dump(data, f)
	os.rename(temp_path, path)

# Measure compression factor, compression error and some other metrics across a range of target errors for a particular dataset and compressor
def measure_rate_distortion(dataset, compressor, threads=1, only_check_for_missing_results=False):
	
	if not only_check_for_missing_results:
		print("Starting experiments with dataset %s and compressor %s"%(dataset.name, compressor.name))
		make_temp_paths_unique("measure-rate-distortion-%s-%s"%(dataset.name, compressor.name), parallel_file_system=compressor.needs_parallel_io)
		cleanup()
	
	# Set up JSON data
	# Results file
	results_file = RESULTS_PATH + "%s_%s.json"%(dataset.name, compressor.name.replace(" ", "_"))
	if not os.path.isfile(results_file):
		if only_check_for_missing_results:
			print("Missing results for dataset %s and compressor %s: results file does not exist"%(dataset.name, compressor.name))
			return
		else:
			data = {dataset.name: {compressor.name: {}}}
			safe_json_file_update(results_file, data)
	else:
		with open(results_file, "r") as f:
			data = json.load(f)
	# Data properties cache file
	data_properties_file = DATA_PATH + "data_properties.json"
	if os.path.isfile(data_properties_file):
		with open(data_properties_file, "r") as f:
			data_properties = json.load(f)
	else:
		data_properties = {}
	if dataset.name not in data_properties.keys():
		data_properties[dataset.name] = {}
		safe_json_file_update(data_properties_file, data_properties)
	
	# Make sure data_range is in cache in the case of only_check_for_missing_results
	property_present = {key: key in data_properties[dataset.name] for key in ["data_range", "data_norm"]}
	if not (only_check_for_missing_results and property_present["data_range"] and property_present["data_norm"]):
		print("Loading original data, properties present: %s"%property_present)
		original = load_data(dataset.path, dataset.dtype)
		if not property_present["data_range"]:
			data_properties[dataset.name]["data_range"] = float(compute_data_range(original))
			safe_json_file_update(data_properties_file, data_properties)
		if not property_present["data_norm"]:
			data_properties[dataset.name]["data_norm"] = float(np.linalg.norm(original))
			safe_json_file_update(data_properties_file, data_properties)
	print_commands = True # For debugging
	
	# Initialize rest of set-up
	margin_factor = 2
	min_psnr = floor(relative_error_to_psnr(margin_factor*1e-1, data_properties[dataset.name]["data_range"], dataset.size(), data_properties[dataset.name]["data_norm"]))
	max_psnr = ceil(relative_error_to_psnr(1e-4/margin_factor, data_properties[dataset.name]["data_range"], dataset.size(), data_properties[dataset.name]["data_norm"]))
	psnr_step = 5
	samples_per_target_error = 5
	
	# Set-up for target error estimation
	target_error = None
	psnr = None
	use_continuous_target_error = compressor.get_valid_target_errors is None
	
	# Set up dataset on the right location for stable/efficient I/O, depending on the compressor's needs
	if only_check_for_missing_results:
		scratch_dataset = deepcopy(dataset)
	else:
		scratch_dataset = get_scratch_dataset(dataset, parallel_file_system=compressor.needs_parallel_io)
	
	for target_psnr_or_error in (range(min_psnr, max_psnr + psnr_step//2, psnr_step) if use_continuous_target_error else compressor.get_valid_target_errors(dataset)):
		
		# target_psnr_or_error represents target PSNR if use_continuous_target_error, else target error in this compressor's metric
		if not only_check_for_missing_results:
			print("Processing target %s %s"%("PSNR" if use_continuous_target_error else "error", target_psnr_or_error))
		
		# Iterate a few times to better approximate particular error
		reached_zero_error = False
		adjust_error = use_continuous_target_error
		for _ in range(3):
			
			# Estimate target error for this compressor
			target_error = compressor.estimate_target_error(target_psnr_or_error, data_range=data_properties[dataset.name]["data_range"], dataset=dataset, prev_target_error=target_error, prev_psnr=psnr) if use_continuous_target_error else target_psnr_or_error
			
			# Add entry for this target error if necessary
			if str(target_error) not in data[dataset.name][compressor.name].keys():
				if only_check_for_missing_results:
					print("Missing results for dataset %s and compressor %s: no samples for target %s %s"%(dataset.name, compressor.name, "PSNR" if use_continuous_target_error else "error", target_psnr_or_error))
					return
				data[dataset.name][compressor.name][str(target_error)] = []
				safe_json_file_update(results_file, data)
			
			# Keep generating samples until the stop condition is triggered
			while True:
				
				# The following block checks the stop condition for this loop
				
				# If there are no known samples, we need to generate at least one sample
				# If there is at least one sample:
				# - If the target error is discrete: we don't do any iterative target error approximation, so we keep sampling this point until we have samples_per_target_error samples
				# - If the target error is continuous: if we are within the desired margin of the target error, we keep sampling this point until we have samples_per_target_error samples. Else, we skip this point and adjust the target error.
				current_samples = len(data[dataset.name][compressor.name][str(target_error)])
				if current_samples > 0:
					
					psnr = data[dataset.name][compressor.name][str(target_error)][0]["psnr"]
					
					if psnr == inf:
						# We've reached the point of no error, so we can stop checking more precise error settings
						# In this case, we also are not within any desired margin of the target PSNR, so we won't take any more samples
						reached_zero_error = True
						adjust_error = False
						break
					
					if use_continuous_target_error:
						if abs(target_psnr_or_error - psnr) > 5:
							# Stop the sample loop but continue to the next iteration in the error adjustment loop
							break
						else:
							# Keep iterating over the sample loop but stop afterwards without adjusting the error
							adjust_error = False
				
				if current_samples >= samples_per_target_error:
					break
				
				if only_check_for_missing_results:
					print("Missing results for dataset %s and compressor %s: no samples for target %s %s"%(dataset.name, compressor.name, "PSNR" if use_continuous_target_error else "error", target_psnr_or_error))
					return
						
				if use_continuous_target_error:
					print("Attempting target error %s"%target_error)
				
				compression_results = compress(compressor, scratch_dataset, target_error, threads=threads, measure_time=True, measure_memory=True, print_commands=print_commands)
				decompression_results = decompress(compressor, scratch_dataset, threads=threads, print_commands=print_commands)
				
				# Compute and store results
				decompressed = load_data(get_decompressed_path(), dataset.dtype)
				psnr = compute_psnr(original, decompressed, data_properties[dataset.name]["data_range"])
				print("Compressor %s resulted in %s dB PSNR and compression factor %s"%(compressor.name, psnr, compression_results["compression_factor"]))
				data[dataset.name][compressor.name][str(target_error)].append({
					"psnr": psnr,
					"relative_error": psnr_to_relative_error(psnr, data_properties[dataset.name]["data_range"], original.size, data_properties[dataset.name]["data_norm"]),
					"compression": compression_results,
					"decompression": decompression_results
				})
				safe_json_file_update(results_file, data)
				
				# Remove decompressed data from memory before next compression to reduce peak memory usage
				del decompressed
				gc.collect()
				
				# Remove files so that they're not used by the following compressor
				cleanup()
			
			if not adjust_error:
				break
		
		if reached_zero_error:
			break

# Parse command-line options

if len(sys.argv) < 2:
	print("Usage:")
	print("- python3 measure-rate-distortion.py <dataset-name> [-c <compressor-names] [-t <threads>]")
	print("- python3 measure-rate-distortion.py --check-missing-results")
	exit(1)
arg_index = 1

if sys.argv[arg_index] == "--check-missing-results":
	# Iterate over all relevant dataset-compressor pairs
	# The following loop should iterate over the datasets first (outer loop) and the compressors second (inner loop), which is what we want
	TuckerMPI_1_process = generate_tuckermpi_compressor(True)
	extended_compressors = compressors + [generate_atc_compressor(rtmss=rtmss) for rtmss in (0.0, 0.2, 0.4, 0.6, 0.8)] + [TuckerMPI_1_process] + ATC_ALL_OPT_LEVELS
	dataset_names = ["foot", "brain", "moffett_field", "isotropic_fine_pressure_time", "isotropic_fine_velocity", "deforest_globe_8", "deforest_globe_33", "hurricane"]
	pairs = [(dataset_name, compressor) for dataset_name in dataset_names for compressor in extended_compressors]
	pairs.remove(("hurricane", TTHRESH)) # Too large for TTHRESH to process
	for dataset_name in dataset_names[-3:]:
		pairs.remove((dataset_name, TuckerMPI_1_process)) # Too large to load for TuckerMPI when only using one process
	for dataset_name, compressor in pairs:
		measure_rate_distortion(globals()[dataset_name], compressor, only_check_for_missing_results=True)
	exit(0)

# Parse dataset name
dataset_name = sys.argv[arg_index]
arg_index += 1
if dataset_name in globals():
	selected_dataset = globals()[dataset_name]
else:
	print("Did not recognize dataset name: %s"%dataset_name)
	exit(1)

# Parse optional compressors argument
selected_compressors = []
if arg_index < len(sys.argv) and sys.argv[arg_index] == "-c":
	# Only select requested compressors
	arg_index += 1
	compressor_names = sys.argv[arg_index].split(",")
	arg_index += 1
	for compressor_name in compressor_names:
		if compressor_name.startswith("ATC-"):
			# Parse ATC parameters
			parts = compressor_name.split("-")[1:]
			parts_index = 0
			rtmss = None
			dtmss = None
			blas = True
			epcm = False
			lpmt = None
			qrem = None
			hwm = None
			optimization_level = ATC_MAX_OPTIMIZATION_LEVEL
			run = 0
			while parts_index < len(parts):
				current_token = parts[parts_index]
				parts_index += 1
				if current_token == "RTMSS":
					rtmss = float(parts[parts_index])
					parts_index += 1
				elif current_token == "DTMSS":
					dtmss = float(parts[parts_index])
					parts_index += 1
				elif current_token == "OPT":
					optimization_level = int(parts[parts_index])
					parts_index += 1
				elif current_token == "NOBLAS":
					blas = False
				elif current_token == "EPCM":
					epcm = True
				elif current_token == "LPMT":
					lpmt = int(parts[parts_index])
					parts_index += 1
				elif current_token == "QREM":
					qrem = float(parts[parts_index])
					parts_index += 1
				elif current_token == "HWM":
					hwm = parts[parts_index]
					parts_index += 1
				elif current_token == "RUN":
					run = int(parts[parts_index])
					parts_index += 1
			selected_compressors.append(generate_atc_compressor(rtmss=rtmss, dtmss=dtmss, epcm=epcm, blas=blas, lpmt=lpmt, qrem=qrem, hwm=hwm, optimization_level=optimization_level, run=run))
		elif compressor_name == "TuckerMPI-1-process":
			selected_compressors.append(generate_tuckermpi_compressor(True))
		elif compressor_name == "TuckerMPI-NOSMF":
			selected_compressors.append(generate_tuckermpi_compressor(shared_memory_flags=False))
		else:
			# Non-ATC compressors
			if compressor_name in globals():
				selected_compressors.append(globals()[compressor_name])
			else:
				print("Did not recognize compressor name: %s"%compressor_name)
				exit(1)
else:
	# Select all compressors
	selected_compressors = compressors
	
# Parse optional threads argument
threads = 1
if arg_index < len(sys.argv) and sys.argv[arg_index] == "-t":
	arg_index += 1
	threads = int(sys.argv[arg_index])
	arg_index += 1

# Call measurement function with right arguments
for selected_compressor in selected_compressors:
	if not (selected_compressor.name == "TTHRESH" and selected_dataset.name == "Hurricane"):
		if selected_compressor.max_modes >= selected_dataset.modes:
			tmp_dataset = selected_dataset
		else:
			# Merge first two modes together
			tmp_dataset = Dataset(selected_dataset.name, selected_dataset.relative_path, [selected_dataset.sizes[0]*selected_dataset.sizes[1],] + selected_dataset.sizes[2:], selected_dataset.dtype, original_dtype=selected_dataset.original_dtype, most_constant_mode=selected_dataset.most_constant_mode, spatial_modes=selected_dataset.spatial_modes)
		measure_rate_distortion(tmp_dataset, selected_compressor, threads=threads)
