/*
 * This file defines the class atc::QuantizedArrayIterator and defines its members.
 */

#pragma once

#include "Options.hpp"

namespace atc {
	
	/*
	 * This class provides an opaque interface to efficient iteration algorithms for several core vectorization methods, such as zigzag ordering or the Z-order.
	 * Lexicographic ordering is also supported but is simply a wrapper around trivial behavior.
	 * 
	 * Note that we are effectively emulating polymorphism for some methods in terms of the flattening parameter.
	 * However, the time cost of finding the right switch-case branch should be negligible compared to the cost of the branch content itself.
	 */
	class QuantizedArrayIterator {
	public:
		// The most general constructor, for both core and factors.
		QuantizedArrayIterator(bool is_factor, const std::vector<size_t>& sizes, size_t size, const std::vector<uint8_t>& current_mode_order, CoreFlattening flattening, bool track_multi_index = false);
		/*
		 * Creates a potentially reordered iterator, meant for the core.
		 * If track_multi_index is false, the multi-index may or may not be stored internally.
		 */
		QuantizedArrayIterator(const std::vector<size_t>& sizes, const std::vector<uint8_t>& current_mode_order, CoreFlattening flattening, bool track_multi_index = false);
		// Advances the iterator by one position. This increments the reordered index and updates the original index accordingly.
		void advance();
		// Returns true if the iterator is not at the end of the array, i.e. original_index < array_size.
		bool unfinished();
		// Gets the current index in the original data (the core, stored in lexicographic order)
		const size_t& get_original_index();
		// Gets the current index in the reordered/quantized array.
		const size_t& get_reordered_index();
		/*
		 * Returns the current index for the mode at mode_index in the current mode storage order.
		 * Should only be used if track_multi_index is true.
		 */
		const size_t& get_multi_index(uint8_t mode_index);
	private:
		uint8_t modes; // Short-hand for transposed_sizes.size()
		size_t size; // Short-hand for product of sizes
		std::vector<size_t> transposed_sizes; // Mode sizes of the tensor, transposed to the current mode order
		const CoreFlattening flattening; // The flattening/vectorization method that is used by this iterator
		const bool track_multi_index; // Stores whether or not the multi-index should be tracked as well
		// State
		size_t original_index; // See get_original_index
		size_t reordered_index; // See get_reordered_index
		std::vector<size_t> multi_index; // See get_reordered_index
		size_t layer_index; // Used only for zigzag encoding
		std::vector<size_t> strides; // Used only for zigzag encoding, derived from transposed_sizes
	};
	
}
