# -*- coding: utf-8 -*-
"""
Tools for the generation of a distributed quantum computer network.

These tools are designed to be controlled using software from
:mod: `~dqc_simulator.software.dqc_control`
"""

import warnings

from netsquid.nodes import Node, DirectConnection, Network
from netsquid.components.cchannel import ClassicalChannel
from netsquid.qubits import ketstates as ks
from netsquid.components.models.delaymodels import FibreDelayModel

from dqc_simulator.hardware.connections import (
    create_black_box_central_source_entangling_link,
    create_classical_fibre_link,
)
from dqc_simulator.hardware.quantum_processors import (
    create_qproc_with_analytical_noise_ionQ_aria_durations_N_standard_lib_gates,
)
from dqc_simulator.util.helper import (
    create_wrapper_with_some_args_fixed,
    filter_kwargs4internal_functions,
)


def link_2_qpus(
    network,
    node_a,
    node_b,
    want_classical_2way_link=True,
    want_extra_classical_2way_link=None,  # can make scheduling
    # easier
    want_entangling_link=True,
    create_entangling_link=None,
    # kwargs for classical links and default entangling link
    node_distance=2e-3,
    # kwargs for default entangling link
    state4distribution=None,
    ent_dist_rate=0,
    # kwargs for non-default entangling_link
    **kwargs4create_entangling_link,
):
    """
    Sets up a link between two QPUs.

    Adds a :class:`~netsquid.netsquid.nodes.connections.Connection` between
    two QPUs and sets up input forwarding between ports on the QPUs so that
    qubits are sent to the right places.

    Parameters
    ----------
    network : netsquid.nodes.network.Network
        The entire network.
    node_a :  object created using netsquid.nodes.Node class
    A network node
    node_b : object created using netsquid.nodes.Node class
    A network node
    create_entangling_link : func, optional
        The function to use in order to create an entangling link. By default,
        :func: `create_black_box_central_source_entangling_link` is used.
    node_distance : float, optional
        Distance between adjacent nodes in km. Used by default value of
        `create_entangling_link`. May not be used, if a different
        `create_entangling_link` function is used.
    want_classical_2way_link : bool, optional
        Whether a two-way classical link should be created between all nodes
        (True) or not (False). Default is True.
    want_extra_classical_2way_link : bool, optional
        Whether an extra two-way classical link should be created between all
        nodes (True) or not (False). Default is True.
    want_entangling_link : bool, optional
        Whether a two-way quantum link should be created between all nodes
        (True) or not (False). Nodes can request entangled pairs using this
        link.
    state4distribution : numpy.ndarray
        The entangled state distributed between nodes when
        an EPR pair is requested. Default is |phi^+> Bell state (formalism not
        fixed to ket). `state4distribution` is used by the default
        `create_entangling_link` function and may be ignored if a different
        entangling link is used.
    ent_dist_rate : float, optional
        The rate of entanglement distribution [Hz]. Used by default value of
        `create_entangling_link`. May not be used, if a different
        `create_entangling_link` function is used. `ent_dist_rate` is used by
        the default `create_entangling_link` function and may be used be
        ignored if a different entangling link is used.
    kwargs4create_entangling_link
        The keyword arguments for the chosen `create_entangling_link` function.
    """
    # instantiating default arguments. Need to do in body of function for mutable
    # objects like a numpy array because default values of functions are defined
    # only when the function is defined (not when it's called) and so mutable
    # default values can be accidentally changed by code outwith the function.
    # See https://stackoverflow.com/questions/10676729/why-does-using-arg-none-
    # fix-pythons-mutable-default-argument-issue
    if want_extra_classical_2way_link is None:
        want_extra_classical_2way_link = want_classical_2way_link
    if state4distribution is None:
        state4distribution = ks.b00
    if (
        create_entangling_link is None
        or create_entangling_link == create_black_box_central_source_entangling_link
    ):
        # The default value (create_black_box_central_source_entangling_link) is
        # wrapped, so that in later code, there is no need to state argument,
        # which is not required by many connections. This keeps the variable
        # name create_entangling_link more generic.
        args2fix = {3: state4distribution}  # fixing 4th argument (index=3) with
        # value state4distribution
        create_entangling_link = create_wrapper_with_some_args_fixed(
            create_black_box_central_source_entangling_link,
            args2fix,
            ent_dist_rate=ent_dist_rate,
        )

    if want_classical_2way_link is False and want_entangling_link is False:
        raise ValueError("""At least one of want_classical_2way_link and
                         want_entangling_link must be True otherwise
                         the function call would be redundant""")
    # Set up classical connection between nodes:
    if want_classical_2way_link:
        # TO DO: give the user option to specify what type of link is used (as
        # they can for the entangling link). Eg, the user may wish to use a
        # free-space classical connection instead.
        create_classical_fibre_link(
            network, node_a, node_b, length=node_distance, label="classical"
        )

    if want_extra_classical_2way_link:
        # This is useful if you want to make scheduling easier by separating out
        # tasks. For example, one link for exchanging measurement results and one
        # for doing handshake protocols.
        # TO DO: give the user option to specify what type of link is used (as
        # they can for the entangling link). Eg, the user may wish to use a
        # free-space classical connection instead.
        create_classical_fibre_link(
            network, node_a, node_b, length=node_distance, label="extra_classical"
        )

    if want_entangling_link:
        create_entangling_link(network, node_a, node_b, **kwargs4create_entangling_link)


