"""
/*********************************************************************
* 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
import networkx as nx
from numba import njit, prange

from qrisp.circuit import QuantumCircuit, Qubit, PTControlledOperation, ControlledOperation, transpile, QubitAlloc, Instruction, fast_append
from qrisp.misc import get_depth_dic, retarget_instructions, reduce_depth



#The purpose of this function is to dynamically (de)allocate qubits when they are need or not needed anymore.
#The qompiler function knows when a qubit is ready to deallocate (ie. it is in |0> state) due to a gate called QubitDealloc.
#After some preprocessing steps, this function will iterate through the data and dynamically (de)allocate if it is required.
#Another feature of this function is that it replaces mcx gates where the method is set to "auto" with implementations that
#fit the situation ie. it takes into consideration how many clean/dirty ancilla qubits are available and synthesizes an
# mcx implementation based on this information.

#The parameter workspace can be given to extend the amount of qubits of the resulting QuantumCircuit by that amount.
#This extension can be helpfull because it gives the compiler more options to pick when allocating a new qubit.
#Since the allocation is based on a depth-reduction heuristic, giving more workspace results in many case in less depth circuits.

#It is however not only depth that can be improved by an increased workspace. Due to more clean/dirty ancillae beeing available,
#in many cases it is also possible to generate more efficient mcx implementations, thus also reducing the gate count.
def qompiler(qs, workspace = 0, disable_uncomputation = True, intended_measurements = [], cancel_qfts = True):
    
    if len(qs.data) == 0:
        return QuantumCircuit(0)
    
    with fast_append():
        qc = qs.copy()
        
        if not disable_uncomputation:
            local_qvs = qs.get_local_qvs()
            
            from qrisp.uncomputation.unqomp import uncompute_qc
            
            for qv in local_qvs:
                try:
                    qc = uncompute_qc(qc, qv.reg)
                except:
                    print(f"Warning: Automatic uncomputation for {qv.name} failed")
            
        if intended_measurements:
            #This function reorders the circuit such that the intended measurements can be executed as early as possible
            #additionally, any instructions that are not needed for the intended measurements are removed
            try:
                qc = measurement_reduction(qc, intended_measurements)
            except Exception as e:
                if "Unitary of operation " not in str(e):
                    raise e
    
        def reordering_transpile_predicate(op):
            if (isinstance(op, PTControlledOperation) and op.base_operation.name in ["x"]) or isinstance(op, LogicSynthGate):
                return False
            return True
    
        
        from qrisp.logic_synthesis import LogicSynthGate
        
        if cancel_qfts:
            # The first step is to cancel adjacent QFT gates, which are inverse to each other
            # This can happen alot because of the heavy use of Fourier arithmetic
            qft_transpile_predicate = lambda op : not "QFT" in op.name and reordering_transpile_predicate(op)
            
            qc = transpile(qc, transpile_predicate = qft_transpile_predicate)
            
            qc = qft_cancellation(qc)
        
        #We now reorder the transpiled QuantumCircuit. Reordering is performed based on the DAG representation
        #of Unqomp. The advantage of this representation is that it abstracts away non-trivial commutation relations
        #between permeable gates. The actual ordering is performed by performing a topological sort on this dag.
        #Contrary to unqomp, we don't use a modified version of Kahns algorithm here (for more information check 
        #the topological sort function)
    
        #The reordering process aims to find an order that minimizes that maximum amount of qubits that is needed.
        #Losely speaking it tries to reorder the data that allocations are performed as late as possible and deallocations
        #are performed as early as possible.
    
        #Note that the more aggressive we transpile, the more optimization is possible through reordering ie. less qubits.
        #By reordering we however also destroy any previous order which might have been intentionally picked to optimize depth.
        #In summary: Transpiling more aggressively leads to less qubits but more depth.#
        
        #Only letting mcx and logic synthesis survive has shown to be a good compromise
        
        transpiled_qc = transpile(qc, transpile_predicate = reordering_transpile_predicate)
        reordered_qc = reorder_qc(transpiled_qc)
        
        #Transpile logic synthesis
        reordered_qc = transpile(reordered_qc, transpile_predicate = lambda op : not (isinstance(op, PTControlledOperation) and op.base_operation.name in ["x"]))
        
        # We combine adjacent single qubit gates
        if not reordered_qc.abstract_params:
            reordered_qc = combine_single_qubit_gates(reordered_qc)
        
        #We now determine the amount of Qubits the circuit will need    
        required_qubits = 0
        max_required_qubits = 0
        
        for instr in reordered_qc.data:
            if instr.op.name == "qb_alloc":
                required_qubits += 1
            elif instr.op.name == "qb_dealloc":
                required_qubits -= 1
            if max_required_qubits < required_qubits:
                max_required_qubits += 1
        
        #Create a the QuantumCircuit. Note that we make sure that the Qubit naming is consistent,
        #since we don't want any non-deterministic elements in the function, as this can hinder bugfixing
        qc = QuantumCircuit()
        for i in range(workspace + max_required_qubits):
            qc.add_qubit(Qubit(str(i)))
    
        
        
        qc.clbits = list(qs.clbits)
    
    
        #This dictionary translates between the qubits of the input QuantumSession and the compiled circuit
        translation_dic = {}
        
        #This list contains the Qubits which are currently not allocated
        free_qb_list = list(qc.qubits)
        
        allocated_qb_list = []
        
        depth_dic = {b : 0 for b in qc.qubits + qc.clbits}
        
        
        #We now iterate through the data of the preprocessed QuantumCircuit
        for i in range(len(reordered_qc.data)):
            
            QuantumCircuit.fast_append = True
            instr = reordered_qc.data[i]
            if instr.op.name == "barrier":
                continue
            
            if instr.op.name == "qb_alloc":
                #We delay the allocation until an actual gate is executed
                allocated_qb_list.append(instr.qubits[0])
                continue
            
            #Check if any of the involved qubits need an allocation
            for qb in instr.qubits:
                
                if not qb in translation_dic:
                    #To pick a good allocation, we determine the depth dic and sort the
                    #available Qubits by their corresponding depth.
                    #Note that we add the hash value in order to prevent non-deterministic behavior
                    free_qb_list.sort(key = lambda x : depth_dic[x] + qc.qubits.index(x)*1E-5)
                    
                    translation_dic[qb] = free_qb_list.pop(0)
                    
                    qc.append(QubitAlloc(), [translation_dic[qb]])
    
            if instr.op.name == "qb_dealloc":
                #For the deallocation, we simply remove the qubits from the translation dic
                #and append it to the free_qb_list
                free_qb_list.append(translation_dic[instr.qubits[0]])
                allocated_qb_list.remove(instr.qubits[0])
                
                qc.append(instr.op, [translation_dic[qb] for qb in instr.qubits], instr.clbits)
                
                del translation_dic[instr.qubits[0]]
                
                continue
    
            if isinstance(instr.op, ControlledOperation) and instr.op.base_operation.name == "x" and instr.op.method == "auto" and len(instr.qubits) > 3:
                
                #This section deals with the automatic recompilation of mcx gates
                #The reason why we do this is, that it might not be clear how many free ancillae are
                #available for mcx gates. If we set the method to auto in the mcx function,
                #this section is called and replaces the generated mcx gate with an implementation
                #that is better fit to suit the amount of available ancillae.
                #We first determine the free ancillae
                
                
                clean_ancillae = list(free_qb_list)
                clean_ancillae.sort(key = lambda x : depth_dic[x] + qc.qubits.index(x)*1E-5)
                
                dirty_ancillae = list(set(translation_dic.values()) - set([translation_dic[qb] for qb in instr.qubits]))
                dirty_ancillae.sort(key = lambda x : depth_dic[x] + qc.qubits.index(x)*1E-5)
                
                QuantumCircuit.fast_append = False
                #This function generates the data for the hybrid implementation
                compiled_mcx_data = gen_hybrid_mcx_data(instr.qubits[:-1], instr.qubits[-1], instr.op.ctrl_state, clean_ancillae, dirty_ancillae)
                
                
                #We now append the data
                for qb in clean_ancillae + dirty_ancillae:
                    translation_dic[qb] = qb
                    
                for instr in compiled_mcx_data:
                    qc.append(instr.op, [translation_dic[qb] for qb in instr.qubits])
                    update_depth_dic(qc.data[-1], depth_dic)
                
                #And free up the qubits
                for qb in clean_ancillae + dirty_ancillae:
                    del translation_dic[qb]
                
                QuantumCircuit.fast_append = True
            
            #Finally if all of the above cases are not met, we simply append the operation to the translated qubits
            else:
                
                try:
                    qc.data.append(Instruction(instr.op, [translation_dic[qb] for qb in instr.qubits], instr.clbits))
                    
                except KeyError:
                    raise Exception("Found operation " + instr.op.name + " on unallocated qubit during compilation.")
                
                update_depth_dic(qc.data[-1], depth_dic)
        
        #Allocate the qubits which had an alloc gate but no actual gates
        for qb in allocated_qb_list:
            if qb not in translation_dic:
                translation_dic[qb] = free_qb_list.pop(0)
        
        
        #The following code is mostly about renaming and ordering the resulting circuit
        #in order to make the compiled circuit still comprehensible
        
        #We rename the allocated qubits to their name from the quantum session
        for i in range(len(qs.qubits)):
            if qs.qubits[i] in translation_dic:
                translation_dic[qs.qubits[i]].identifier = qs.qubits[i].identifier
                translation_dic[qs.qubits[i]].hash_value = qs.qubits[i].hash_value
    
        #We also want the qubits to be sorted according the order of the QuantumVariables
        #ie. the order of creation
        sorted_qubit_list = []
        for qv in qs.qv_list:
            
            #Due to the measurement reduction feature, not all qubits of quantum variables
            #that are live, are guaranteed to be represented in the translation dic
            #Therefore we need the try - except structure here
            try:
                sorted_qubit_list.extend([translation_dic[qb] for qb in qv.reg])
            except KeyError:
                pass
        
        #Furthermore, not all qubits that have been deallocated in the QuantumSession
        #are guaranteed to be removed from the translation dic, since the 
        #measurement_reduction function might have removed their uncomputation/deallocation gates
        temp = list(set(translation_dic.values()) - set(sorted_qubit_list))
        temp.sort(key = lambda x : x.identifier)
        sorted_qubit_list.extend(temp)
    
        
        #Finally, we rename the deallocated qubits to "workspace"    
        workspace_naming_counter = 0
        td_values = list(translation_dic.values())
        for i in range(len(qc.qubits)):
            if not qc.qubits[i] in td_values:
                qc.qubits[i].identifier = "workspace_" + str(workspace_naming_counter)
                workspace_naming_counter += 1
                sorted_qubit_list.append(qc.qubits[i])
    
        qc.qubits = sorted_qubit_list
        
        reduced_qc = reduce_depth(qc)
        
    if reduced_qc.depth() > qc.depth():
        return qc
    else:
        return reduced_qc


def gen_hybrid_mcx_data(controls, target, ctrl_state, clean_ancillae, dirty_ancillae):
    #This function generates the data for the hybrid mcx implementation
    
    from qrisp.circuit.multi_cx import hybrid_mcx
    from qrisp.core import QuantumVariable
    #Specify QuantumVariables to call mcx function
    control_qv = QuantumVariable(len(controls), name = "control")
    target_qv = QuantumVariable(1, name = "target")
    
    # dirty_ancillae = dirty_ancillae + clean_ancillae
    hybrid_mcx(control_qv, target_qv, ctrl_state = ctrl_state, num_ancilla = len(clean_ancillae), num_dirty_ancilla = len(dirty_ancillae))
    
    #Get the list of used ancillae    
    used_ancillae_set = set(control_qv.qs.qubits) - set(control_qv.reg) - set(target_qv.reg)
    
    #If we used the list() function to transform the set, this introduces a non-deterministic
    #element in the compilation algorithm, which can hamper bugfixing
    used_clean_ancillae = []
    used_dirty_ancillae = []
    
    for qb in control_qv.qs.qubits:
        if qb in used_ancillae_set:
            if "dirty" in qb.identifier:
                used_dirty_ancillae.append(qb)
            else:
                used_clean_ancillae.append(qb)
    
    
    depth_dic = get_depth_dic(control_qv.qs)
    
    used_clean_ancillae.sort(key = lambda x : -depth_dic[x])
    used_dirty_ancillae.sort(key = lambda x : -depth_dic[x])
    
    
    # ancilla_list = list(clean_ancillae)
    
    # for i in range(len(used_ancillae)):
    #     if "yong" in used_ancillae[i].identifier:
    #         used_ancillae.insert(0, used_ancillae.pop(i))
    #         ancilla_list.insert(0, dirty_ancillae[0])
    #         break
    
    #Now retarget the instructions such that they use the appropriate qubits
    data = control_qv.qs.data
    
    retarget_instructions(data, list(control_qv), controls)
    retarget_instructions(data, list(target_qv), [target])
    retarget_instructions(data, used_clean_ancillae, clean_ancillae)
    retarget_instructions(data, used_dirty_ancillae, dirty_ancillae)
    
    i = 0
    #Remove (de)allocation gates
    while i < len(data):
        if data[i].op.name in ["qb_dealloc", "qb_alloc"]:
            data.pop(i)
            continue
        i += 1
        
    return data



def reorder_qc(qc):
    
    from qrisp.uncomputation import dag_from_qc
    G = dag_from_qc(qc, remove_init_nodes = True)
    
    qc_new = qc.clearcopy()
    
    dealloc_identifier = lambda x : x.op.name == "qb_dealloc"
    alloc_identifer = lambda x : x.op.name == "qb_alloc"
    
    # mcx_identifier = lambda x : isinstance(x.op, PTControlledOperation) and x.op.base_operation.name == "x"
    # nmcx_identifier = lambda x : not mcx_identifier(x)
    # sub_sort = lambda G : topological_sort(G, prefer = mcx_identifier, delay = nmcx_identifier)
    # for n in topological_sort(G, prefer = dealloc_identifier, delay = alloc_identifer, sub_sort = sub_sort):
        
    for n in topological_sort(G, prefer = dealloc_identifier, delay = alloc_identifer):
        qc_new.append(n.instr)
    
    return qc_new



#This function performs a topological sort of the graph G where we try to execute any deallocation gates
#as early as possible while still adhering to the topological order.
#We to perform a depth-first search as described here: https://en.wikipedia.org/wiki/Topological_sorting
#According to the Wikipedia page, we are allowed to pick any node as a "starting point" of the DF-search,
#which allows us to have modify the algorithm such that it optimizes the (de)allocation order.
#The general idea is to pick the deallocation nodes as starting points, where we order them,
#such that those deallocation nodes that "require" the least allocation nodes are executed first.
#"Require" here means that there is a causal relationship between the allocation and deallocation nodes,
#ie. there is a path from the allocation node to the deallocation node.

#We can thus determine the amount of allocation nodes required for a deallocation node n by counting, the amount
#of allocation nodes in the "ancestors" subgraph of n.
def topological_sort(G, prefer = None, delay = None, sub_sort = nx.topological_sort):
    """
    Function to perform a topological sort on an Unqomp DAG which allows prefering/delaying specific types of nodes

    Parameters
    ----------
    G : nx.DiGraph
        The Unqomp DAG.
    prefer : function, optional
        Function which returns True, when presented with an Instruction, that should be prefered. The default is the function that returns False on all Operations
    delay : function, optional
        Function which returns True, when presented with an Instruction, that should be delayed. The default is the function that returns False on all Operations
    sub_sort : function, optional
        A function which performs a topological sort, which can sorting preferences of secondary importance. The default is nx.topological_sort.

    Returns
    -------
    lin : list[UnqompNode]
        The linearized list of UnqompNodes. The init nodes are not included.

    """
    
    if prefer is None:
        prefer = lambda x : False
    
    if delay is None:
        delay = lambda x : False
    
    G = G.copy()
    #Collect the prefered nodes
    prefered_nodes = []
    
    for n in G.nodes():
        if prefer(n.instr):
            
            prefered_nodes.append(n)
            
        n.processed = False
    
    #For large scales, finding the ancestors is a bottle neck. We therefore use a jitted version
    if len(G)*len(prefered_nodes) > 10000:
        anc_lists = ancestors(G, prefered_nodes)
        node_ancs = {prefered_nodes[i] : anc_lists[i] for i in range(len(prefered_nodes))}
    else:
        node_ancs = {prefered_nodes[i] : list(nx.ancestors(G, prefered_nodes[i])) for i in range(len(prefered_nodes))}
    #Determine the required delay nodes for each prefered nodes
    required_delay_nodes = {n : [] for n in prefered_nodes}
    
    for n in prefered_nodes:
        for k in node_ancs[n]:
            if delay(k.instr):
                required_delay_nodes[n].append(k)
    
    required_delay_nodes = {n : set(required_delay_nodes[n]) for n in prefered_nodes}
            
    #Generate linearization
    lin = []
    
    while prefered_nodes:
        
        #Sort nodes accordingly
        prefered_nodes.sort(key = lambda x : len(required_delay_nodes[x]))
        
        
        node = prefered_nodes.pop(0)
        ancs = []
        
        for n in node_ancs[node] + [node]:
            if n.processed:
                continue
            else:
                n.processed = True
                ancs.append(n)
        
        sub_graph = G.subgraph(ancs)
        
        lin += list(sub_sort(sub_graph))
        
        continue
        for n in prefered_nodes:
            required_delay_nodes[n] = required_delay_nodes[n] - required_delay_nodes[node]
        
    #Linearize the remainder
    remainder = []
    for n in G.nodes():
        if n.processed:
            continue
        else:
            n.processed = True
            remainder.append(n)
    
    # lin += list(sub_sort(G))
    lin += list(sub_sort(G.subgraph(remainder)))
    
    return lin


#Function to combine any sequences of single qubit gates into a single U3
def combine_single_qubit_gates(qc):
    
    def apply_combined_gates(qc_new, gate_list, qb):
        
        if not len(gate_list):
            return
        
        n = len(gate_list)
        
        m = np.eye(2)
        while gate_list:
            gate = gate_list.pop(-1)
            m = np.dot(m, gate.get_unitary())
        
        if np.linalg.norm(m - np.eye(2)) < 1E-10:
            return
            
        if n == 1:
            qc_new.append(gate, [qb])
            return
        
        qc_new.unitary(m, [qb])
        
    qb_dic = {qb : [] for qb in qc.qubits}
    
    qc_new = qc.clearcopy()
    
    for instr in qc.data:
        if len(instr.qubits) > 1 or instr.op.name in ["qb_alloc", "qb_dealloc"] or len(instr.clbits) > 0:
            
            for qb in instr.qubits:
                apply_combined_gates(qc_new, qb_dic[qb], qb)
            qc_new.append(instr)
        else:
            qb_dic[instr.qubits[0]].append(instr.op)
            
    for qb in qc.qubits:
        apply_combined_gates(qc_new, qb_dic[qb], qb)
    
    return qc_new

   

#This function reorders the circuit such that the intended measurements can be executed as early as possible.
#Additionally, any instructions that are not needed for the intended measurements are removed

#Intended measurements has to be a list of qubits

#The strategy is similar to the one presented in reorder_circuit:
#We bring the circuit in the dag representation and perform a topological sort with
#the intended measurement as prefered instructions

#After that we investigate the circuit for instructions that can be removed
def measurement_reduction(qc, intended_measurements):
    
    qc = qc.copy()
    
    #Insert intended measurements into the circuit
    for qb in intended_measurements:
        qc.measure(qb)

    #Generate dag representation    
    from qrisp.uncomputation import dag_from_qc
    G = dag_from_qc(qc, remove_init_nodes = True)
    
    #Create result qc
    qc_new = qc.clearcopy()

    #Define prefered instructions    
    measure_identifier = lambda x : x.op.name == "measure" and x.qubits[0] in intended_measurements
    
    
    #Perform topological sort
    for n in topological_sort(G, prefer = measure_identifier):
        qc_new.append(n.instr)

    #Check which instructions come after the final measurement        
    for i in range(len(qc_new.data))[::-1]:
        if measure_identifier(qc_new.data[i]):
            break
    
    redundant_qc = qc_new.clearcopy()
    
    redundant_qc.data = qc_new.data[i+1:]
    
    

    G = dag_from_qc(redundant_qc, remove_init_nodes = True)
    
    
    # alloc_graph = allocation_graph(redundant_qc)
    
    
    for node in G.nodes():
        
        if node.instr.op.name == "qb_dealloc":
            ancs = nx.ancestors(G, node)
            for pred in ancs:
                if pred.instr.op.name == "qb_alloc":
                    break
            else:
                redundant_qc.data.remove(node.instr)
                # print(f"removed {node.instr}")
                for pred in ancs:
                    
                    try:
                        redundant_qc.data.remove(pred.instr)
                    except ValueError:
                        pass
                    
                    
        
    # print(redundant_qc)    
    
    redundant_instructions = redundant_qc.data
    # print(redundant_instructions)
    
    
    # redundanct_instructions = qc_new.data[i+1:]
    # #Now we need to make sure we don't remove deallocation gates from the data
    # #because this would inflate the qubit count of the compiled circuit
    
    # #The strategy here is that if we find a deallocation gate
    # #we remove any instruction involving the deallocated qubit from the list 
    # #redundant instructions.
    
    # #If we however find an allocation gate, this chain of instructions can be
    # #safely removed and therefore stay in the redundant instructions
    
    # i = 0
    # while i < len(redundant_instructions):
    #     if redundant_instructions[i].op.name == "qb_dealloc":
    #         dealloc_qubit = redundant_instructions[i].qubits[0]
            
    #         j = 0
    #         while j < len(redundant_instructions):
    #             if dealloc_qubit in redundant_instructions[j].qubits:
    #                 if redundant_instructions[j].op.name == "qb_alloc":
    #                     i += 1
    #                     break
                    
    #                 redundant_instructions.pop(j)
    #                 continue
    #             j += 1
    #         continue
    #     i += 1
            
    #We now remove the redundant instructions and the inserted 
    #measurements from the circuit data
    i = 0
    while i < len(qc_new.data):
        
        if measure_identifier(qc_new.data[i]):
            qc_new.data.pop(i)
            continue
        if qc_new.data[i] in redundant_instructions:
            qc_new.data.pop(i)
            continue
        i += 1
    
    
    return qc_new


def allocation_graph(qc):
    from qrisp.uncomputation.unqomp import dag_from_qc
    dag = dag_from_qc(qc, remove_init_nodes = True)
    
    res = nx.DiGraph()
    
    dealloc_nodes = []
    for n in dag.nodes():
        if n.instr.op.name in ["qb_alloc", "qb_dealloc"]:
            res.add_node(n)
        if n.instr.op.name == "qb_dealloc":
            dealloc_nodes.append(n)
    
    for n in dealloc_nodes:
        for anc in nx.ancestors(n, dag):
            if anc.instr.op.name == "qb_alloc":
                res.add_edge(anc, n)
                
    return n
    
    
def update_depth_dic(instruction, depth_dic):
    
    if instruction.op.definition:
        qc = QuantumCircuit()
        qc.qubits = instruction.qubits
        qc.append(instruction)
        instr_list = qc.transpile().data
    else:
        instr_list = [instruction]
    
    # Assign each bit in the circuit a unique integer
    # to index into op_stack.
    # If no bits, return 0


    # Here we are playing a modified version of
    # Tetris where we stack gates, but multi-qubit
    # gates, or measurements have a block for each
    # qubit or cbit that are connected by a virtual
    # line so that they all stacked at the same depth.
    # Conditional gates act on all cbits in the register
    # they are conditioned on.
    # We treat barriers or snapshots different as
    # They are transpiler and simulator directives.
    # The max stack height is the circuit depth.
    
    for instr in instr_list:
        
        if instr.op.name in ["qb_alloc", "qb_dealloc"]:
            continue
        qargs = instr.qubits
        cargs = instr.clbits
        
        levels = []
        # If count then add one to stack heights
        
        for b in qargs + cargs:
            # Add to the stacks of the qubits and
            # cbits used in the gate.
            levels.append(depth_dic[b] + 1)
    
        max_level = max(levels)
            
        for b in qargs + cargs:
            depth_dic[b] = max_level


@njit(cache = True)
def ancestors_jitted(start_index, indptr, indices, node_amount):
    
    to_do_array = np.zeros(node_amount, dtype = np.byte)
    to_do_array[start_index] = 1
    done_array = np.zeros(node_amount, dtype = np.byte)
    
    stack = 1
    while stack:
        
        node = np.argmax(to_do_array)
        to_do_array[node] = 0
        
        for i in range(indptr[node], indptr[node+1]):
            new_node = indices[i]
            if done_array[new_node] == 0:
                to_do_array[new_node] = 1
                stack += 1
        
        done_array[node] = 1
        stack -= 1
        
    return np.nonzero(done_array)[0]

@njit(parallel = True, cache = True)
def ancestors_jitted_wrapper(start_indices, indptr, indices, node_amount):
    
    res = [np.zeros(1, dtype = np.int64)]*len(start_indices)
    for i in prange(len(start_indices)):
        start_index = start_indices[i]
        res[i] = ancestors_jitted(start_index, indptr, indices, node_amount)
    
    return res
        

def ancestors(dag, start_nodes):
    
    node_list = list(dag.nodes())
    
    sprs_mat = nx.to_scipy_sparse_array(dag, format = "csc")
    
    start_indices = []
    for i in range(len(dag)):
        if node_list[i] in start_nodes:
            start_indices.append(i)
    
    res_list_indices = ancestors_jitted_wrapper(np.array(start_indices).astype(np.int32), sprs_mat.indptr, sprs_mat.indices.astype(np.int32), len(dag))
    
    node_list = [[node_list[j] for j in anc_indices] for anc_indices in res_list_indices]
    
    return node_list


# Function to cancel adjacent QFT, which are inverse to each other
# Due to the heavy use of Fourier arithmetic, this can happen alot
# especially if multiple arithmetic operation on a single target are executed
def qft_cancellation(qc):
    
    # The idea is to iterate through the instructions of the circuit
    # and save the last instruction acting on each qubit.
    # If a QFT is appended, we check if the last instruction was an inverse QFT
    # or an allocation and denote the index if so
    # Afterwards, we remove the QFTs or replace them with H gates
    
    last_instr_dic = {qb : None for qb in qc.qubits}
    cancellation_indices = []
    h_replacements = []
    dealloc_replacements = []
    from numpy.linalg import norm
    
    for i in range(len(qc.data)):
        
        instr = qc.data[i]
        if "QFT" in instr.op.name:
            previous_instruction_type = []
            
            for qb in instr.qubits:
                
                previous_instruction = qc.data[last_instr_dic[qb]]
                
                if previous_instruction.op.num_qubits != instr.op.num_qubits:
                    break
                
                if "QFT" in previous_instruction.op.name:
                    
                    if instr.op.num_qubits < 8:
                        
                        unitary_self = instr.op.get_unitary()
                        inv_unitary_other = previous_instruction.op.get_unitary().transpose().conjugate()

                        if bool(norm(unitary_self - inv_unitary_other) < 10**-4):
                            previous_instruction_type.append("QFT")
                        else:
                            break
                        
                    elif hash(instr.op.definition) == hash(previous_instruction.op.definition.inverse()):
                        previous_instruction_type.append("QFT")
                    else:
                        break
                    
                elif previous_instruction.op.name == "qb_alloc":
                   previous_instruction_type.append("alloc")
                else:
                    break
            else:
                if len(set(previous_instruction_type)) == 1:
                    if previous_instruction_type[0] == "QFT":
                        cancellation_indices.append(i)
                        cancellation_indices.append(last_instr_dic[qb])
                    else:
                        h_replacements.append(i)
        
        if instr.op.name == "qb_dealloc":
            
            deallocated_qubit = instr.qubits[0]
            previous_instruction = qc.data[last_instr_dic[deallocated_qubit]]
            
            if "QFT" in previous_instruction.op.name:
                try:
                    previous_instruction.deallocated_qubits[deallocated_qubit] = True
                except AttributeError:
                    previous_instruction.deallocated_qubits = {qb : False for qb in previous_instruction.qubits}
                    previous_instruction.deallocated_qubits[deallocated_qubit] = True
                    dealloc_replacements.append(last_instr_dic[deallocated_qubit])
                        
        for qb in instr.qubits:
            last_instr_dic[qb] = i
    
    
    for i in dealloc_replacements:
        instr = qc.data[i]
        for qb in instr.qubits:
            if not instr.deallocated_qubits[qb]:
                break
        else:
            h_replacements.append(i)
    
    new_qc = qc.clearcopy()
    
    # print(len(qc.data))
    for i in range(len(qc.data)):
        
        if i in h_replacements:
            for qb in qc.data[i].qubits:
                new_qc.h(qb)
            # print("H replacement successfull")
            continue
                
        if i in cancellation_indices:
            continue
        else:
            new_qc.append(qc.data[i])
    # print(len(new_qc.data))    
    # print("====")
    return new_qc
    