/*
 * This file implements the class atc::QuantizedArrayIterator.
 */

#include <numeric>

#include "QuantizedArrayIterator.hpp"
#include "util.hpp"

namespace atc {
	
	QuantizedArrayIterator::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) : QuantizedArrayIterator(is_factor ?
		QuantizedArrayIterator({size}, {0}, CoreFlattening::Unfold, false)
		: QuantizedArrayIterator(sizes, current_mode_order, flattening, track_multi_index)) { }
	
	QuantizedArrayIterator::QuantizedArrayIterator(const std::vector<size_t>& sizes, const std::vector<uint8_t>& current_mode_order, CoreFlattening flattening, bool track_multi_index)
	: flattening(flattening), track_multi_index(track_multi_index) {
		// Initialize helper variables
		modes = sizes.size();
		size = product(sizes);
		transposed_sizes = std::vector<size_t>(modes);
		for (uint8_t mode_index = 0; mode_index < modes; mode_index++)
			transposed_sizes[mode_index] = sizes[current_mode_order[mode_index]];
		// Initialize state
		original_index = 0;
		reordered_index = 0;
		multi_index = std::vector<size_t>(modes, 0);
		
		// Encoding-order-dependent stuff
		if (flattening == CoreFlattening::Zigzag || flattening == CoreFlattening::ZOrder ) {
			// Pre-compute strides
			strides = std::vector<size_t>(modes);
			strides[0] = 1;
			for (uint8_t mode = 0; mode < modes - 1; mode++)
				strides[mode + 1] = strides[mode]*transposed_sizes[mode];
			if (flattening == CoreFlattening::Zigzag)
				layer_index = 0;
		}
		
	}
	
	void QuantizedArrayIterator::advance() {
		
		reordered_index++;
		if (unfinished()) {
			
			switch (flattening) {
			case CoreFlattening::Unfold:
				if (track_multi_index) {
					original_index++;
					multi_index[0]++;
					uint8_t mode_index = 0;
					while (multi_index[mode_index] >= transposed_sizes[mode_index] && mode_index < modes - 1) {
						multi_index[mode_index] = 0;
						mode_index++;
						multi_index[mode_index]++;
					}
				} else {
					original_index++;
				}
				break;
			case CoreFlattening::Zigzag:
				
				for (uint8_t mode = 1; ; mode++) {
					if (multi_index[0] > 0 && multi_index[mode] < transposed_sizes[mode] - 1) {
						// Can advance at least on this zigzag line, so don't check for further overflow
						// Check for furthest index that can be decremented (necessary for not skipping certain elements at the edges)
						for (uint8_t decrement_mode = mode; decrement_mode-- > 0; ) {
							if (multi_index[decrement_mode] > 0) {
								multi_index[decrement_mode]--;
								multi_index[mode]++;
								original_index += strides[mode] - strides[decrement_mode];
								break;
							}
						}
						break;
					} else {
						// Can't advance further on this zigzag line
						if (mode < modes - 1) {
							// Still modes left to increment within the same zigzag layer, so reset these inner loops by concentrating as much energy as possible into the first indices
							size_t sum_left_to_allocate = std::accumulate(multi_index.begin(), multi_index.begin() + mode + 1, 0);
							size_t sum_left_to_remove = sum_left_to_allocate;
							for (uint8_t fill_mode = 0; fill_mode < mode + 1; fill_mode++) {
								size_t new_index = std::min(sum_left_to_allocate, transposed_sizes[fill_mode] - 1);
								// Because of wrap-around, overflow/underflow is not a problem in the following steps
								sum_left_to_remove -= multi_index[fill_mode];
								sum_left_to_allocate -= new_index;
								original_index += (new_index - multi_index[fill_mode])*strides[fill_mode];
								multi_index[fill_mode] = new_index;
								if (sum_left_to_remove == 0) {
									// All indices beyond here up till and including mode mode have value 0, so they can be skipped
									break;
								}
							}
						} else {
							// Last mode couldn't advance either, so we need to go to the next layer
							layer_index++;
							// Start at first valid position (within tensor bounds) in layer by allocating as much of the plane sum to the first indices
							size_t sum_left_to_allocate = layer_index;
							std::fill(multi_index.begin(), multi_index.end(), 0);
							original_index = 0;
							for (uint8_t fill_mode = 0; fill_mode < modes; fill_mode++) {
								multi_index[fill_mode] = std::min(sum_left_to_allocate, transposed_sizes[fill_mode] - 1);
								original_index += multi_index[fill_mode]*strides[fill_mode];
								sum_left_to_allocate -= multi_index[fill_mode];
								if (sum_left_to_allocate == 0)
									break;
							}
							break;
						}
					}
				}
				
				break;
				
			case CoreFlattening::ZOrder:
				
				for (uint8_t level = 0; level < 8*sizeof(size_t); level++) {
					const size_t mask = 1 << level;
					bool backtrack = true;
					for (uint8_t mode_index = 0; mode_index < modes; mode_index++) {
						// Increment index for this mode for this bit, so definitely flip it and then deal with potential carrying
						multi_index[mode_index] ^= mask;
						if (multi_index[mode_index] & mask) {
							// Bit just became 1, so no carry occurred
							if (multi_index[mode_index] >= transposed_sizes[mode_index]) {
								// Increment is finished because there is no carry now, but more increments will need to happen because we went out-of-bounds
								// No increments to lower modes will fix this out-of-bounds violation, so we can perform 2^modes increments right away in order to fix this violation
								// This corresponds to flipping the current bit for multi_index[mode_index] back to 0 and carrying a 1 to the next mode
								// Therefore, this can also not affect the fact that the lower modes all have indices in-bounds (the loop invariant)
								// If this makes multi_index[mode_index + 1] go out-of-bounds, it will be handled in the next iteration using the same check
								multi_index[mode_index] ^= mask;
							} else {
								// Normal situation, increase index in original core
								original_index += mask*strides[mode_index];
								backtrack = false;
								break;
							}
						} else {
							// Overflow happened, 1 is carried to the next mode's corresponding bit, increase index in original core
							original_index -= mask*strides[mode_index];
						}
					}
					
					if (!backtrack)
						break;
					
					// We will nog backtrack and perform the same procedure on the level of the parent node
					// The bits on this level were incremented past 111...1 so they are now 000...0
					// Therefore, all lower bits are already initialized correctly for when a valid parent node is found
					
				}
				
				break;
			}
			
		}
	}
	
	bool QuantizedArrayIterator::unfinished() {
		return reordered_index < size;
	}
	
	const size_t& QuantizedArrayIterator::get_original_index() { return original_index; }
	const size_t& QuantizedArrayIterator::get_reordered_index() { return reordered_index; }
	const size_t& QuantizedArrayIterator::get_multi_index(uint8_t mode_index) { return multi_index[mode_index]; }
	
}
