"""
/*********************************************************************
* 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
**********************************************************************/
"""

import numpy as np

#Checks if a gate is permeable on a set of qubits
#Permeable means that the gate has the form
# U = \sum_{x = 0}^{2^n} |x><x| U_x
#where the braket expression represents the qubits
#of the list "qubit_indices"
#{U_x | x <= 2^n} is a set of operations on the remaining qubits

#Applying a gate which is permeable on a register R to a computational basis state in R
#leaves the state of R invariant ie. if R is in state |x>
# U |x>|y> = |x> U_x |y>
#this property will for instance be used to make sure that uncomputations,
#which are depending on the state of R indeed stil recieve the original state of R
#when the uncomputation is performed. An error will be thrown if in the meantime,
#there have been non-permeable operations on R

#But what is the general strategy in determining this property?
#To understand this, note that if we order the computational basis vectors,
#such that all the qubits which are supposed to be checked for permeability are at the end

#In this basis, the unitary is block diagonal, where the blocks on the diagonal are the U_x matrices
#we check the block diagonality by first calculating the block matrices and then create a boolean 2^n x 2^n
#matrix, where the entry is False if the corresponding block matrix is 0

#We then check the diagonality of this matrix and return True if so
def is_permeable(gate, qubit_indices):
    
    for i in qubit_indices:
        if gate.permeability[i] == False:
            return False
        elif gate.permeability[i] is None:
            break
    else:
        return True
    
    from qrisp.circuit import PTControlledOperation
    from qrisp.simulator import calc_embedded_unitary
    
    if gate.definition:
        
        qc = gate.definition
        for instr in qc.data:
            relevant_qubits = set([qc.qubits[k] for k in qubit_indices]).intersection(instr.qubits)
            
            if not relevant_qubits:
                continue
            
            if instr.op.name == "measure":
                return False
            
            local_qubit_indices = [instr.qubits.index(qb) for qb in relevant_qubits]
            
            if not is_permeable(instr.op, local_qubit_indices):
                break
        
        else:
            for i in qubit_indices:
                gate.permeability[i] = True
        
            return True

    #Create qubit order
    qubit_order = []
    for i in range(gate.num_qubits):
        if not i in qubit_indices:
            qubit_order.append(i)
            
    #The qubits which are to be checked come first
    qubit_order = qubit_indices + qubit_order
    
    #Invert permutation
    qubit_order = [qubit_order.index(i) for i in range(len(qubit_order))]
    
    #Calculate unitary
    unitary = calc_embedded_unitary(gate, gate.num_qubits, qubit_order)
    
    if not isinstance(unitary, np.ndarray):
        unitary = unitary.to_array()
    
    #Calculate boolean matrix
    bbm = get_boolean_block_matrix(unitary, 2**(gate.num_qubits - len(qubit_indices)))
    #Check non zero off diagonal elements
    off_diagonal_non_zero_count = np.count_nonzero(bbm - np.diag(np.diagonal(bbm)))
    #Return result
    
    
    from numpy.linalg import norm
    result = not bool(off_diagonal_non_zero_count)
    
    if result:
        for i in qubit_indices:
            gate.permeability[i] = True
    
    
    return result


#Tests if a gate is qfree where qfree is the property that is introduced in SILQ
#meaning that an operation maps any computational basis state to a computational basis state
def is_qfree(gate):
    
    
    if not gate.is_qfree is None:
        return gate.is_qfree
    
    if gate.definition:
        
        for instr in gate.definition.data:
            if not is_qfree(instr.op):
                break
        else:
            gate.is_qfree = True
            return True
    
    unitary = np.around(gate.get_unitary(), 5)
    #Check if the unitary has only one non-zero entry in each line
    for i in range(unitary.shape[0]):
        if np.count_nonzero(unitary[i,:]) != 1:
            gate.is_qfree = False
            return False
        
    gate.is_qfree = True
    return True



    
#Splits the given matrix into blocks of size block_size
#Returns an object array where each entry contains the corresponding block
def get_block_matrix(matrix, block_size):
    
    matrix = np.array(matrix)
    if matrix.shape[0] != matrix.shape[1]:
        raise Exception("Tried to decompose non-square blockmatrix")
        
    if matrix.shape[0]%block_size:
        raise Exception("Matrix shape is no integer multiple of block size")
        
    
    block_amount = int(matrix.shape[0]/block_size)
    result = np.zeros((block_amount, block_amount), dtype = "object")
    
    for i in range(block_amount):
        for j in range(block_amount):
            result[i,j] = matrix[i*block_size:(i+1)*block_size, j*block_size:(j+1)*block_size]
    
    return result


#Recieves a matrix which is turned into a block matrix and then
#returns a matrix of boolean entries where the entry is False, if
#the corresponding block is the 0 matrix and True otherwise
def get_boolean_block_matrix(matrix, block_size):
    
    
    block_matrix = get_block_matrix(matrix, block_size)
        
    block_amount = int(matrix.shape[0]/block_size)
    
    result = np.zeros((block_amount, block_amount))
    
    for i in range(block_amount):
        for j in range(block_amount):
            
            if np.linalg.norm(block_matrix[i,j]) > 1E-4:
                result[i,j] = 1
    
    return result