"""
/*********************************************************************
* Copyright (c) 2023 the Qrisp Authors
*
* This program and the accompanying materials are made
* available under the terms of the Eclipse Public License 2.0
* which is available at https://www.eclipse.org/legal/epl-2.0/
*
* SPDX-License-Identifier: EPL-2.0
**********************************************************************/
"""

# -*- coding: utf-8 -*-

import threading
import time


import numpy as np
from scipy.sparse import coo_array, bsr_matrix, coo_matrix, csr_matrix, csc_matrix, csr_array, csc_array

import qrisp.simulator.bi_array_helper as hlp


try:
    #sparse_dot_mkl seems to be only faster in situations, where the shape of the matrices,
    #results in better parallelization. For many quantum circuits, this was not the case
    #so this feature is disabled until further notice.
    raise Exception
    #Install command: conda install -c conda-forge sparse_dot_mkl
    from sparse_dot_mkl import dot_product_mkl
    def sparse_matrix_mult(a,b):
        
        if not isinstance(a, (np.ndarray, csr_array, csc_array)):
            a = a.tocsr()
        if not isinstance(b, (np.ndarray, csc_array, csr_array)):
            b = b.tocsc()
        
        return dot_product_mkl(a, b)
except:
    # print("Failed to import mkl sparse matrix multiplication. Install: with conda install -c conda-forge sparse_dot_mkl. Using scipy algorithm.")
    sparse_matrix_mult = lambda a,b : (a @ b)
    
#A quick helper function to evaluate the product of an iterable
def prod(iter):
    res = 1
    for i in range(len(iter)):
        res *= iter[i]
    return res


class BiArray:
    
    def v_dot(self, other):
        
        original_shape_self = self.shape
        orignal_shape_other = other.shape
        
        self.reshape((self.size,))
        other.reshape((other.size,))
        
        res = tensordot(self, other, (0,0))
        
        self.reshape(original_shape_self)
        other.reshape(orignal_shape_other)
        
        return res

#The classes inside this file aim describe so called BiArrays.
#A BiArray is an nd-array, where all the dimensions is are a power of 2, ie. 2, 4, 8, or 1024...
#This property can be used to efficiently manipulate indices of sparse matrices such that
#they can be processed quickly enough to use sparse matrix multiplication libraries for tensor contraction.
#This functionality is summarized in the SparseBiArray class.
#The DenseBiArray class is an interface to numpy arrays, that also use the restriction 
#of the dimension being a power of two. For instance does the DenseBiArray class support
#constant time reshaping and axis swapping. The reshaping and axis swappings are "collected"
#and executed together once the array is required for contraction.

#Both classes automatically make use of potential parallelism in the code.
#This is achieved by returning arrays where the content is "under construction" by a certain thread.
#If the "under construction" array is required in a contraction, the finishing of the calculation process is awaited.

#The tensordot function is a numpy-esque interface for performing tensor contractions. This function also 
#Automatically converts between sparse and dense BiArrays once a specified sparsity threshold is reached.

#A third class, the DummyBiArray, is a class that provides the same methods as the previous two,
#however the contractions are not actually executed, but instead the amount of floating point operations
#is logged in a class attribute. This allows to quickly evaluate the performance of a contraction order
#without actually contracting



multithreading_threshold = np.inf

#This SparseBiArray class provides a numpy-esque rank n tensor interface, which uses sparse matrix multiplication
#for tensor contraction. The fact that only powers of two are allowed for the size, allow efficient
#index manipulations, that permit faster construction and processing of said sparse matrices compared
#to for instance the rank n tensor interface of the "sparse" package.

#The idea is to store multi dimensional sparse arrays as the coo-array of their raveled version
#ie. we ravel the array and store the indices of the non-zero entries together with their data
#So for instance the array 
# array([[0, 2, 0, 0],
#        [3, 0, 0, 0]])
#is stored as with the non-zero indices nz_indices = [1, 4] and data = [2,3].
#If we now want to swap the axes of the array in this representation, we can do bit-level manipulations
#on the nz_indices array:
#In binary we have
#nz_indices = ["001", "100"]
#An axis swap can now be performed by manipulating the binary strings (which is possible rather efficiently)
#In this case the manipulation consists of moving the last two digits of the bitstring to the front
#After the axis swap we have
#nz_indices_axis_swapped = ["010", "001"] = [2, 1]
#If we now unravel the array with updated shape, we get
# array([[0, 3],
#        [2, 0],
#        [0, 0],
#        [0, 0]])
#Which confirms that this bit-level operation indeed represented an axis-swap