def create_dqc_network(
    *args4qproc,
    node_list=None,
    quantum_topology=None,
    classical_topology=None,
    want_classical_2way_link=True,
    want_extra_classical_2way_link=None,
    want_entangling_link=True,
    num_qpus=2,
    name="linear network",
    create_entangling_link=None,
    custom_qpu_func=None,
    # parameter for default classical AND entangling connection
    node_distance=2e-3,
    # parameters for default entangling connection
    state4distribution=None,  # ks.b00 defined in function body
    ent_dist_rate=182,  # Hz
    # parameters for default quantum processor
    num_comm_qubits=2,
    p_depolar_error_cnot=0,
    comm_qubit_depolar_rate=0,
    proc_qubit_depolar_rate=0,
    # for any non-default entangling connection or quantum processor
    **kwargs,
):
    """
    Creates a network suitable for distributed quantum computing

    Creates a network with QPUs node_0 to node_nn
    connected by entangling connections and/or classical connections.
    In other words this function sets up the physical components of the
    distributed quantum computer. The
    default topology of the network is a linear topology of the form
    node_1 <--> node_2 <--> node_3 <--> ...
    If not manually specified with the optional node_list argument, all nodes
    will be identical.

    Parameters
    ----------
    *args4qproc :
        The arguments for the custom_qpu_func specified. By default,
        there are none of these.
    node_list :  list of netsquid.components.nodes.Node objects, optional
        List of nodes to be included in network. Overrides num_qpus if
        specified.
    num_qpus : int (>=2), optional
        The number of identical nodes to be included in the network. Each node
        has 5 data qubits and two communication qubits. Must be at least 2 to
        avoid trivial scenario (if you have no connections there is no need)
        for this function, just create a node.
    Is overridden if a node_list is specified.
    node_distance : float, optional
        Distance between adjacent nodes.
    quantum_toplogy : list of tuple pairs of integers, optional
        The quantum network topology. Default is linear. The indices refer to
        the nodes.to be connected should be connected in the tuple.
    classical_topology: list of tuple pairs of integers, optional
            The classical network topology. Default is linear. The indices
            refer to the nodes.to be connected should be connected in the tuple.
    want_classical_2way_link : bool, optional
        Whether a two-way classical link should be created between all nodes
        (True) or not (False). Default is True. Unused if `classical_topology`
        or `quantum_topology` is not None.
    want_extra_classical_2way_link : bool or None, optional
        Whether an extra two-way classical link should be created between all
        nodes (True) or not (False). Default is None. If None, it will take
        the same value as `want_classical_2way_link`. Unlike,
        `want_classical_2way_link` this will be used if `classical_topology` is
        not None (in which case an extra classical link will be created
        everywhere a classical link is created).
    want_entangling_link : bool, optional
        Whether a two-way quantum link should be created between all nodes
        (True) or not (False). Nodes can request entangled pairs using this
        link. Unused if `classical_topology` or `quantum_topology` is not None.
    create_entangling_link : func, optional
        The function to use in order to create an entangling link. By default,
        :func: `create_black_box_central_source_entangling_link` is used.
    custom_qpu_func: function or None, optional
        Creates the quantum processor object to use on each node. It must
        be able to run with no positional arguments. If unspecified then
        a default noiseless processor will be used
    state4distribution : numpy.ndarray
        The entangled state distributed between nodes when
        an EPR pair is requested. Default is |phi^+> Bell state (formalism not
        fixed to ket). This argument is used by the default
        `create_entangling_link` function but may be ignored if another
         function is specified using `create_entangling_link`.
    ent_dist_rate : float, optional
        The rate of entanglement distribution [Hz]. Used by default value of
        `create_entangling_link`. May not be used, if a different
        `create_entangling_link` function is used. `ent_dist_rate` is used by
        the default `create_entangling_link` function and may be used be
        ignored if a different entangling link is used.
    num_comm_qubits : int, optional
        The number of communication qubit positions on the QPU. Default is 2.
        This argument is used by the default custom_qpu_func but may be ignored
        if a different function is specified using the `create_qpu_func`
        parameter and that function does not use this argument.
    p_depolar_error_cnot : float, optional
        The probability of a depolarisation error after each CNOT gate.
        The default value is 0.
        This argument is used by the default custom_qpu_func but may be ignored
        if a different function is specified using the `create_qpu_func`
        parameter and that function does not use this argument.
    comm_qubit_depolar_rate : float, optional
        The depolarisation rate for communication qubits.
        The default value is 0.
        This argument is used by the default custom_qpu_func but may be ignored
        if a different function is specified using the `create_qpu_func`
        parameter and that function does not use this argument.
    proc_qubit_depolar_rate : float
        The depolarisation rate for processing qubits.
        The default value is 0.
        This argument is used by the default custom_qpu_func but may be ignored
        if a different function is specified using the `create_qpu_func`
        parameter and that function does not use this argument.
    name : str, optional
        Name of the network
    **kwargs :
        The name-value arguments for the custom_qpu_func and
        create_entangling_link functions specified.

    Returns
    -------
    network : :class:`~netsquid.nodes.network.Network`
        The simulated hardware for a distributed quantum computer.
    """

    # initialising default arguments
    if want_extra_classical_2way_link is None:
        want_extra_classical_2way_link = want_classical_2way_link
    if state4distribution is None:
        state4distribution = ks.b00
    if custom_qpu_func is None:
        custom_qpu_func = (
            create_qproc_with_analytical_noise_ionQ_aria_durations_N_standard_lib_gates
        )
    if create_entangling_link is None:
        create_entangling_link = create_black_box_central_source_entangling_link
    # instantiating keyword arguments for default custom_qpu_func
    # function, this is necessary because I have explicitly included those
    # keyword arguments in the function definition above. This is done to
    # make the documentation clearer (ie, explicitly make the user aware of
    # the default options)
    kwargs["num_comm_qubits"] = num_comm_qubits
    kwargs["p_depolar_error_cnot"] = p_depolar_error_cnot
    kwargs["comm_qubit_depolar_rate"] = comm_qubit_depolar_rate
    kwargs["proc_qubit_depolar_rate"] = proc_qubit_depolar_rate
    # instantiating keyword arguments for default create_entangling_link
    # function, this is necessary because I have explicitly included those
    # keyword arguments in the function definition above. This is done to
    # make the documentation clearer (ie, explicitly make the user aware of
    # the default options)
    kwargs["state4distribution"] = state4distribution
    kwargs["node_distance"] = node_distance
    kwargs["ent_dist_rate"] = ent_dist_rate

    kwargs4qproc, kwargs4create_entangling_link = filter_kwargs4internal_functions(
        [custom_qpu_func, create_entangling_link], kwargs
    ).values()

    if num_qpus < 2:
        raise ValueError("num_qpus must be at least 2")
    if type(num_qpus) is not int:
        raise TypeError("num_qpus must be an integer")

    def _create_qpus(num_qpus2create=num_qpus, starting_index=0):
        qpu_list = []
        for ii in range(num_qpus2create):
            # need to instantiate a new processor object each time with a
            # new function call. Otherwise, the code is trying to add the
            # same processor object to different nodes, as opposed to a copy
            # of it. Therefore, it makes most sense to define this here within
            # create_dqc_network
            qpu = custom_qpu_func(*args4qproc, **kwargs4qproc)
            qpu.name = f"qproc4node_{ii + starting_index}"
            qpu_node = Node(f"node_{ii + starting_index}", qmemory=qpu)
            qpu_list = qpu_list + [qpu_node]  # appending node to node_list
        return qpu_list

    def _handle_even_num_qpus(network, num_qpus_considered):
        """
        Subroutine for creating linear network with even number of QPUs

        Parameters
        ----------
        network : TYPE
            DESCRIPTION.
        num_qpus_considered : TYPE
            DESCRIPTION.
        """
        for ii in range(num_qpus_considered - 1):
            node_1 = network.get_node(f"node_{ii}")
            node_2 = network.get_node(f"node_{ii + 1}")
            link_2_qpus(
                network,
                node_1,
                node_2,
                want_classical_2way_link=want_classical_2way_link,
                want_extra_classical_2way_link=want_extra_classical_2way_link,
                want_entangling_link=want_entangling_link,
                create_entangling_link=create_entangling_link,
                **kwargs4create_entangling_link,
            )

    network = Network(name)

    if node_list is None:
        node_list = _create_qpus(num_qpus2create=num_qpus)
    # NetSquid natively raises error if you try adding nodes with the same
    # name to a network, so it is safe to assume each node has a unique name in
    # what follows despite the possibility of mixing and matching auto-generated
    # and user-defined names
    network.add_nodes(node_list)
    if quantum_topology is None and classical_topology is None:
        # create network with linear topology
        if (num_qpus % 2) == 0:  # if num_qpus is even
            _handle_even_num_qpus(network, num_qpus)
        else:  # if num_qpus is odd
            # handle first num_qpus-1 nodes
            num_qpus_considered = num_qpus - 1
            _handle_even_num_qpus(network, num_qpus_considered)
            # handle the rest
            node_1 = network.get_node(f"node_{num_qpus - 2}")
            node_2 = network.get_node(f"node_{num_qpus - 1}")  # num_qpus-1 is
            # max index
            link_2_qpus(
                network,
                node_1,
                node_2,
                want_classical_2way_link=want_classical_2way_link,
                want_extra_classical_2way_link=want_entangling_link,
                want_entangling_link=want_entangling_link,
                create_entangling_link=create_entangling_link,
                **kwargs4create_entangling_link,
            )
    else:
        if quantum_topology is not None:
            if (quantum_topology[-1])[1] > (len(network.nodes) - 1):
                # if index is to large for number of nodes in network
                nodes_needed = (quantum_topology[-1])[1] - (len(network.nodes) - 1)
                extra_nodes = _create_qpus(
                    num_qpus2create=nodes_needed, starting_index=num_qpus
                )
                network.add_nodes(extra_nodes)
                warnings.warn(
                    "WARNING: num_qpus was overwritten because there"
                    " were not enough nodes to realise the requested"
                    " topology"
                )

            for ii in range(len(quantum_topology)):
                indices = quantum_topology[ii]
                node_a = network.get_node(f"node_{indices[0]}")
                node_b = network.get_node(f"node_{indices[1]}")
                link_2_qpus(
                    network,
                    node_a,
                    node_b,
                    want_classical_2way_link=False,
                    want_extra_classical_2way_link=False,
                    want_entangling_link=True,
                    create_entangling_link=create_entangling_link,
                    **kwargs4create_entangling_link,
                )
        if classical_topology is not None:  # if classical_topology is not None
            if (classical_topology[-1])[1] > (len(network.nodes) - 1):
                # if index is too large for number of QPUs in network
                nodes_needed = (classical_topology[-1])[1] - (len(network.nodes) - 1)
                extra_nodes = _create_qpus(
                    num_qpus2create=nodes_needed, starting_index=num_qpus
                )
                network.add_nodes(extra_nodes)
                warnings.warn(
                    "WARNING: num_qpus was overwritten because there"
                    " were not enough nodes to realise the requested"
                    " topology"
                )

            for ii in range(len(classical_topology)):
                indices = classical_topology[ii]
                node_a = network.get_node(f"node_{indices[0]}")
                node_b = network.get_node(f"node_{indices[1]}")
                link_2_qpus(
                    network,
                    node_a,
                    node_b,
                    want_classical_2way_link=True,
                    want_extra_classical_2way_link=want_extra_classical_2way_link,
                    want_entangling_link=False,
                    create_entangling_link=create_entangling_link,
                    **kwargs4create_entangling_link,
                )
    return network


