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


from qrisp.interface.circuit_converter import convert_circuit
from qrisp.quantum_network.interface.codegen import QuantumNetworkService
from qrisp.quantum_network.interface.codegen.ttypes import Message
from qrisp.interface import PortableQubit
from qrisp import Qubit

class QuantumNetworkClient(QuantumNetworkService.Client):
    """
    This class allows to connect to Qrisp quantum networks.
    
    Parameters
    ----------
    name : str
        A nickname such as "Alice" or "Bob", which is used to direct the communication of other users to you.
    socket_ip : str
        The IP adress of the QuantumNetworkServer.
    port : int, optional
        The port of the QuantumNetworkServer. The default is 7070.
        
    Examples
    --------
    
    We create a QuantumNetworkServer listening on the localhost IP-address an connect the client.
    
    >>> from qrisp.quantum_network import QuantumNetworkServer, QuantumNetworkClient
    >>> local_server = QuantumNetworkServer("127.0.0.1", background = True)
    >>> local_server.start()
    >>> client = QuantumNetworkClient(name = "alice", socket_ip = "127.0.0.1")
    
    
    """
    
    
    def __init__(self, name, socket_ip, port = None):
        
        from thrift.transport import TSocket, TTransport
        from thrift.protocol import TBinaryProtocol
        
        if port == None:
            port = 7070
        #Create the transport for the User Interface to the server
        self.transport = TSocket.TSocket(socket_ip, port)
        # Buffering is critical. Raw sockets are very slow
        self.transport = TTransport.TBufferedTransport(self.transport)
        # Wrap in a protocol
        protocol = TBinaryProtocol.TBinaryProtocol(self.transport)
        self.socket_ip = socket_ip        
        # Create a client to use the protocol encoder
                        
        super().__init__(protocol)
        
        # Connect!
        self.transport.open()
        
        
        self.name = name
        self.register_client(name)
    
    #Destructor closes the transport
    def __del__(self):
    
        self.transport.close()
        
    def run(self, qc, shots = None):
        r"""
        Runs a QuantumCircuit on the client's backend. Note that QuantumNetwork simulations do not support multiple shots, as the quantum state of the network is stored and updated everytime a client sends a query. Multiple shots could yield differing measurement outcomes, which implies an ambiguous quantum state of the network. Nevertheless the results are returned in the form of a dictionary in order to comply with the quantum circuit execution backend interface.
        
        Note that it is possible to submit QuantumCircuits which contain qubits, that have not been requested previously. In this case, the qubits names are internally extended by the string "@client_name"  (if they aren't extended in this way already). This is to allow multiple clients to submit circuits with matching qubit names.

        Parameters
        ----------
        qc : QuantumCircuit
            The QuantumCircuit to run.

        Returns
        -------
        res : dict
            A dictionary containing a single key/value pair where the key represents the measurement outcome.
            
        
        Examples
        --------
        
        We create a local QuantumNetworkServer, connect a client and run a QuantumCircuit
        
        >>> from qrisp.quantum_network import QuantumNetworkServer, QuantumNetworkClient
        >>> local_server = QuantumNetworkServer("127.0.0.1", background = True)
        >>> local_server.start()
        >>> client = QuantumNetworkClient(name = "alice", socket_ip = "127.0.0.1")
        >>> from qrisp import QuantumCircuit
        >>> qc = QuantumCircuit(2)
        >>> qc.h(0)
        >>> qc.cx(0,1)
        >>> qc.measure([0,1])
        >>> client.run(qc)
        {'11': 1}
        
        We perform another shot
        
        >>> client.run(qc)
        {'01': 1}

        After applying the first run command, the quantum state is

        .. math::        
            \ket{\psi} = \ket{1}\ket{1}
        
        Applying the Hadamard yields
        
        .. math::        
            \text{H}_0 \ket{\psi} = \frac{1}{\sqrt{2}}(\ket{0} - \ket{1})\ket{1}
        
        The CX yields
        
        .. math::        
            \text{CX}_{01} \text{H}_0 \ket{\psi} = \frac{1}{\sqrt{2}}(\ket{0}\ket{1} - \ket{1}\ket{0})
        
        Finally, the measurement collapsed the state into the first summand.
        
        """
        converted_circuit = convert_circuit(qc, "thrift")
        
        res = super().run(converted_circuit, self.name)
        
        return res
    

        
        qc = super().get_clear_qc(self.name)
        
        return convert_circuit(qc, "qrisp")

    def send_qubits(self, recipient, qubits, annotation):
        """
        Sends the specified qubits to another participant of the network.

        Parameters
        ----------
        recipient : str
            The recipients name.
        qubits : list[Qubit]
            The list of qubits to send.
        annotation : str
            A string containing an arbitrary message that is available for the reciepient.

        
        Examples
        --------
        
        We create a QuantumNetworkServer, connect two clients and distribute a bell pair another client.
        
        >>> from qrisp.quantum_network import QuantumNetworkServer, QuantumNetworkClient
        >>> local_server = QuantumNetworkServer("127.0.0.1", background = True)
        >>> local_server.start()
        >>> alice_client = QuantumNetworkClient(name = "alice", socket_ip = "127.0.0.1")
        >>> bob_client = QuantumNetworkClient(name = "bob", socket_ip = "127.0.0.1")
        
        Now we create the bell pair:
        
        >>> from qrisp import QuantumCircuit
        >>> qc = QuantumCircuit(2)
        >>> qc.h(0)
        >>> qc.cx(0,1)
        >>> alice_client.run(qc)
        {'': 1}
        
        Send one of the qubits to Bob:
            
        >>> alice_client.send_qubits("bob", [qc.qubits[0]], annotation = "Happy birthday!")
        
        Now Bob can measure
        
        >>> messages = bob_client.inbox()
        >>> received_qubits, annotation = messages[0]
        >>> qc = bob_client.get_clear_qc()
        >>> qc.measure(received_qubits)
        >>> bob_client.run(qc)
        {'1': 1}
        
        After Bob's measurement, we expect Alice's measurement to yield the same due to `spukhafter Fernwirkung <https://en.wikipedia.org/wiki/Quantum_entanglement#History>`_:
            
        >>> qc = alice_client.get_clear_qc()
        >>> qc.measure(qc.qubits)
        >>> alice_client.run(qc)
        {'1': 1}
        


        """
        
        msg = Message([PortableQubit(qb.identifier) for qb in qubits], annotation)
        
        super().send_qubits(self.name, recipient, msg)

    
    def request_qubits(self, amount):
        """
        Creates the specified amount of qubits for the client's backend.

        Parameters
        ----------
        amount : int
            The amount of qubits to create.

        Returns
        -------
        list[Qubit]
            The qubits created.
            
        Examples
        --------

        >>> from qrisp.quantum_network import QuantumNetworkServer, QuantumNetworkClient
        >>> local_server = QuantumNetworkServer("127.0.0.1", background = True)
        >>> local_server.start()
        >>> client = QuantumNetworkClient(name = "alice", socket_ip = "127.0.0.1")            
        >>> qb_list = client.request_qubits(4)
        >>> print(qb_list)        
        [Qubit(qb_0@alice), Qubit(qb_1@alice), Qubit(qb_2@alice), Qubit(qb_3@alice)]
        
        """
        
        qb_list = super().request_qubits(amount, self.name)
        
        return [Qubit(qb.identifier) for qb in qb_list]

    def get_clear_qc(self):
        """
        Returns a QuantumCircuit containing all the qubits that belong to the client at the moment.
        
        Note that the qubit names of the run QuantumCircuit are internally extended by the string "@client_name"  (if they aren't extended in this way already). This is to allow multiple clients to submit circuits with matching qubit names.
        
        Note that 

        Returns
        -------
        QuantumCircuit
            An empty QuantumCircuit containing the qubits, this client may operate on.
            
        Examples
        --------
        
        We set up a quantum network, run some QuantumCircuits and retrieve the clear QuantumCircuit:

        >>> from qrisp.quantum_network import QuantumNetworkServer, QuantumNetworkClient
        >>> local_server = QuantumNetworkServer("127.0.0.1", background = True)
        >>> local_server.start()
        >>> client = QuantumNetworkClient(name = "alice", socket_ip = "127.0.0.1")            
        >>> from qrisp import QuantumCircuit
        >>> qc_0 = QuantumCircuit(1)
        >>> qc_0.x(0)
        >>> print(qc_0)
              ┌───┐
        qb_8: ┤ X ├
              └───┘
        >>> client.run(qc_0)
        >>> qc_1 = QuantumCircuit(1)
        >>> qc_1.h(0)
        >>> print(qc_1)
                ┌───┐
         qb_26: ┤ H ├
                └───┘
        >>> client.run(qc_1)
        >>> print(client.get_clear_qc().qubits)
        [Qubit(qb_58@alice), Qubit(qb_77@alice)]
        
        """
        qc = super().get_clear_qc(self.name)
        
        return convert_circuit(qc, "qrisp")


    def get_overall_qc(self):
        """
        Retrieves the QuantumCircuit of all operations that happened in the network so far.

        Returns
        -------
        QuantumCircuit
            The QuantumCircuit containing all operations that happened in this QuantumNetwork.

        Examples
        --------
        
        We assume that the commands of the example of the :meth:`send_qubits <qrisp.quantum_network.QuantumNetworkClient.send_qubits>` method have been executed.
        
        >>> print(alice_client.get_overall_qc())
                      ░ ┌───┐      ░ ┌─┐ ░    
        qb_34@alice: ─░─┤ H ├──■───░─┤M├─░────
                      ░ └───┘┌─┴─┐ ░ └╥┘ ░ ┌─┐
        qb_35@alice: ─░──────┤ X ├─░──╫──░─┤M├
                      ░      └───┘ ░  ║  ░ └╥┘
           cb_bob_0: ═════════════════╩═════╬═
                                            ║ 
         cb_alice_0: ═══════════════════════╩═
                                              
                                    
        """
        
        
        qc = super().get_overall_qc()
        
        return convert_circuit(qc, "qrisp")
    
    def inbox(self):
        """
        Returns a list of tuples containing the received qubits and their annotations.

        Returns
        -------
        res_list : list[tuple[list[Qubit], str]]
            A list of tuples containing the send qubits and the annotations.
        
        
        Examples
        --------
        
        We send a qubit from one client to another and inspect the inbox
        
        >>> from qrisp.quantum_network import QuantumNetworkServer, QuantumNetworkClient
        >>> local_server = QuantumNetworkServer("127.0.0.1", background = True)
        >>> local_server.start()
        >>> alice_client = QuantumNetworkClient(name = "alice", socket_ip = "127.0.0.1")
        >>> bob_client = QuantumNetworkClient(name = "bob", socket_ip = "127.0.0.1")
        
        Prepare the qubit which will be sent to bob:
        
        >>> from qrisp import QuantumCircuit
        >>> qc = QuantumCircuit(1)
        >>> qc.h(0)
        >>> alice_client.run(qc)
        {'': 1}
        
        Send the qubit to Bob:
            
        >>> alice_client.send_qubits("bob", [qc.qubits[0]], annotation = "Merry christmas!")
        
        Now Bob can check his inbox
        
        >>> messages = bob_client.inbox()
        >>> received_qubits, annotation = messages[0]
        >>> print(received_qubits)
        [Qubit(qb_20@alice)]
        >>> print(annotation)
        Merry christmas!

        """
        
        inbox_list = super().inbox(self.name)
        
        res_list = []
        
        for i in range(len(inbox_list)):
            
            res_list.append(([Qubit(qb.identifier) for qb in inbox_list[i].qubits], inbox_list[i].annotation))
        return res_list           