class SparseBiArray(BiArray):
    
    #This attribute decides after how many contractions, there should be a check for duplicates and zero entries
    #int the sparse matrix data
    contraction_counter_threshold = np.inf
    
    
    #Constructor of the SparseBiArray class.
    #init_object can be 
    #- a numpy array
    #- a tuple of arrays where init_object[0] is an array of non-zero indices and init_object[1] the corresponding data
    #- a sparse array in coo format
    def __init__(self, init_object, shape = None, contraction_counter = 0):
        #This attribute tracks how many sparse matrix contractions, this object has been through
        self.contraction_counter = contraction_counter
        
        
        #Handle the case of an array 
        if isinstance(init_object, np.ndarray):
            array = init_object
            
            
            #Store shape
            self.shape = array.shape
            
            #Ravel init array
            raveled_array = array.ravel()
            
            # raveled_array = np.around(raveled_array, 6, out = raveled_array)

            #Find non-zero indices
            self.nz_indices = np.nonzero(raveled_array)[0]
            
            #Store data
            self.data = raveled_array[self.nz_indices]
            
            # self.nz_indices, self.data = hlp.elim_zeros(self.nz_indices, self.data)
            
        #Handle the case that the init_object is a tuple of nz_indices/data arrays
        elif isinstance(init_object, tuple):
            
            if shape == None:
                raise Exception("Tried to initialise SparseBiArray from sparse data without providing a shape")
            
            self.nz_indices = init_object[0].ravel()
            self.data = init_object[1]
            
            self.nz_indices, self.data = hlp.elim_zeros(self.nz_indices, self.data)
        
        #Handle the case that the init_object is a sparse matrix/array
        elif isinstance(init_object, (coo_array, coo_matrix)):
            
            if shape == None:
                raise Exception("Tried to initialise SparseBiArray from sparse matrix without providing a shape")
            
            self.nz_indices = hlp.gen_flat_coords(init_object.col, init_object.row, init_object.shape)
            self.data = init_object.data
        
        elif isinstance(init_object, SparseBiArray):
            copy = init_object.copy()
            
            for var in copy.__dict__.keys():
            
                self.__dict__[var] = copy.__dict__[var]
                
            return
        else:
            raise Exception("Tried to initialize SparseBiArray with unknown type " + str(type(init_object)))
        if shape != None:
            self.shape = shape
        
        
        #Store the size
        self.size = prod(self.shape)
        
        if not ((self.size & (self.size-1) == 0) and self.size != 0):
            raise Exception("Tried to initialize SparseBiArray with size " + str(self.size) + " which is not a power of 2")
        
        #Store the sparsity
        self.sparsity = len(self.nz_indices)/self.size
        
        #This list keeps track of the axes permutations that have been applied
        #Once the SparseBiArray is needed for a contraction, all these swaps are applied within a single
        #function call and then the corresponding matrix is build from that
        self.index_bit_permutation = list(range(int(np.log2(self.size))))
        
        self.dtype = self.data.dtype
        
        #If the contraction counter is above the threshold, we sum the duplicates and eliminate the zeros
        if False:
        # if contraction_counter > self.contraction_counter_threshold:
            self.sum_duplicates()
            self.eliminate_zeros()
            self.contraction_counter = 0
    
    #Method for copying
    def copy(self):
        self.apply_swaps()
        return SparseBiArray((self.nz_indices.copy(), self.data.copy()), shape = tuple(self.shape))
    
    
    #This method checks if the sparse array has a thread object, which means that it's data
    #is under construction. If so, it waits until the thread is finished
    def catch_up(self):
        if hasattr(self, "thread"):
            start_time = time.time()
            self.thread.join()
            
            duration = time.time()-start_time
            if duration > 0.001:
                # print(duration)
                pass
            del self.thread
    
    
    #Reshaping method - note that no data is manipulated, ie. constant time
    def reshape(self, shape):
        
        if isinstance(shape, int):
            shape = (shape,)
        
        shape = tuple(shape)
        if prod(shape) != prod(self.shape):
            raise Exception("Tried to reshape with invalid shape")
        
        self.shape = shape
        
        return self
    
    
    #Method for swapping two axes
    def swapaxes(self, i, j):
        
        if i == j:
            return
        if i > j:
            i, j = j, i
        
        
        #In order to store the swap in the index_bit_permutation attribute,
        #we need to swap the corresponding digit ranges inside this list
        
        #For this we first determine the digit range of the index bitstrings
        #For instance imagine we have the index 215 and the shape (4,8,8)
        #and we want to swap the axes i = 0 and j = 2.
        #In binary, 215 is 11010111.
        #We first determine the layout of the indices inside this bitstring
        bit_shape = np.log2(self.shape).astype(np.int64)
        
        #In our case the bit shape is (2,3,3)
        
        i_interval_start = np.sum(bit_shape[:i])
        i_interval_end = i_interval_start + bit_shape[i]

        #In our example i = 0, the i_interval starts at the first (most significant) digit and ends after the third
        #ie |11|010111

        j_interval_start = np.sum(bit_shape[:j])
        j_interval_end = j_interval_start + bit_shape[j]
        
        #The j interval now starts at 6-th digit and ends at the last digit
        #ie. |11|010|111|
        
        #these two intervals have to be swapped.
        #However in order to makes this method constant in time, the actual swapping
        #of the indices is not performed here - rather we log the permutation of the indices
        
        #After creating of our SparseBiArray the attribute .index_bit_permutation simply holds
        #the numbers from 0 to log2(size)-1, ie. in the case of our example
        #[0,1,2,3,4,5,6,7]
        #In order to log the axis swap, we perform this permutation on the .index_bit_permutation attribute
        #ie. after the application of this method, we want it to be
        #[5,6,7,2,3,4,0,1]
        #ie. the interval of the first two and the last three swapped
        
        
        #In order to do this we split the index_bit_permutation list into 5 parts
        #1. the interval before the i-interval
        #2. the i-interval
        #3. the interval between i and j
        #4. the j interval
        #5. the interval after j
        
        #in the case of our example, we would have
        #1. = []
        #2. = [0,1]
        #3. = [2,3,4]
        #4. = [5,6,7]
        #5. = []
        
        #We than acquire the permuted list by concatenating these list in the correct order, ie.
        # 1. + 4. + 3. + 2. + 5.
        
        
        #Determine the first interval
        init_interval = self.index_bit_permutation[:i_interval_start]
        
        #Determine the i-interval
        i_interval = self.index_bit_permutation[i_interval_start:i_interval_end]
        
        #Determine the middle-interval
        middle_interval = self.index_bit_permutation[i_interval_end:j_interval_start]
        
        #Determine the j-interval
        j_interval = self.index_bit_permutation[j_interval_start:j_interval_end]
        
        #Determine the final interval
        final_interval = self.index_bit_permutation[j_interval_end:]
        
        #Concatenate intervals
        self.index_bit_permutation = init_interval + j_interval + middle_interval + i_interval + final_interval
        
        
        #Create the tuple for the new shape
        new_shape = list(self.shape)
                
        new_shape[i], new_shape[j] = new_shape[j], new_shape[i]
        
        self.shape = tuple(new_shape)
        
        return self
    
    #This method applies the index permutation generated too
    def apply_swaps(self):
        
        #Wait for any contraction processes to finish
        self.catch_up()
        
        #Check if any permutations have to be performed
        if self.index_bit_permutation == list(range(int(np.log2(self.size)))):
            return
        


        
        #Apply permutations
        #The bit level manipulations are done in the helper file
        self.nz_indices = hlp.permute_axes(self.nz_indices, np.array(self.index_bit_permutation)).ravel()
        
        #Reset the index_bit_permutation attribute, to 
        #indicate that no permutation has to be executed now
        self.index_bit_permutation = list(range(int(np.log2(self.size))))
    
    #This method sums duplicates ie. if an index appears twice in nz_indices,
    #the corresponding data entries are summed
    def sum_duplicates(self):
        self.nz_indices, self.data = hlp.sum_duplicates(self.nz_indices, self.data)
    
    #If after a contraction a new zero appeared, this entry does not need to tracked anymore
    #This method removes the zeros
    def elim_zeros(self):
        self.nz_indices, self.data = hlp.elim_zeros(self.nz_indices, self.data)
    
    
    #Builds up the sparse matrix in coo format
    #The coo format saves an array of row/column indices and the corresponding entries
    #Shape is the shape the sparse matrix should have
    def build_sr_matrix(self, shape, transpose = False):
        

        if prod(shape) != self.size:
            raise Exception("Tried to build sparse matrix with invalid shape")
        
        #Perform potential swaps
        self.apply_swaps()
        
        #Retrieve coordinates
        row, col = hlp.get_coordinates(self.nz_indices, shape)
        
        #Create sparse matrix
        if not transpose:
            res = coo_array((self.data, (row, col)), shape = shape)
        else:
            res = coo_array((self.data, (col, row)), shape = shape[::-1])
        
        
        #This attribute seems to be False which triggers a sorting algorithm,
        #which is in many cases among the most resource costly operations during simulation.
        #Setting this to True had no influence on the results tested so far.
        res.has_canonical_format = True
        
        return res
    
    #This method allows moving an index to the first position without
    #changing the order of the remaining indices
    def move_index_to_front(self, index):
        
        #if the index is already at the front, we don't need to do anything
        if index == 0:
            return
        
        #The logic for creating the ne self.index_bit_permutation attribute is similar as in swapaxes
        
        bit_shape = [int(np.log2(s)) for s in self.shape]
        
        new_shape = list(self.shape)
        new_shape.insert(0, new_shape.pop(index))
        
        self.shape = tuple(new_shape)
        
        index_range_start = sum(bit_shape[:index])
        index_range_end = index_range_start + bit_shape[index]
        
        init_perm = self.index_bit_permutation[:index_range_start]
        index_perm = self.index_bit_permutation[index_range_start:index_range_end]
        final_perm = self.index_bit_permutation[index_range_end:]
        self.index_bit_permutation = index_perm + init_perm + final_perm
        
    
    #This function performs tensor contraction along the specified axes
    #Axes can be contracted if they have the same dimension
    #Imagine we have two tensors with shapes (2,4,64,8,16) and (8,32,4,2)
    #then we could for instance contract the following axes
    #0 with 3
    #1 with 2
    #3 with 0
    #From this we will then get a tensor, which has the shape of the
    #first concatenated with the second but with the contraction indices removed
    #ie. (64,16,32)
    #The idea of this function is to move all the contraction indices to the front
    #reshape to matrices, transpose the first matrix, and then call some matrix 
    #multiplication algorithm. In the case of the example the shapes after 
    #moving to the front look like
    #(2,4,8,64,16) and (2,4,8,32)
    #After reshaping to matrices the shapes look like 
    #(2*4*8, 64*16) and (2*4*8, 32)
    #Transpose the first matrix
    #(64*16, 2*4*8) and (2*4*8, 32)
    #Call matrix multiplication to get a matrix of shape
    #(64*16,32)
    #Finally, reshape the result
    #(64,16,32)
    def contract(self, other, axes_self, axes_other):
        
        #Check if the contraction size of both tensors is allowed
        contraction_size = prod([self.shape[i] for i in axes_self])
        contraction_size_other= prod([other.shape[i] for i in axes_other])
        
        if contraction_size != contraction_size_other:
            raise Exception("Tried to contract tensor with differently sized contraction indices")
        
        #Copy axes list in order to prevent modification
        axes_self = list(axes_self)
        axes_other = list(axes_other)
        
        #Move axes to the front
        for i in range(len(axes_self)):
            self.move_index_to_front(axes_self[i])
            
            #Note that we need to increase the index of the axes
            #greater than the axes we just moved, because
            #the insertion at the front moved the all the indices
            #before axes_self[i]
            for j in range(i + 1, len(axes_self)):
                if axes_self[i] > axes_self[j]:
                    axes_self[j] += 1 
        
        #Perform the same for the other tensor
        for i in range(len(axes_other)):
            other.move_index_to_front(axes_other[i])
            
            for j in range(i + 1, len(axes_other)):
                if axes_other[i] > axes_other[j]:
                    axes_other[j] += 1
    
        
        #Determine new shape (now that the contraction axes are at the front, we can simply slice the shapes)
        res_shape = list(self.shape[len(axes_self):]) + list(other.shape[len(axes_other):])
        
        #In order to feature multiprocessing, we create an "empty" result array, which will be returned
        #after the calculation has been started. If the content is needed before the calculation finished,
        #the catch_up method will make sure the calculation finishes before further processing 
        res = SparseBiArray((np.zeros(1, dtype = np.int64), np.zeros(1, dtype = self.data.dtype)), shape = res_shape, contraction_counter = self.contraction_counter + other.contraction_counter)
        
        #Define multi processing wrapper
        def mp_wrapper():
            
            #Build up sparse matrices
            sr_matrix_self = self.build_sr_matrix((contraction_size, self.size//contraction_size), transpose = True)
            sr_matrix_other = other.build_sr_matrix((contraction_size, other.size//contraction_size))
            
            #Perform sparse matrix multiplication
            
            res_sr_matrix = sparse_matrix_mult(sr_matrix_self, sr_matrix_other).tocoo()
            
            
            #Acquire flattened coordinates from the helper function   
            res.nz_indices = hlp.gen_flat_coords(res_sr_matrix.col, res_sr_matrix.row, res_sr_matrix.shape)
            
            #Set the data
            res.data = res_sr_matrix.data

            # res.nz_indices, res.data = hlp.elim_zeros(res.nz_indices, res.data)

            
            #Set the sparsity
            res.sparsity = len(res.nz_indices)/res.size
            
            #Perform consolidation operations
            if False:
            # if res.contraction_counter > self.contraction_counter_threshold:
                res.sum_duplicates()
                res.eliminate_zeros()
                res.contraction_counter = 0
        
        #Start the thead
        if res.size > multithreading_threshold:
            res.thread = threading.Thread(target = mp_wrapper)
            res.thread.start()
        else:
            mp_wrapper()
        
        return res 

    #This method returns self as a numpy array        
    def to_array(self):
        #Apply all swapping operations
        self.apply_swaps()
        
        #Construct the flat array
        res = hlp.construct_flat_array(self.nz_indices, self.data, self.size)
        
        #Reshape
        return res.reshape(self.shape)

    
    
    def ravel(self):
        return self.reshape(self.size)
    
    def __repr__(self):
        return str(self.to_array())
    
    
    #This function allows to split the array into it's lower and upper half
    #ie. if this SparseBiArray represents the array [0,1,2,3,4,5,6,7]
    #we return [0,1,2,3] and [4,5,6,7] as SparseBiArrays
    #This feature is important when measuring QuantumState as the norms of the
    #lower and upper half of the array determine the probability of measuring a 1 or 0
    def split(self):
        
        #Check if self is in a suitable shape (multi dimensional arrays can not be splitted as intended)
        if len(self.shape) != 1:
            raise Exception("Tried to split multi dimensional sparse array")
            
        #Applya swaps
        self.apply_swaps()    
        
        
        #Find the boolean mask of indizes that are less than half of the size
        mask = self.nz_indices < self.size//2
        
        #Acquire the lower half nz_indices
        lower_half_nz_indices = self.nz_indices[mask].copy()
        #Acquire the lower half data
        lower_half_data = self.data[mask].copy()
        
        #Create lower half result
        lower_half = SparseBiArray((lower_half_nz_indices, lower_half_data), shape = (self.shape[0]//2,))
        
        #Perform the same process for the upper half but with inverted mask
        mask = np.invert(mask)
        upper_half_nz_indices = self.nz_indices[mask] - self.size//2
        upper_half_data = self.data[mask].copy()
        
        upper_half = SparseBiArray((upper_half_nz_indices, upper_half_data), shape = (self.shape[0]//2,))
        
        #Return result
        return lower_half, upper_half
    
    def multi_split(self, indices):
        
        sprs = self.build_sr_matrix([2**(len(indices)), self.size//2**len(indices)]).tocsr()
        
        col_indices = sprs.indices
        row_ptr = sprs.indptr
        data = sprs.data
        
        new_bi_arrays = []
        for r in range(len(row_ptr)-1):
            if row_ptr[r] == row_ptr[r+1]:
                new_bi_array = None
            else:
                new_bi_array = SparseBiArray((col_indices[row_ptr[r]:row_ptr[r+1]].astype(np.int64), data[row_ptr[r]:row_ptr[r+1]]), shape = (self.size//2**len(indices),))
            
            new_bi_arrays.append(new_bi_array)
            
            
        
        return new_bi_arrays
    
    #Calculate the squared norm, ie. the sesquilinear scalar product of self with itself
    def squared_norm(self):
        return np.abs(np.vdot(self.data, self.data))
    
    
    #Return self as a DenseBiArray
    def to_dense(self):
        
        res = DenseBiArray(np.zeros(1), sparsity = self.sparsity, shape = self.shape)
        #We perform the same technique to enable multithreading as in contract,
        #ie. return an unfinished result that is waited on if required
        def mp_wrapper():
            res.data = self.to_array()
        
        
        #Start the thead
        if self.size > multithreading_threshold:
            res.thread = threading.Thread(target = mp_wrapper)
            res.thread.start()
        else:
            mp_wrapper()
        
        return res
    
    def to_sparse(self):
        return self


#This class serves mainly as an interface to numpy arrays using the same methods as the SparseBiArray class
#this way the algorithm using these classes doesn't need to care about wether it's treating
#a sparse array or a dense array.
#Furthermore some of the multithreading techniques from the SparseBiArray class are deployed
class DenseBiArray(BiArray):
    
    #The constructor can only be called with numpy arrays
    def __init__(self, array, sparsity = None, shape = None):
        
        if not isinstance(array, np.ndarray):
            raise Exception("Tried to initialize DenseBiArray with type " + str(type(array)))
        
        #Save data
        self.data = array
        
        #Set shape
        if shape:
            self.shape = shape
        else:
            self.shape = array.shape
            
        #Determine size
        self.size = prod(self.shape)
        
        #If not given, estimate sparsity. In many situations the sparsity
        #of a contraction result can be estimated (and then given to the constructor), 
        #so it is helpfull to track it
        if sparsity == None:
            self.sparsity = np.count_nonzero(array)/array.size
        else:
            # self.sparsity = np.count_nonzero(array)/array.size
            self.sparsity = sparsity
        
        #This attribute serves a similar purpose as in the SparseBiArray class
        self.index_bit_permutation = list(range(int(np.log2(self.size))))
        
        self.dtype = self.data.dtype
    
    #Copy method
    def copy(self):
        self.apply_swaps()
        return DenseBiArray(self.data.copy(), sparsity = self.sparsity)
    
    #Catch up method (same purpose as in SparseBiArray)
    def catch_up(self):
        if hasattr(self, "thread"):
            start_time = time.time()
            self.thread.join()
            
            duration = time.time()-start_time
            if duration > 0.001:
                # print("Waited ", duration)
                pass
            del self.thread
        
    #The shape of the DenseBiArray is tracked only as a tuple.
    #Once the data is required for a contraction, the actual data will be reshaped
    def reshape(self, shape):
        
        if isinstance(shape, int):
            shape = (shape,)
        
        if prod(shape) != prod(self.shape):
            raise Exception("Tried to reshape with invalid shape")
        
        self.shape = tuple(shape)
        return self
    
    #This method works similarly as it's equivalent in SparseBiArray
    def swapaxes(self, i, j):
        
        
        if i == j:
            return
        if i > j:
            i, j = j, i
        
        
        bit_shape = np.log2(self.shape).astype(np.int64)
        
        i_interval_start = sum(bit_shape[:i])
        i_interval_end = i_interval_start + bit_shape[i]

        j_interval_start = sum(bit_shape[:j])
        j_interval_end = j_interval_start + bit_shape[j]
        
        init_interval = self.index_bit_permutation[:i_interval_start]
        i_interval = self.index_bit_permutation[i_interval_start:i_interval_end]
        
        middle_interval = self.index_bit_permutation[i_interval_end:j_interval_start]
        
        j_interval = self.index_bit_permutation[j_interval_start:j_interval_end]
        
        final_interval = self.index_bit_permutation[j_interval_end:]
        
        self.index_bit_permutation = init_interval + j_interval + middle_interval + i_interval + final_interval
        
        new_shape = list(self.shape)
                
        new_shape[i], new_shape[j] = new_shape[j], new_shape[i]
        
        self.shape = tuple(new_shape)
        
        return self
    
    #This method works similarly as it's equivalent in SparseBiArray
    def move_index_to_front(self, index):
        
        
        if index == 0:
            return
        
        bit_shape = [int(np.log2(s)) for s in self.shape]
        
        new_shape = list(self.shape)
        new_shape.insert(0, new_shape.pop(index))
        
        self.shape = tuple(new_shape)
        
        index_range_start = sum(bit_shape[:index])
        index_range_end = index_range_start + bit_shape[index]
        
        init_perm = self.index_bit_permutation[:index_range_start]
        index_perm = self.index_bit_permutation[index_range_start:index_range_end]
        final_perm = self.index_bit_permutation[index_range_end:]
        self.index_bit_permutation = index_perm + init_perm + final_perm
        
    #This method works similarly as it's equivalent in SparseBiArray
    def apply_swaps(self):
        
        #Wait for any computation processes to finish
        self.catch_up()
        
        #Check if axes need to be swapped
        if self.index_bit_permutation == list(range(int(np.log2(self.size)))):
            
            #Check if data needs to be reshaped
            if self.data.shape != self.shape:
                self.data = self.data.reshape(self.shape)
        
            return
        
        #The plan is now to reshape into "tensor format" ie. log2(self.size) axes of dimension
        #two and permute the axes
        
        #Numpy however does not allow arrays with more than 32 axes, so it would not be possible
        #to manipulate arrays describing more than 32 qubits.
        #In order to reduce the number of axes that need to be permuted,
        #we identify intervals of indices that do not need permutation on them
        #and group them into one axis
        
        
        
        
        #Get the permutation of the axes
        bi_index_perm = [self.index_bit_permutation.index(i) for i in range(int(np.log2(self.size)))]
        
        
        
        if False:
            
            
            # result = np.empty(self.data.size, dtype = self.data.dtype)

            index_perm = hlp.invert_permutation(np.array(bi_index_perm))[::-1]
            index_perm = np.abs(len(index_perm)-1-index_perm)
            
            
            #flattened index array
            f_index_array = np.empty(2**len(index_perm), dtype = np.int64)
            
            f_index_array[0] = 0
            # result[0] = self.data.ravel()[0]
            
            
            #Unfortunately, this is not faster (maybe if implemented in C?)
            self.data = hlp.bi_array_moveaxis(self.data.ravel(), index_perm, f_index_array)
            # temp = hlp.bi_array_moveaxis(self.data.ravel(), index_perm, f_index_array)
            
            # permutated_indices = hlp.bi_array_moveaxis_jitted(index_perm, f_index_array)
            # self.data = self.data.ravel()[temp.ravel()]
            
            self.data.shape = self.shape
        
        else:
            perm = list(bi_index_perm)
            #Initialize the shape
            perm_shape = len(self.index_bit_permutation)*[2]
            
            
            i = 1
            while i < len(perm):
                #If this condition is met, both indices form an interval,
                #that does not need further permutation
                if perm[i]-1 == perm[i-1]:
                    
                    #Remove the index from the perm and perm_shape list
                    perm.pop(i)
                    perm_shape.pop(i)
                    
                    #Update the shape
                    perm_shape[i-1] *= 2
                    
                    #Update the permutation indices
                    for j in range(len(perm)):
                        if perm[j] > perm[i-1]:
                            perm[j] -= 1
                else:
                    i += 1
            
                
            #Reshape data to the required shape
            self.data = self.data.reshape(perm_shape)
            
            #Move the axes on the data
            self.data = np.moveaxis(self.data, list(range(len(perm))), perm)
        
            #Reshape the data
            self.data = self.data.reshape(self.shape)
        
        #Set the index_bit_permutation to identity
        self.index_bit_permutation = list(range(int(np.log2(self.size))))
    
    #Raveling method
    def ravel(self):
        self.reshape(self.size)
        self.apply_swaps()
        return self
        
    #This function works similar as it's equivalent in SparseBiArray
    #ie. identify contraction axes, move to the front, transpose one matrix
    #and then apply matrix multiplication
    def contract(self, other, axes_self, axes_other):
        
        
        #Check if the contraction size of both tensors is allowed
        contraction_size = prod([self.shape[i] for i in axes_self])
        contraction_size_other= prod([other.shape[i] for i in axes_other])
        
        if contraction_size != contraction_size_other:
            raise Exception("Tried to contract tensor with differently sized contraction indices")
        
        #Copy axes list in order to prevent modification
        axes_self = list(axes_self)
        axes_other = list(axes_other)
        
        #Move axes to the front
        for i in range(len(axes_self)):
            self.move_index_to_front(axes_self[i])
            
            #Note that we need to increase the index of the axes
            #greater than the axes we just moved, because
            #the insertion at the front moved the all the indices
            #before axes_self[i]
            for j in range(i + 1, len(axes_self)):
                if axes_self[i] > axes_self[j]:
                    axes_self[j] += 1 
        
        #Perform the same for the other tensor
        for i in range(len(axes_other)):
            other.move_index_to_front(axes_other[i])
            
            for j in range(i + 1, len(axes_other)):
                if axes_other[i] > axes_other[j]:
                    axes_other[j] += 1
    
        
        #Determine new shape (now that the contraction axes are at the front, we can simply slice the shapes)
        res_shape = list(self.shape[len(axes_self):]) + list(other.shape[len(axes_other):])
        
        
        #The sparsity can be estimated using statistical methods:
        #In an regular tensor contraction, any entry of the resulting tensor,
        #is the sum of n = contraction_size products of entries of the two source tensors
        #Each of these products is non-zero with probability p_nz = self.sparsity*other.sparsity
        
        # p_nz = self.sparsity*other.sparsity
        
        #The probabilty for the product being zero is therefore p_z = 1 - self.sparsity*other.sparsity
        
        # p_z = 1 - self.sparsity*other.sparsity
        
        #therefore we perform n = contraction_size random experiments with probability p_z
        #The result entry is zero if all products are zero, so
        
        # p_res_z = p_z**contraction_size
        
        #To get the sparsity, we finally need the probability of the result entry being non-zero
        
        # res_sparsity = 1-p_res_z
        
        #We tested this formula and it works extremely well.
        
        res_sparsity = 1-(1-self.sparsity*other.sparsity)**contraction_size
        
        #Prepare "unfinished" result
        res = DenseBiArray(np.array([self.data.ravel()[0]*other.data.ravel()[0]]), sparsity = res_sparsity, shape = res_shape)
        
        #Prepare multithreading wrapper
        #Performs the same logic as it's SparseBiArray equivalent
        def mp_wrapper():
            
            
            original_shape_self = self.shape
            original_shape_other = other.shape
            
            self.reshape([contraction_size, self.size//contraction_size])
            other.reshape([contraction_size, other.size//contraction_size])
            
            self.swapaxes(0,1)
            
            self.apply_swaps()
            other.apply_swaps()
            
            res.data = np.matmul(self.data, other.data)
            
            self.swapaxes(0,1)
            self.reshape(original_shape_self)
            other.reshape(original_shape_other)
            
            
        if res.size > multithreading_threshold:
            #Start the wrapper    
            res.thread = threading.Thread(target = mp_wrapper)
            res.thread.start()
        else:
            mp_wrapper()
        #Return the result
        return res
    
    #This method works similarly as it's equivalent in SparseBiArray
    def split(self):
        
        if len(self.shape) != 1:
            raise Exception("Tried to split multi rank BiArray")
        self.apply_swaps()    
        
        return DenseBiArray(self.data[:self.size//2].copy(), self.sparsity), DenseBiArray(self.data[self.size//2:].copy(), sparsity = self.sparsity)
    
        
    def multi_split(self, indices):
        
        original_shape = tuple(self.shape)
        self.reshape([2**(len(indices)), self.size//2**len(indices)])
        
        np_array = self.to_array()
        
        new_bi_arrays = []
        
        for r in range(2**len(indices)):
            new_bi_arrays.append(DenseBiArray(np_array[r,:]))
        
        self.reshape(original_shape)
        
        return new_bi_arrays

    #This method works similarly as it's equivalent in SparseBiArray
    def squared_norm(self):
        self.apply_swaps()
        return np.abs(np.vdot(self.data, self.data))
    
    #This method works similarly as it's equivalent in SparseBiArray
    def to_array(self):
        self.apply_swaps()
        return self.data
    
    
    #This method works similarly as it's equivalent in SparseBiArray    
    def to_dense(self):
        return self
    
    #Generates the SparseBiArray version of self using multithreading
    def to_sparse(self):
        
        res = SparseBiArray((np.zeros(1, dtype = np.int64), np.zeros(1)), shape = self.shape)

        #Prepare the wrapper       
        def mp_wrapper():
            self.apply_swaps()
            
            #Ravel init array
            raveled_array = self.data.ravel()
            
            
            #Find non-zero indices
            res.nz_indices = np.nonzero(raveled_array)[0]

            #Store data
            res.data = raveled_array[res.nz_indices]
            
            res.nz_indices, res.data = hlp.elim_zeros(res.nz_indices, res.data)
            
            res.sparsity = len(res.nz_indices)/res.size
        
        #Start the
        if res.size > multithreading_threshold:
            res.thread = threading.Thread(target = mp_wrapper)
            res.thread.start()
        else:
            mp_wrapper()        
        return res
    
    def __repr__(self):
        return str(self.to_array())

#This class is used to estimate the complexity of a contraction order
#The contract method does not actually execute a contraction but only logs
#the estimated amount of floating point operations
class DummyBiArray(BiArray):
    
    current_memory = np.zeros(1)
    max_memory = np.zeros(1)
    flops = np.zeros(1)
    
    def __init__(self, data = None, shape = None, sparsity = 1):
        
        self.sparsity = sparsity
        if not isinstance(data, type(None)):
            self.shape = data.shape
            self.sparsity = np.count_nonzero(data)/data.size
        if not isinstance(shape, type(None)):
            self.shape = tuple(shape)
        
        
        self.size = prod(self.shape)
        
    def reshape(self, shape):
        self.shape = shape
        
    def swapaxes(self, axis_0, axis_1):
        
        self.shape = list(self.shape)
        self.shape[axis_0], self.shape[axis_1] = self.shape[axis_1], self.shape[axis_0]
        self.shape = tuple(self.shape)
    
    def contract(self, other, axes_self, axes_other):
        
        contraction_size = prod([self.shape[i] for i in axes_self])
        
        temp_shape_self = list(self.shape)
        temp_shape_other = list(other.shape)
        
        delete_multiple_element(temp_shape_self, axes_self)
        delete_multiple_element(temp_shape_other, axes_other)
        
        res_sparsity =  1-(1-self.sparsity*other.sparsity)**contraction_size
        res_shape = temp_shape_self + temp_shape_other
        try:
            self.flops += prod(res_shape)*contraction_size
        except TypeError:
            self.flops += np.inf
            
        res = DummyBiArray(shape = res_shape, sparsity = res_sparsity)
        self.current_memory[0] += res.size - self.size - other.size
        
        if self.current_memory > self.max_memory:
            self.max_memory[0] = self.current_memory[0]
        
        return res
        
    def copy(self):
        return DummyBiArray(shape = self.shape)
    @classmethod
    def reset_stat_counter(self):
        
            
        self.current_memory = np.zeros(1)
        self.max_memory = np.zeros(1)
        self.flops = np.zeros(1)


#This function emulates the behavior of np.tensordot
def tensordot(a,b, axes, contract_sparsity_threshold = 0.01):
    #Turn axes object into a list
    axes = list(axes)
    if isinstance(axes[0], int):
        axes[0] = [axes[0]]
    if isinstance(axes[1], int):
        axes[1] = [axes[1]]
    
    #Treat the case of a dummy contraction    
    if isinstance(a, DummyBiArray) and isinstance(b, DummyBiArray):
        return a.contract(b, axes[0], axes[1])
    
    #Treat the case of a numpy contraction
    if isinstance(a, np.ndarray) and isinstance(b, np.ndarray):
        return np.tensordot(a,b, axes)
    
    if a.data.dtype == np.dtype("O") or b.data.dtype == np.dtype("O"):
        a = a.to_dense()
        b = b.to_dense()
        
        return a.contract(b, axes[0], axes[1])
    
    if len(axes[0]) != 0:
        
        contraction_size = prod([a.shape[i] for i in axes[0]])
        
        #Estimate resulting sparsity (for more information on this estimate check the contract method of DenseBiArray)
        res_sparsity = 1-(1-a.sparsity*b.sparsity)**contraction_size
        
        res_size = a.size*b.size/contraction_size
    else:
        res_sparsity = a.sparsity*b.sparsity
        
        res_size = a.size*b.size
    
    
    
    if (res_sparsity > contract_sparsity_threshold or a.size*b.size < 2**12) and not res_size > 2**32:
        
    #     if ((isinstance(a, SparseBiArray) and isinstance(b, DenseBiArray)) or (isinstance(b, SparseBiArray) and isinstance(a, DenseBiArray))):
            
    #         return mixed_contractor(a, b, axes)
    
            a = a.to_dense()
            b = b.to_dense()
        
    else:
            
            a = a.to_sparse()
            b = b.to_sparse()
        
        
    #Treat the case of a "contractionless" tensor product    
    if len(axes[0]) == 0:
        
        #We utilize the multithreading pattern presented in the previous contraction methods
        if isinstance(a, SparseBiArray) and isinstance(b, SparseBiArray):
            #Prepare result arrays            
            res = SparseBiArray((np.zeros(1, dtype = np.int64), np.ones(1)), shape = list(a.shape) + list(b.shape))
            
            #Prepare wrapper
            def mp_wrapper():
                
                #Apply any potential swaps on the source arrays
                a.apply_swaps()
                b.apply_swaps()
                
                #Imagine we have two arrays with flattened nz_indices [1,2,3], [2,3,6] and sizes 16, 8  respectively
                #the nz_indices of the result can be ordered in a matrix
                #[1,2,3] + 2*16
                #[1,2,3] + 3*16
                #[1,2,3] + 6*16
                #(where the plus is executed on all entries)
                
                #This idea is encapsulated in the following command
                nz_indices = np.tensordot(b.size*a.nz_indices, np.ones(len(b.data), dtype = np.int64), ((),())) + np.tensordot(np.ones(len(a.data), dtype = np.int64), b.nz_indices, ((),()))
                
                #The corresponding data can be calculated as
                data = np.tensordot(a.data.ravel(), b.data.ravel(), ((),())).ravel()
                
                #Set nz_indices and data
                res.nz_indices = nz_indices.ravel()
                
                res.data = data
            
            if res.size > multithreading_threshold:
                res.thread = threading.Thread(target = mp_wrapper)
                res.thread.start()
            else:
                mp_wrapper()
                
            return res
            
        else:
            #For the DenseBiArray, we can simply use the numpy tensordot
            a.catch_up()
            b.catch_up()
            
            res = DenseBiArray(np.array([a.data.ravel()[0]*b.data.ravel()[0]]), sparsity = a.sparsity*b.sparsity, shape = list(a.shape) + list(b.shape))
            
            def mp_wrapper():
                a.apply_swaps()
                b.apply_swaps()
                
                res.data = np.tensordot(a.to_array(), b.to_array(), axes = ((),()))
            
            if res.size > multithreading_threshold:
                res.thread = threading.Thread(target = mp_wrapper)
                res.thread.start()
            else:
                mp_wrapper()
            
            return res
    
    return a.contract(b, axes[0], axes[1])


def mixed_contractor(a, b, axes):
    
    if not ((isinstance(a, SparseBiArray) and isinstance(b, DenseBiArray)) or (isinstance(b, SparseBiArray) and isinstance(a, DenseBiArray))):
        raise Exception("Tried to call mixed contraction on non-mixed input")
        
    #Copy axes list in order to prevent modification
    axes_a = list(axes[0])
    axes_b = list(axes[1])
        
        
    #Check if the contraction size of both tensors is allowed
    contraction_size = prod([a.shape[i] for i in axes_a])
    contraction_size_b= prod([b.shape[i] for i in axes_b])
    
    if contraction_size != contraction_size_b:
        raise Exception("Tried to contract tensor with differently sized contraction indices")
    
    
    #Move axes to the front
    for i in range(len(axes_a)):
        a.move_index_to_front(axes_a[i])
        
        #Note that we need to increase the index of the axes
        #greater than the axes we just moved, because
        #the insertion at the front moved the all the indices
        #before axes_a[i]
        for j in range(i + 1, len(axes_a)):
            if axes_a[i] > axes_a[j]:
                axes_a[j] += 1 
    
    #Perform the same for the b tensor
    for i in range(len(axes_b)):
        b.move_index_to_front(axes_b[i])
        
        for j in range(i + 1, len(axes_b)):
            if axes_b[i] > axes_b[j]:
                axes_b[j] += 1

    
    #Determine new shape (now that the contraction axes are at the front, we can simply slice the shapes)
    res_shape = list(a.shape[len(axes_a):]) + list(b.shape[len(axes_b):])
    
    
    res_sparsity = 1-(1-a.sparsity*b.sparsity)**contraction_size
    
    #Prepare "unfinished" result
    res = DenseBiArray(np.empty(1, dtype = a.data.dtype), sparsity = res_sparsity, shape = res_shape)
    
    
    #Define multi processing wrapper
    def mp_wrapper():
        
        #Build up sparse matrices
        if isinstance(a, SparseBiArray):
            mult_matrix_a = a.build_sr_matrix((contraction_size, a.size//contraction_size), transpose = True)
            original_shape_b = b.shape
            b.reshape([contraction_size, b.size//contraction_size])
            b.apply_swaps()
            mult_matrix_b = b.data
            
        else:
            
            original_shape_a = a.shape
            a.reshape([contraction_size, a.size//contraction_size])
            a.swapaxes(0,1)
            a.apply_swaps()
            mult_matrix_a = a.data
            mult_matrix_b = b.build_sr_matrix((contraction_size, b.size//contraction_size))
        
        
        #Set the data
        res.data = sparse_matrix_mult(mult_matrix_a, mult_matrix_b)
    
    #Start the thead
    if res.size > multithreading_threshold:
        res.thread = threading.Thread(target = mp_wrapper)
        res.thread.start()
    else:
        mp_wrapper()
    
    return res

def delete_multiple_element(list_object, indices):
    indices = sorted(indices, reverse=True)
    for idx in indices:
        if idx < len(list_object):
            list_object.pop(idx)