class DQC(Network):
    """
    Encapsulates all of the hardware needed for distributed quantum computing.

    Creates `num_qpus` QPU nodes, with quantum memories defined by `qpu_class`,
    connected by quantum and classical Connections according to the classical
    and quantum topologies specified by `classical_topology` and
    `quantum_topology`, respectively.


    Parameters
    ----------
    entangling_connection_class : subclass of :class:`~netsquid.nodes.connections.Connection`
        The subclass to use for creating a connections between QPU nodes. Can
        be used to define a specific connection. Pre-made Connections can
        be found in `dqc_simulator.hardware.connections`.
    num_qpus : int
        The number of quantum processing units in the DQC.
    quantum_topology: list of (int, int)
        The topology of quantum connections between QPU nodes.
    classical_topology : list of (int, int)
        The topology of classical connections between QPU nodes. These are
        optical fibre connections.
    qpu_class : subclass of :class:`~dqc_simulator.hardware.quantum_processors.QPU` or None, optional
        The class to use when creating the QPUs. Allows a specific type of
        QPU to be specied. The default is None. If this is None, `nodes`
        must not be None and the QPUs used will be defined there.
    classical_connection_class : subclass of :class:`~netsquid.nodes.connections.Connection` or None, optional
        The classical connection to use. The default is None. If None, a
        classical fibre connection is used
    want_extra_classical_link : bool, optional
        Whether to create two classical connections of class
        `classical_connection_class` for every connection indicated in
        `classical_topology`. The default is True.
    name : str, optional
        The name of the DQC object. The default is "DQC".
    nodes : list of :class:`~netsquid.nodes.node.Node` instances or None, optional
        Overwritten if qpu_class is specified. If this is None, `nodes` must
        not be None.
    node_separation : float or None, optional
        The distance between QPU nodes in km. Used when instantiating the default
        `classical_connection_class`. Default is None. If None, the value of
        2e-3 (corresponding to 2m) is used.
    **kwargs
        The kwargs needed by `entangling_connection_class`, `qpu_class` and
        `classical_connection_class`. Note that if these classes have
        keyword arguments with the same name, an error will be raised.


    Returns
    -------
    network : :class:`~netsquid.nodes.network.Network`
        The simulated hardware for a distributed quantum computer.
    """

    def __init__(
        self,
        entangling_connection_class,
        num_qpus,
        quantum_topology,
        classical_topology,
        qpu_class=None,
        classical_connection_class=None,
        want_extra_classical_link=True,
        name="DQC",
        nodes=None,
        node_separation=None,
        **kwargs,
    ):
        super().__init__(name, nodes=nodes)
        # Instantiating mutable default values.
        if node_separation is None:
            self.node_separation = 2e-3  # 2m
        else:
            self.node_separation = node_separation

        # Splitting the kwargs by the class that they should be applied to
        class_tuple = (
            entangling_connection_class,
            qpu_class,
            classical_connection_class,
        )
        classes = [value for value in class_tuple if value is not None]
        kwargs4each_class = filter_kwargs4internal_functions(classes, kwargs)
        kwarg_types = {key: kwargs4each_class[key] for key in classes}

        self.qpu_class = qpu_class
        self.num_qpus = num_qpus

        if nodes is None and qpu_class is None:
            raise ValueError("nodes and qpu_class cannot both have value None")
        if qpu_class is not None:
            qpu_nodes = self._create_qpu_nodes(**kwarg_types[qpu_class])
            self.add_nodes(qpu_nodes)
        else:
            qpu_nodes = list(self.nodes.values())

        for pair in quantum_topology:
            entangling_connection = entangling_connection_class(
                **kwarg_types[entangling_connection_class]
            )
            label = "entangling"
            self._add_link(qpu_nodes, pair, entangling_connection, label)
        for pair in classical_topology:
            label = "classical"
            classical_connection = self._get_classical_connection(
                classical_connection_class, kwarg_types
            )
            self._add_link(qpu_nodes, pair, classical_connection, label)
            if want_extra_classical_link:
                label = "extra_classical"
                classical_connection = self._get_classical_connection(
                    classical_connection_class, kwarg_types
                )
                self._add_link(qpu_nodes, pair, classical_connection, label)

    def _create_qpu_nodes(self, **kwargs4qpu):
        qpu_list = []
        for ii in range(self.num_qpus):
            qpu = self.qpu_class(**kwargs4qpu)
            qpu.name = f"qproc4node_{ii}"
            qpu_node = Node(f"node_{ii}", qmemory=qpu)
            qpu_list = qpu_list + [qpu_node]  # appending node to node_list
        return qpu_list

    def _get_classical_connection(self, classical_connection_class, kwarg_types):
        if classical_connection_class is None:
            return DirectConnection(
                "DirectConnection",
                channel_AtoB=ClassicalChannel(
                    "Channel_A2B",
                    length=self.node_separation,
                    models={"delay_model": FibreDelayModel()},
                ),
                channel_BtoA=ClassicalChannel(
                    "Channel_B2A",
                    length=self.node_separation,
                    models={"delay_model": FibreDelayModel()},
                ),
            )
        else:
            return classical_connection_class(**kwarg_types[classical_connection_class])

    def _add_link(self, qpu_nodes, nodes2connect, connection, label):
        """


        Parameters
        ----------
        qpu_nodes : TYPE
            DESCRIPTION.
        qpus2connect : (int, int)
            The indices of the QPU nodes to connect
        connection : TYPE
            DESCRIPTION.
        label : TYPE
            DESCRIPTION.

        Returns
        -------
        None.

        """
        node_a = qpu_nodes[nodes2connect[0]]
        node_b = qpu_nodes[nodes2connect[1]]
        node_a_port_name = node_a.connection_port_name(node_b.name, label=label)
        node_b_port_name = node_b.connection_port_name(node_a.name, label=label)
        self.add_connection(
            node_a,
            node_b,
            connection=connection,
            label=label,
            port_name_node1=node_a_port_name,
            port_name_node2=node_b_port_name,
        )


# CHANGELOG:
# create_processor changed to have only optional arguments.
