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


#This module contains the necessary tools to merge QuantumSessions

#Due to the interplay with QuantumEnvironments and the behavior
#that QuantumEnvironments automatically detect QuantumSessions that operated
#within them there is quite a bit of complexity involved here.

#The idea to merge two session qs_0 and qs_1 is to first add the appropriate 
#amount of qubits/clbits to qs_0, append all the instructions and 
#fix the other attributes Then we let all the attributes of qs_1 point 
#to the corresponding attributes of qs_0 such that any change on one quantum session
#also appears in the other quantum session

#Therefore two QuantumSessions are declared as identical (or merged) if their attributes
#(like .data, .qubits) point to the same lists as the other one. 
#The == operator for QuantumSessions checks if the .data attribute points to the same list

#If a merge is requested, the first step is to "engage" the QuantumSessions
#in the environment stack. By default QuantumEnvironments create their own QuantumSession
#and "engaging" here means that if merge is called on a QuantumSession within an environment,
#this QuantumSession will be merged with all the QuantumSessions of the environment stack.
def merge_sessions(qs_0, qs_1):
    
    #We now iterate through every QuantumEnvironment and insert it's data into the correct original_data attribute
    for qs in [qs_0, qs_1]:
        
        #We need to find the environment where the env_qs quantum session is not merged into qs.
        #This implies that the instructions of qs have been appended in this environment's parent.
        #Therefore all the data needs to go into the original_data attribute of this environment.
        for env in qs.env_stack:
            if not env.env_qs == qs:
                env.original_data.extend(qs.data)
                qs.data = []
                merge_sessions_inner(qs, env.env_qs)

    merge_sessions_inner(qs_0, qs_1)
    
    

