# This file shows how we can compress and decompress in-memory buffers from Python.

import numpy as np
from ctypes import *
from numpy.ctypeslib import ndpointer
from enum import IntEnum

# ATC PRELIMINARIES

atc = cdll.LoadLibrary("../lib/libatc.so")

ATC_VERSION = c_char_p.in_dll(atc, "ATC_VERSION").value.decode()
print("ATC v%s"%ATC_VERSION)

# Enums

# From https://v4.chriskrycho.com/2015/ctypes-structures-and-dll-exports.html
class CtypesEnum(IntEnum):
	"""A ctypes-compatible IntEnum superclass."""
	@classmethod
	def from_param(cls, obj):
		return int(obj)

class DataType(CtypesEnum):
	
	UINT8 = 0
	UINT16 = 1
	UINT32 = 2
	UINT64 = 3
	INT8 = 4
	INT16 = 5
	INT32 = 6
	INT64 = 7
	FLOAT32 = 8
	FLOAT64 = 9
	ERROR = 10
	
	@staticmethod
	def to_string(data_type):
		if data_type == DataType.ERROR:
			raise Exception("Cannot convert ERROR DataType, this value was normally obtained because an error occurred in ATC.")
		return {
			DataType.UINT8: "uint8",
			DataType.UINT16: "uint16",
			DataType.UINT32: "uint32",
			DataType.UINT64: "uint64",
			DataType.INT8: "int8",
			DataType.INT16: "int16",
			DataType.INT32: "int32",
			DataType.INT64: "int64",
			DataType.FLOAT32: "float32",
			DataType.FLOAT64: "float64"
		}[data_type]

class Reduction(CtypesEnum):
	DOWNSAMPLING = 0
	BOX = 1
	LANCZOS = 2

class StorageOrder(CtypesEnum):
	MSI = 0
	LSI = 1

class FloatType(CtypesEnum):
	FLOAT32 = 0
	FLOAT64 = 1

class ErrorMetric(CtypesEnum):
	RELATIVE_ERROR = 0
	RMSE = 1
	PSNR = 2

class IntType(CtypesEnum):
	UINT32 = 0
	UINT63 = 1

class CoreFlattening(CtypesEnum):
	UNFOLD = 0
	ZIGZAG = 1
	ZORDER = 2

class EntropyCodingMethod(CtypesEnum):
	ARITHMETIC = 0
	ANS = 1

# Structures

class CustomStruct(Structure):
	# Defines common functionality for all ctypes structs we will use
	def __repr__(self):
		if len(self._fields_) == 0:
			return "{}"
		else:
			string = "{\n"
			for field_name, _ in self._fields_:
				# Take representation of value of field, split into lines, indent each line apart from the first one, then join the lines into a single str
				value_repr = "\n".join([(line if i == 0 else "\t" + line) for i, line in enumerate(repr(getattr(self, field_name)).split("\n"))])
				string += "\t%s: %s,\n"%(field_name, value_repr)
			string += "}"
			return string

class Slice(CustomStruct):
	_fields_ = [
		("start", c_long),
		("stop", c_long),
		("step", c_long),
		("reduction", c_int), # Reduction
		("description", c_char_p),
	]

class Options(CustomStruct):
	
	_fields_ = [
		("verbose", c_bool),
		("suppress_warnings", c_bool),
		("timings", c_bool),
		("threads", c_int),
		("storage_order", c_int), # StorageOrder
		("float_type", c_int), # FloatType
		("skip_bytes", c_size_t),
		("error_metric", c_int), # ErrorMetric
		("apply_dct", c_bool),
		("dct_truncation_max_sse_share", c_double),
		("rank_truncation_max_sse_share", c_double),
		("mode_processing_order", POINTER(c_ubyte)),
		("precise_svd", c_bool),
		("print_squared_singular_values", c_bool),
		("householder", c_bool),
		("int_type", c_int), # IntType
		("core_flattening", c_int), # CoreFlattening
		("mode_storage_order", POINTER(c_ubyte)),
		("quantization_blocks", c_size_t),
		("last_pass_max_threads", c_size_t),
		("quantization_rel_error_margin", c_double),
		("split_bitplane_truncation", c_bool),
		("entropy_coding_method", c_int), # EntropyCodingMethod
		("cutout_length", c_size_t),
		("cutout", POINTER(Slice)),
		("autocrop", c_bool),
		("autocrop_threshold", c_double),
		("calculate_error", c_bool)
	]
	
	def __init__(self):
		# Copies contents of ATC_OPTIONS_DEFAULT
		# deepcopy and such is not sufficient, causes segfault
		default_options = Options.in_dll(atc, "ATC_OPTIONS_DEFAULT")
		for field_name, _ in self._fields_:
			setattr(self, field_name, getattr(default_options, field_name))

class Timing(CustomStruct):
	_fields_ = [
		("name", c_char_p),
		("timing", c_double)
	]

