import netsquid as ns
import numpy as np
from netsquid.nodes import Node
from netsquid.protocols import NodeProtocol

from nb_protocol import NBProtocolAlice, NBProtocolBob
from utils import EntanglingConnectionOnDemand, create_qprocessor


class BAI:
    '''The base class of the Best Arm Identification algorithm.'''

    def __init__(self, arms):
        self.arms = arms


class QuantumNode(Node):
    '''A quantum node with a quantum port and quantum memory.'''

    def __init__(self, name):
        super().__init__(name)
        # For our purpose, we only need to store 1 qubit
        self.qmemory = create_qprocessor(num_positions=1)
        # Add quantum port used for receiving qubits generated by quantum source
        self.add_ports(["qport"])


class QuantumNodeProtocol(NodeProtocol):
    '''The protocol of each quantum node
       It scans the quantum port and see whether there is an input
       If yes, it stores the qubit in the node's quantum memory
    '''

    def run(self):
        while True:
            qport = self.node.ports["qport"]
            yield self.await_port_input(qport)
            msg = qport.rx_input()
            if msg is not None:
                self.node.qmemory.put(msg.items[0])


class QuantumNetwork:
    '''Simulate the quantum links between two parties, Alice and Bob.
    '''

    def __init__(self, path_num, fidelity_list, noise_model):
        '''Initialize `path_num` number of paths between Alice and Bob.
           The fidelity of each path is provided in `fidelity_list`.
        '''
        super().__init__()
        assert path_num == len(fidelity_list)
        self.best_path = np.argmax(fidelity_list) + 1
        self.noise_model = noise_model

        # Initialize Alice and Bob
        self.alice = QuantumNode(name="Alice")
        self.bob = QuantumNode(name="Bob")
        alice_protocol = QuantumNodeProtocol(self.alice)
        bob_protocol = QuantumNodeProtocol(self.bob)
        alice_protocol.start()
        bob_protocol.start()

        # Initialize quantum channels between Alice and Bob
        # Note that the channels are not connected to Alice and Bob yet
        self.quantum_channels = []
        for i in range(path_num):
            qconn = EntanglingConnectionOnDemand(noise_model, fidelity=fidelity_list[i])
            self.quantum_channels.append(qconn)

    def benchmark_path(self, path, bounces, sample_times):
        """Run network benchmarking along `path`.

           Parameters
           ----------
           path :int
               Path id
           bounces : [int]
               List of number of bounces
           sample_times : dict, int -> int
               Map from number of bounces to its repetition times
        """

        # Get the quantum channel of the corresponding path
        qconn = self.quantum_channels[path - 1]
        # print("Benchmarking channel with fidelity:", qconn.fidelity)

        # Connect Alice and Bob via the quantum channel
        self.alice.ports["qport"].disconnect()
        self.bob.ports["qport"].disconnect()
        self.alice.ports["qport"].connect(qconn.ports["A"])
        self.bob.ports["qport"].connect(qconn.ports["B"])

        while True:
            # Run NB protocol using quantum channel `qconn`
            alice_protocol = NBProtocolAlice(self.alice, bounce=bounces, num_samples=sample_times, qconn=qconn)
            bob_protocol = NBProtocolBob(self.bob)

            # Let Alice and Bob know each other
            alice_protocol.set_target_protocol(bob_protocol)
            bob_protocol.set_target_protocol(alice_protocol)

            alice_protocol.start()
            bob_protocol.start()
            ns.sim_run()

            p, cost = alice_protocol.data_processing()
            # print(f"Estimated parameter p: {p}, cost: {cost}")
            # Only accept the result if the estimated parameter is in a reasonable range
            if p >= 0 and p < 1.5:
                return p, cost