#   Copyright 2018 National Technology & Engineering Solutions of Sandia,
#   LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS,
#   the U.S. Government retains certain rights in this software.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

"""
These functions build distance graphs for logical gates of qeccs.
"""

import networkx as nx


def precompute(instr):
    """

    Args:
        instr:

    Returns:

    """

    qecc = instr.qecc

    if qecc.name == '4.4.4.4 Surface Code' and qecc.circuit_compiler.name == 'Check2Circuits':
        precomputed_data = code_surface4444(instr)

    elif qecc.name == 'Medial 4.4.4.4 Surface Code' and qecc.circuit_compiler.name == 'Check2Circuits':
        precomputed_data = code_surface4444medial(instr)

    else:
        raise Exception('Can only handle the non-medial surface code!')

    return precomputed_data


def code_surface4444(instr):
    """
    Pre-computing for surface4444 class.

    This decoder is for 2D slices. It is assumed that it can decode logical instruction by logical instruction.

    :param logical_gate:
    :param precomputed_data: A dictionary of precomputed data used to decode syndromes of logical instructions.
    :return:
    """

    if instr.symbol == 'instr_syn_extract':

        # In the future go through different instructions
        decoder_data = surface4444_identity(instr)

    else:
        raise Exception('Can currently only handle "instr_init_zero".')

    return decoder_data


def code_surface4444medial(instr):
    """
    Pre-computing for surface4444 class.

    This decoder is for 2D slices. It is assumed that it can decode logical instruction by logical instruction.

    :param logical_gate:
    :param precomputed_data: A dictionary of precomputed data used to decode syndromes of logical instructions.
    :return:
    """

    if instr.symbol == 'instr_syn_extract':

        # In the future go through different instructions
        decoder_data = surface4444medial_identity(instr)

    else:
        raise Exception('Can currently only handle "instr_init_zero".')

    return decoder_data