def merge_sessions_inner(qs_0, qs_1, merge_env_stack_ = True):
    if qs_0 == qs_1:
        return
    
    from qrisp.quantum_network import QuantumNetworkSession
    
    
    if isinstance(qs_0, QuantumNetworkSession):
        
        if not qs_0.backend is qs_1.backend:
            if isinstance(qs_1, QuantumNetworkSession):
                raise Exception("Tried to merge two QuantumNetworkSessions containing differing clients")
            
            if not qs_1.backend is None:
                raise Exception("Tried to merge a QuantumNetworkSession with a QuantumSession with a non-trivial backend")
            
        qs_1.backend = qs_0.backend
        qs_1.inbox = qs_0.inbox
        qs_1.__class__ = qs_0.__class__
    
    if isinstance(qs_1, QuantumNetworkSession):
        
        if not qs_1.backend is qs_0.backend:
            
            if isinstance(qs_0, QuantumNetworkSession):
                raise Exception("Tried to merge two QuantumNetworkSessions containing differing clients")
            
            if not qs_0.backend is None:
                raise Exception("Tried to merge a QuantumNetworkSession with a QuantumSession with a non-trivial backend")
            
        qs_0.backend = qs_1.backend
        qs_0.inbox = qs_1.inbox
        qs_0.__class__ = qs_1.__class__
    
    
    
    #The problem we face here is the following:
    #If two sessions a, b are merged and b is merged with another session c
    #at a later point, the attributes of b are updated by that merge, 
    #but the attributes of a are not. To solve this problem, we declare one of the merge
    #participants the "parent" of the merge. The parent has an attribute called "shadow_sessions"
    #that contain the sessions that the parent has been merged with. 
    #The attribute updates are then performed on all shadow sessions.
    
    #To make sure we only operate on parent sessions, we check the parent attribute
    while hasattr(qs_0, "parent"):
        qs_0 = qs_0.parent
    
    
    while hasattr(qs_1, "parent"):
        qs_1 = qs_1.parent
    
    resolve_naming_collisions(qs_0, qs_1)
    
    intersecting_qubits = set([qb.identifier for qb in qs_0.qubits]).intersection([qb.identifier for qb in qs_1.qubits])
    
    if len(intersecting_qubits):
        raise Exception(f"Tried to merge two QuantumSessions containing identically named qubits {intersecting_qubits}")
    
    if len(qs_0.env_stack) < len(qs_1.env_stack):
        qs_0, qs_1 = qs_1, qs_0
        
        
    if merge_env_stack_:
        merge_env_stack(qs_0, qs_1)
    
        if qs_0 == qs_1:
            return
    
    #The qs_tracker list should only contain one Session instance per equivalence class.
    #Ie. If two sessions are merged only one of them should appear in the the qs_tracker
    i = 0
    while i < len(qs_0.qs_tracker):
        if qs_0.qs_tracker[i]() is None:
            qs_0.qs_tracker.pop(i)
        elif qs_0.qs_tracker[i]() == qs_1:
            qs_0.qs_tracker.pop(i)
        else:
            i+= 1

    #Each environment contains a list of QuantumSessions that when it was entered.
    #This list determines if a QuantumSession is engaged into an QuantumEnvironment, if merge
    #is called from inside the QuantumEnvironment.
    #We only want to have one of the two QuantumSessions that are merged in this list, which is  
    #resolved by the identify_sessions_in_environment
    for i in range(len(qs_1.env_stack)):
        identify_sessions_in_environment(qs_1.env_stack[-i-1], qs_0, qs_1)
    
    
    if not qs_0.backend is None and not qs_1.backend is None:
        if id(qs_0.backend) != id(qs_1.backend):
            raise Exception("Tried to merge QuantumSessions with differing, non-trivial default backends.")
    
    if qs_0.backend is None:
        qs_0.backend = qs_1.backend

    if set(qs_0.clbits).intersection(qs_1.clbits):
        raise Exception("Tried to merge sessions with common classical bits")
    
    for qb in qs_1.qubits: qb.qs = weakref.ref(qs_0)
    for cb in qs_1.clbits: cb.qs = weakref.ref(qs_0)
    
    qs_0.qubits.extend(qs_1.qubits)
    qs_0.clbits.extend(qs_1.clbits)
    
    qs_0.data.extend(qs_1.data)
    
    object.__setattr__(qs_1, "data", qs_0.data)    
    
    #Update qv_list
    while len(qs_1.qv_list):
        
        qv = qs_1.qv_list.pop()
        #Patch quantum session attribute
        qv.qs = qs_0
        
        qs_0.qv_list.append(qv)
    
    
    qs_0.deleted_qv_list.extend(qs_1.deleted_qv_list)
    
    reorder_quantum_variables(qs_0)
    
    #Add variables to the uncomputation stack
    qs_0.uncomp_stack.extend(qs_1.uncomp_stack)
    
    qs_0.shadow_sessions.extend(qs_1.shadow_sessions + [weakref.ref(qs_1)])
    
    i = 0
    while i < len(qs_0.shadow_sessions):
        qs = qs_0.shadow_sessions[i]()
        if qs is None:
            qs_0.shadow_sessions.pop(i)
            continue
        
        qs.parent = qs_0
        
        # Update all the attributes
        for var in qs_0.__dict__.keys():
        
            qs.__dict__[var] = qs_0.__dict__[var]
            
        i += 1
            
            

def resolve_naming_collisions(qs_0, qs_1):
    
    qs_names_0 = [qv.name for qv in qs_0.qv_list + qs_0.deleted_qv_list]
    qs_names_1 = [qv.name for qv in qs_1.qv_list + qs_1.deleted_qv_list]
    
    intersecting_qv_names = set(qs_names_0).intersection(qs_names_1)

    if len(intersecting_qv_names):
        
        for qv_name in intersecting_qv_names:
            
            
            qv_0_index = qs_names_0.index(qv_name)
            
            qv_0 = (qs_0.qv_list + qs_0.deleted_qv_list)[qv_0_index]
            
            qv_1_index = qs_names_1.index(qv_name)
            qv_1 = (qs_1.qv_list + qs_1.deleted_qv_list)[qv_1_index]
        
            if qv_1.user_given_name:
                
                if qv_0.user_given_name:
                    raise Exception("Tried to merge QuantumSession containing "
                                    f"identically named QuantumVariables {qv_1.name}")
                
                qv_0, qv_1 = qv_1, qv_0
                
            else: 
                if qv_0.creation_time > qv_1.creation_time:
                    if not qv_1.user_given_name:
                        qv_0, qv_1 = qv_1, qv_0
        
        
            proposed_new_name = qv_1.name + "_1"
            k = 1
            
            while proposed_new_name in qs_names_0 + qs_names_1:
                k += 1
                proposed_new_name = qv_1.name + "_" + str(k)
            
            qv_1.name = proposed_new_name
            
            
            for k in range(len(qv_1.reg)):
                qv_1.reg[k].identifier = qv_1.name + "." + str(k)