# Implementation is private in C, so we just provide an empty structure so we can construct opaque pointers
class ResultsImpl(CustomStruct):
	pass

# This struct is read-only, if the Python user modifies its values this may cause memory management problems and segmentation faults
class Results(CustomStruct):
	
	_fields_ = [
		("success", c_bool),
		("compressed_size", c_size_t),
		("compressed_data", POINTER(c_ubyte)), # Better to treat this as a pointer to an unsigned byte than as a char pointer, since the latter assumes it represents a null-terminated string
		("compression_factor", c_double),
		("decompressed_sizes_length", c_size_t),
		("decompressed_sizes", POINTER(c_size_t)),
		("error", c_double),
		("timings_length", c_size_t),
		("timings", POINTER(Timing)),
		("atc_results_impl", POINTER(ResultsImpl)) # Only for private use so this shouldn't be used by the Python user, but needs to be declared to ctypes to match the full C struct
	]
	
	def __del__(self):
		# Destructor, automatically called when this Python object goes out of scope
		# Destroys the associated C object
		destructor = atc.atc_destroy_results
		destructor.argtypes = [POINTER(Results),]
		destructor.restype = None
		destructor(pointer(self))

def print_timings(results):
	# Check for null pointer
	if results.timings:
		for i in range(results.timings_length):
			print("- %s: %s"%(results.timings[i].name.decode(), results.timings[i].timing))

# END OF ATC PRELIMINARIES

# Data generation
shape = [100, 100, 100]
size = np.prod(shape)
data = np.array(range(size), dtype=np.float32).reshape(shape)*((np.random.rand(*shape) - 0.5)*0.1 + 1) # Very low-rank data with a bit of noise

# Compression + decompression

def compress(data, rel_error):
	
	# Define options
	options = Options()
	options.threads = 4
	options.error_metric = ErrorMetric.RELATIVE_ERROR
	options.timings = True
	
	# Call compress
	compress_fun = getattr(atc, "atc_compress_mm_results_options_%s"%data.dtype.name)
	# We mark the original data and sizes pointers as ndpointers because they will be passed as NumPy arrays
	compress_fun.argtypes = [ndpointer(dtype=data.dtype, ndim=1, flags="C_CONTIGUOUS"), c_ubyte, ndpointer(c_size_t, ndim=1, flags="C_CONTIGUOUS"), c_double, POINTER(Options)]
	compress_fun.restype = Results
	results = compress_fun(data.flatten(), data.ndim, np.array(data.shape, dtype=np.uint64), rel_error, pointer(options))
	if not results.success:
		raise Exception("Compression was not successful!")
	
	return results

def decompress(compress_results):
	
	# Assumes that compress_results stores the compressed data
	
	# Determine decompressed tensor type
	read_data_type = atc.atc_read_data_type_m
	read_data_type.argtypes = [c_size_t, POINTER(c_ubyte)]
	read_data_type.restype = c_int
	data_type = read_data_type(compress_results.compressed_size, compress_results.compressed_data)
	dtype_string = DataType.to_string(data_type)
	
	# Determine decompressed tensor size
	read_decompressed_size = atc.atc_read_decompressed_size_m
	read_decompressed_size.argtypes = [c_size_t, POINTER(c_ubyte)]
	read_decompressed_size.restype = c_size_t
	decompressed_size = read_decompressed_size(compress_results.compressed_size, compress_results.compressed_data)
	
	# Allocate empty decompressed tensor
	decompressed = np.empty(decompressed_size, dtype=dtype_string)
	
	# Define options
	options = Options()
	options.timings = True
	
	# Decompress data into allocated tensor
	decompress_fun = getattr(atc, "atc_decompress_mm_results_options_%s"%dtype_string)
	# We mark the decompressed data pointer as an ndpointer because it will be a NumPy array
	decompress_fun.argtypes = [c_size_t, POINTER(c_ubyte), c_size_t, ndpointer(dtype=dtype_string, flags="C_CONTIGUOUS"), POINTER(Options)]
	decompress_fun.restype = Results
	decompress_results = decompress_fun(compress_results.compressed_size, compress_results.compressed_data, decompressed_size, decompressed, pointer(options))
	if not decompress_results.success:
		raise Exception("Decompression was not successful!")
	
	# Reshape data using decompressed results
	decompressed = decompressed.reshape(decompress_results.decompressed_sizes[:decompress_results.decompressed_sizes_length])
	
	return decompressed, decompress_results

print("Compression ...")
target_rel_error = 1e-3
compress_results = compress(data, target_rel_error)
print("Finished compression! Timings:")
print_timings(compress_results)

print("Decompression ...")
decompressed, decompress_results = decompress(compress_results)
print("Finished decompression! Timings:")
print_timings(decompress_results)

print("Original data size:", data.nbytes)
print("Compressed data size:", compress_results.compressed_size)
print("Compression factor:", compress_results.compression_factor)
print("Target relative error:", target_rel_error)
print("Relative error:", np.linalg.norm(decompressed - data)/np.linalg.norm(data))