def surface4444_identity(instr):
    """
    For X and Z decoding separately:

    - Create dictionary:

    - Determine how virtual nodes connect to data qubits.
    - find syndromes(and virtual) edges -> data
    - Generate distance graph
    - Determine syn -> closest v and weight

    :param instr:
    :return:
    """

    # In the end... need:
    # For x and z separately:
    #
    # syn -> closest virtual node
    # edge -> data
    # distance graph

    qecc = instr.qecc

    virtual_edge_data_x = {}
    virtual_edge_data_z = {}

    # Create a dictionary to store precomputed information that will be used for decoding
    info = {
        'X': {
            'dist_graph': nx.Graph(),
            'closest_virt': {},
            'virtual_edge_data': virtual_edge_data_x,
        },
        'Z': {
            'dist_graph': nx.Graph(),
            'closest_virt': {},
            'virtual_edge_data': virtual_edge_data_z,
        },
    }

    # Record what data qudits the syndrome to syndrome edges correspond to.
    # (syn, syn/virt): data
    edges_x = {}
    edges_z = {}

    # syndrome-to-syndrome, fully-connected graph
    # (syn, syn): weight
    graph_x = info['X']['dist_graph']
    graph_z = info['Z']['dist_graph']

    # The closest virtual node to each syndrome
    # syn: virt
    closest_x = info['X']['closest_virt']
    closest_z = info['Z']['closest_virt']

    # The sides of the QECC patch
    sides = qecc.sides  # t, r, b, l

    # checks of the logical instruction
    instr = qecc.instruction('instr_syn_extract')
    abs_circ = instr.abstract_circuit

    # Dictionary of data qudit to syndrome-to-syndrome edge
    d2edge_x = {}
    d2edge_z = {}

    # Temporary graphs that will store the direct syndrome-to-syndrome edges. This will be used to create the fully
    # connected, distance graph.
    temp_graph_x = nx.Graph()
    temp_graph_z = nx.Graph()

    # Assume the QECC uses checks
    # add edges based on checks
    for gate_symbol, _, params in abs_circ.items(params=True):

        ancilla = params['ancillas']
        data_qubits = params['datas']

        for data in data_qubits:

            if gate_symbol == 'X check':
                edges = d2edge_x
            elif gate_symbol == 'Z check':
                edges = d2edge_z
            else:
                raise Exception('This decoder can only handle check of purely X or Z type rather than %s!'
                                % gate_symbol)

            syn_list = edges.setdefault(data, [])
            syn_list.append(ancilla)

    # ---- Create virtual nodes ----- #

    virt_x = set([])
    virt_z = set([])

    # side: top, right, bottom, left
    # For the non-medial surface code patch, a x virtual node is paired with data qubits on the right and left sides
    # and a z virtual node is paired with data qubits on the top and bottom sides (should say what Pauli type the side
    # is.

    vi = 0

    for side_label, side_qubits in sides.items():

        for data in side_qubits:

            vi += 1

            virt_node = 'v' + str(vi)

            # X virtual nodes (sides left and right)
            if side_label in ['left', 'right']:  # 1 for i = 1 and 3 => left and right

                syn_list = d2edge_x.setdefault(data, [])
                syn_list.append(virt_node)
                virt_x.add(virt_node)

            # Z virtual nodes (sides top and bottom)
            elif side_label in ['top', 'bottom']:  # 0 for i = 0 and 2 => top and bottom

                syn_list = d2edge_z.setdefault(data, [])
                syn_list.append(virt_node)
                virt_z.add(virt_node)
            else:
                raise Exception('side_label "%s" not understood!' % side_label)

    # invert data -> edge and make sure len(edge) = 2
    for check_type in ['X', 'Z']:

        if check_type == 'X':
            edge_dict = d2edge_x
            edges = edges_x
            temp_graph = temp_graph_x
        else:
            edge_dict = d2edge_z
            edges = edges_z
            temp_graph = temp_graph_z

        for data, edge in edge_dict.items():

            if len(edge) != 2:
                raise Exception('There should be exactly two syndromes (virtual or not) connected to each data '
                                'qudit. Instead, q: %s edge: %s' % (data, edge))

            edges[tuple(edge)] = data
            edges[(edge[1], edge[0])] = data
            temp_graph.add_edge(edge[0], edge[1])

    # Create distance graph
    for check_type in ['X', 'Z']:

        if check_type == 'X':

            temp_graph = temp_graph_x
            G = graph_x
            closest = closest_x
            virt = virt_x
            edge2d = edges_x
            virtual_edge_data = virtual_edge_data_x

        else:
            temp_graph = temp_graph_z
            G = graph_z
            closest = closest_z
            virt = virt_z
            edge2d = edges_z
            virtual_edge_data = virtual_edge_data_z

        pathes = nx.shortest_path(temp_graph)

        for n1, wdict in pathes.items():
            for n2, syn_path in wdict.items():

                weight = len(syn_path) - 1

                if weight != 0:

                    # print 'n1', n1, 'n2', n2, 'AP', syn_path, 'weight', weight

                    # Get list of datas corresponding to the connected path between syndromes
                    data_path = []
                    s1 = syn_path[0]
                    for s2 in syn_path[1:]:
                        # print 's1', s1, 's2', s2
                        # print 'data', edge2d[(s1, s2)]
                        data = edge2d[(s1, s2)]
                        data_path.append(data)
                        s1 = s2

                    # G.add_edge(n1, n2, weight=1.0/float(weight), syn_path=syn_path, data_path=data_path)
                    # G.add_edge(n1, n2, weight=float(weight), syn_path=syn_path, data_path=data_path)
                    if (n1 not in virt) and (n2 not in virt):
                        G.add_edge(n1, n2, weight=-weight, syn_path=syn_path, data_path=data_path)
                    # elif (n1 not in virt) and (n2 in virt):
                    #    virtual_edge_data[n2] = {'weight': -weight, 'syn_path': syn_path, 'data_path': data_path}

        syn = set(G.nodes())
        syn -= virt

        # Find closest virtual node
        # print 'syn', syn
        # print 'virt', virt

        for s in syn:
            shortest_len = float('inf')
            closest_v = None
            for v in virt:
                sv_len = len(pathes[s][v])
                # print '%s -> %s: %s' % (s, v, sv_len)
                if sv_len < shortest_len:
                    shortest_len = sv_len
                    closest_v = v
            closest[s] = closest_v
            # print 'C', s, closest_v, shortest_len
            # print 'closest', closest_v

        for s, v in closest.items():
            # print s, '->', v

            syn_path = pathes[s][v]
            weight = len(syn_path) - 1

            data_path = []
            s1 = syn_path[0]
            for s2 in syn_path[1:]:
                # print 's1', s1, 's2', s2
                # print 'data', edge2d[(s1, s2)]
                data = edge2d[(s1, s2)]
                data_path.append(data)
                s1 = s2

            virtual_edge_data[s] = {'virtual_node': v, 'weight': -weight, 'syn_path': syn_path, 'data_path': data_path}
            # print s, '>>', virtual_edge_data[s]

    # print 'closest X', closest_x
    # print 'closest Z', closest_z

    return info