def merge_env_stack(qs_0, qs_1):
    
    #We need to find the environment where the env_qs quantum session is not merged into qs.
    #This implies that the instructions of this session have been appended in this environment's parent.
    #Therefore all the data needs to go into the original_data attribute of this environment.
    for env in qs_0.env_stack:
        
        if env.env_qs == qs_0 and not env.env_qs == qs_1:
            if env in qs_1.env_stack:
                env.original_data.extend(qs_1.data)
                qs_1.data = []
            
        if env.env_qs == qs_1 and not env.env_qs == qs_0:
            for env_ in qs_0.env_stack:
                if not env_.env_qs == qs_0:
                    
                    env_.original_data.extend(qs_0.data)
                    qs_0.data = []
                    merge_sessions_inner(qs_0, env_.env_qs, False)
                if env.env_qs == env_.env_qs:
                    break

def reorder_quantum_variables(qs):

    overall_qv_list = qs.qv_list + qs.deleted_qv_list
    
    qs.qv_list.sort(key = lambda x : x.creation_time)
    
    overall_qv_list.sort(key = lambda x : x.creation_time)
    
    sorted_qubit_list = []
    for qv in overall_qv_list:
        sorted_qubit_list += qv.reg
        
    sorted_qubit_list.extend(list(set(qs.qubits) - set(sorted_qubit_list)))
    
    qs.qubits = sorted_qubit_list


def identify_sessions_in_environment(env, qs_0, qs_1):
    
    i = 0
    while i < len(env.active_qs_list):
        if env.active_qs_list[i] == qs_1:
            env.active_qs_list.pop(i)
            continue
        i += 1
    


def multi_session_merge(input_list):
    
    session_list = input_list
    session_list.sort(key = lambda x : -len(x.data))
    
    if not len(session_list):
        from qrisp.core import QuantumSession
        return QuantumSession()
    
    
    for i in range(1, len(session_list)):
        merge_sessions(session_list[0], session_list[i])
    
    return session_list[0]

def recursive_qs_search(input):
        
    if isinstance(input, str):
        return []
    
    from qrisp.core import QuantumSession
    
    if hasattr(input, "__iter__"):
        iterable = True
        
    #This case would in principle allow for also searching objects
    #this however requires a considerable chunk of resources
    elif hasattr(input, "__dict__") and not isinstance(input, QuantumSession):
        iterable = False
        # input = input.__dict__
    else:
        iterable = False
    
    result = []
    if iterable:
        if isinstance(input, dict):
            for key in input.keys():
                result += recursive_qs_search(key)
                result += recursive_qs_search(input[key])
        else:
            for i in range(len(input)):
                result += recursive_qs_search(input[i])
    else:
        if isinstance(input, QuantumSession):
            result = [input]
        else:
            try:
                if isinstance(input(), QuantumSession):
                    result = [input()]
            except TypeError:
                pass
            
            try:
                if isinstance(input.qs(), QuantumSession):
                    result = [input.qs()]
            except AttributeError:
                pass
            
            

    return result

def recursive_qv_search(input):
        
    if isinstance(input, str):
        return []
    from qrisp.core import QuantumVariable
    
    if hasattr(input, "__iter__"):
        iterable = True
    elif hasattr(input, "__dict__") and not isinstance(input, QuantumVariable):
        iterable = False
        # input = input.__dict__
    else:
        iterable = False
        
    result = []
    if iterable:
        if isinstance(input, dict):
            for key in input.keys():
                result += recursive_qv_search(key)
                result += recursive_qv_search(input[key])
        else:
            for i in range(len(input)):
                result += recursive_qv_search(input[i])
    else:
        if isinstance(input, QuantumVariable):
            result = [input]

    return result


def merge(*args):
    session_list = recursive_qs_search(args)
    return multi_session_merge(session_list)