def surface4444medial_identity(instr):
    """
    For X and Z decoding separately:

    - Create dictionary:

    - Determine how virtual nodes connect to data qubits.
    - find syndromes(and virtual) edges -> data
    - Generate distance graph
    - Determine syn -> closest v and weight

    :param instr:
    :return:
    """

    # In the end... need:
    # For x and z separately:
    #
    # syn -> closest virtual node
    # edge -> data
    # distance graph

    qecc = instr.qecc

    virtual_edge_data_x = {}
    virtual_edge_data_z = {}

    # Create a dictionary to store precomputed information that will be used for decoding
    info = {
        'X': {
            'dist_graph': nx.Graph(),
            'closest_virt': {},
            'virtual_edge_data': virtual_edge_data_x,
        },
        'Z': {
            'dist_graph': nx.Graph(),
            'closest_virt': {},
            'virtual_edge_data': virtual_edge_data_z,
        },
    }

    # Record what data qudits the syndrome to syndrome edges correspond to.
    # (syn, syn/virt): data
    edges_x = {}
    edges_z = {}

    # syndrome-to-syndrome, fully-connected graph
    # (syn, syn): weight
    graph_x = info['X']['dist_graph']
    graph_z = info['Z']['dist_graph']

    # The closest virtual node to each syndrome
    # syn: virt
    closest_x = info['X']['closest_virt']
    closest_z = info['Z']['closest_virt']

    # The sides of the QECC patch
    sides = qecc.sides  # t, r, b, l

    # checks of the logical instruction
    instr = qecc.instruction('instr_syn_extract')
    abs_circ = instr.abstract_circuit

    # Dictionary of data qudit to syndrome-to-syndrome edge
    d2edge_x = {}
    d2edge_z = {}

    # Temporary graphs that will store the direct syndrome-to-syndrome edges. This will be used to create the fully
    # connected, distance graph.
    temp_graph_x = nx.Graph()
    temp_graph_z = nx.Graph()

    # Assume the QECC uses checks
    # add edges based on checks
    for gate_symbol, _, params in abs_circ.items(params=True):

        data_qudits = params['datas']
        ancilla = params['ancillas']

        for data in data_qudits:

            if gate_symbol == 'X check':
                edges = d2edge_x
            elif gate_symbol == 'Z check':
                edges = d2edge_z
            else:
                raise Exception('This decoder can only handle check of purely X or Z type rather than %s!' %
                                gate_symbol)

            syn_list = edges.setdefault(data, [])
            syn_list.append(ancilla)

    # print('x ed', d2edge_x)
    # print('z ed', d2edge_z)

    # VIRTUAL NODES...
    # ---- Create virtual nodes ----- #
    virt_x = set([])
    virt_z = set([])

    # side: top, right, bottom, left
    # For the non-medial surface code patch, a x virtual node is paired with data qubits on the right and left sides
    # and a z virtual node is paired with data qubits on the top and bottom sides (should say what Pauli type the side
    # is.

    distance_width = qecc.width
    distance_height = qecc.height

    vi = 0
    # print('sides', sides)
    for side_label, side_qubits in sides.items():

        # print('side qudbits: %s' % side_qubits)

        for i, data in enumerate(side_qubits):

            if side_label == 'top':

                if distance_width % 2 == 1:  # odd

                    if i == 0 or i % 2 == 1:
                        vi += 1
                        virt_node = 'v' + str(vi)

                else:  # even

                    if i % 2 == 0:
                        vi += 1
                        virt_node = 'v' + str(vi)

                syn_list = d2edge_z.setdefault(data, [])
                syn_list.append(virt_node)
                virt_z.add(virt_node)

            elif side_label == 'bottom':

                if i == 0 or i % 2 == 1:
                    vi += 1
                    virt_node = 'v' + str(vi)

                syn_list = d2edge_z.setdefault(data, [])
                syn_list.append(virt_node)
                virt_z.add(virt_node)

            elif side_label == 'left':

                if i == 0 or i % 2 == 1:
                    vi += 1
                    virt_node = 'v' + str(vi)

                syn_list = d2edge_x.setdefault(data, [])
                syn_list.append(virt_node)
                virt_x.add(virt_node)

            elif side_label == 'right':

                if distance_height % 2 == 1:
                    if i == 0 or i % 2 == 1:
                        vi += 1
                        virt_node = 'v' + str(vi)
                else:

                    if i % 2 == 0:
                        vi += 1
                        virt_node = 'v' + str(vi)

                syn_list = d2edge_x.setdefault(data, [])
                syn_list.append(virt_node)
                virt_x.add(virt_node)
            else:
                raise Exception('side_label "%s" not understood!' % side_label)

    # print('vx ed', d2edge_x)
    # print('vz ed', d2edge_z)

    # invert data -> edge and make sure len(edge) = 2
    for check_type in ['X', 'Z']:

        if check_type == 'X':
            edge_dict = d2edge_x
            edges = edges_x
            temp_graph = temp_graph_x
        else:
            edge_dict = d2edge_z
            edges = edges_z
            temp_graph = temp_graph_z

        for data, edge in edge_dict.items():

            if len(edge) != 2:
                raise Exception('There should be exactly two syndromes (virtual or not) connected to each data '
                                'qudit. Instead, q: %s edge: %s' % (data, edge))

            edges[tuple(edge)] = data
            edges[(edge[1], edge[0])] = data
            temp_graph.add_edge(edge[0], edge[1])

    # Create distance graph
    for check_type in ['X', 'Z']:

        if check_type == 'X':

            temp_graph = temp_graph_x
            G = graph_x
            closest = closest_x
            virt = virt_x
            edge2d = edges_x
            virtual_edge_data = virtual_edge_data_x

        else:
            temp_graph = temp_graph_z
            G = graph_z
            closest = closest_z
            virt = virt_z
            edge2d = edges_z
            virtual_edge_data = virtual_edge_data_z

        pathes = nx.shortest_path(temp_graph)

        for n1, wdict in pathes.items():
            for n2, syn_path in wdict.items():

                weight = len(syn_path) - 1

                if weight != 0:

                    # Get list of datas corresponding to the connected path between syndromes
                    data_path = []
                    s1 = syn_path[0]
                    for s2 in syn_path[1:]:
                        data = edge2d[(s1, s2)]
                        data_path.append(data)
                        s1 = s2

                    if (n1 not in virt) and (n2 not in virt):
                        G.add_edge(n1, n2, weight=-weight, syn_path=syn_path, data_path=data_path)

        syn = set(G.nodes())
        syn -= virt

        # Find closest virtual node

        for s in syn:
            shortest_len = float('inf')
            closest_v = None
            for v in virt:
                sv_len = len(pathes[s][v])
                # print '%s -> %s: %s' % (s, v, sv_len)
                if sv_len < shortest_len:
                    shortest_len = sv_len
                    closest_v = v
            closest[s] = closest_v

        for s, v in closest.items():

            syn_path = pathes[s][v]
            weight = len(syn_path) - 1

            data_path = []
            s1 = syn_path[0]
            for s2 in syn_path[1:]:
                data = edge2d[(s1, s2)]
                data_path.append(data)
                s1 = s2

            virtual_edge_data[s] = {'virtual_node': v, 'weight': -weight, 'syn_path': syn_path, 'data_path': data_path}

    return info
