# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.3.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import numpy as np
import copy
from qiskit import QuantumCircuit, QuantumRegister, Aer, execute
np.set_printoptions(linewidth=10000, threshold=10000)


# +
def theoretical_prob(initial, step, ptran, nq):
        Pi_op = Pi_operator(ptran)
        swap = swap_operator(nq)
        operator = (2*Pi_op) - np.identity(len(Pi_op))
        Szegedy = np.dot(operator, swap)
        Szegedy_n = copy.copy(Szegedy)
        print(id(Szegedy), id(Szegedy_n))
        if step == 0:
            init_prob = np.array([abs(i)**2 for i in initial], dtype=np.float)
            return init_prob
        elif step == 1:
            prob = np.array([abs(i)**2 for i in np.dot(Szegedy, initial)],
                            dtype=np.float)
            return prob
        else:
            for n in range(step-1):
                Szegedy_n = np.dot(Szegedy_n, Szegedy)
            probs = np.array([abs(i)**2 for i in np.dot(initial, Szegedy_n)],
                             dtype=np.float)
            return probs

def swap_operator(n_qubit):
    q1 = QuantumRegister(n_qubit//2)
    q2 = QuantumRegister(n_qubit//2)
    qc = QuantumCircuit(q1, q2)
    for c, t in zip(q1, q2):
        qc.swap(c, t)
    # FIXME
    backend = Aer.get_backend('unitary_simulator')
    job = execute(qc, backend=backend)
    swap = job.result().get_unitary(qc)
    return swap

def Pi_operator(ptran):
    '''
    This is not a quantum operation,
    just returning matrix
    '''
    lg = len(ptran)
    psi_op = []
    count = 0
    for i in range(lg):
        psi_vec = [0 for _ in range(lg**2)]
        for j in range(lg):
            psi_vec[count] = np.sqrt(ptran[j][i])
            count += 1
        psi_op.append(np.kron(np.array(psi_vec).T,
                      np.conjugate(psi_vec)).reshape((lg**2, lg**2)))
    Pi = psi_op[0]
    for i in psi_op[1:]:
        Pi = np.add(Pi, i)
    return Pi

def is_unitary(operator, tolerance=0.0001):
    h, w = operator.shape
    if not h == w:
        return False
    adjoint = np.conjugate(operator.transpose())
    product1 = np.dot(operator, adjoint)
    product2 = np.dot(adjoint, operator)
    ida = np.eye(h)
    return np.allclose(product1, ida) & np.allclose(product2, ida)



# +
alpha = 0.85
target_graph = np.array([[0, 1, 0, 0, 0, 0, 0, 1],
                                [0, 0, 1, 0, 0, 0, 0, 0],
                                [0, 1, 0, 1, 0, 0, 0, 0],
                                [0, 0, 1, 0, 1, 0, 0, 0],
                                [0, 0, 0, 1, 0, 1, 0, 0],
                                [0, 0, 0, 0, 1, 0, 1, 0],
                                [0, 0, 0, 0, 0, 1, 0, 1],
                                [0, 0, 0, 0, 0, 0, 1, 0]])

E = np.array([[1/8, 1/2, 0, 0, 0, 0, 0, 1/2],
                [1/8, 0, 1/2, 0, 0, 0, 0, 0],
                [1/8, 1/2, 0, 1/2, 0, 0, 0, 0],
                [1/8, 0, 1/2, 0, 1/2, 0, 0, 0],
                [1/8, 0, 0, 1/2, 0, 1/2, 0, 0],
                [1/8, 0, 0, 0, 1/2, 0, 1/2, 0],
                [1/8, 0, 0, 0, 0, 1/2, 0, 1/2],
                [1/8, 0, 0, 0, 0, 0, 1/2, 0]])
# use google matrix
lt = len(target_graph)
prob_dist = alpha*E + ((1-alpha)/lt)*np.ones((lt, lt))
init_state_eight = 1/2*np.array([np.sqrt(prob_dist[j][i]) for i in range(lt) for j in range(lt)])
# -

Pi_op = Pi_operator(prob_dist)
operator = (2*Pi_op) - np.identity(len(Pi_op))

lo = len(operator)
initial = np.array([0 if i != 32 else 1 for i in range(lo)])
# initial = np.array([np.sqrt(1/lo) for i in range(lo)])

ideal = theoretical_prob(initial, 2, prob_dist, 6)
for i, v in enumerate(ideal):
    print(i, v)

import networkx as nx
import matplotlib.pyplot as plt

target_graph = np.array([[0, 1, 0, 1],
                                [0, 0, 1, 0],
                                [0, 1, 0, 1],
                                [0, 0, 1, 0]])
Gs = nx.from_numpy_matrix(target_graph)
G = nx.MultiDiGraph()
pos = [(1, 1), (5, 1), (5, 5), (1, 5)]
G.add_edges_from(Gs.edges())
plt.figure(figsize=(8,8))
nx.draw_networkx(G, pos, node_size=1000, width=3, arrowsize=50)
plt.show()

# +
target_graph = np.array([[0, 1, 0, 0, 0, 0, 0, 1],
                                [0, 0, 1, 0, 0, 0, 0, 0],
                                [0, 1, 0, 1, 0, 0, 0, 0],
                                [0, 0, 1, 0, 1, 0, 0, 0],
                                [0, 0, 0, 1, 0, 1, 0, 0],
                                [0, 0, 0, 0, 1, 0, 1, 0],
                                [0, 0, 0, 0, 0, 1, 0, 1],
                                [0, 0, 0, 0, 0, 0, 1, 0]])

Gs = nx.from_numpy_matrix(target_graph)
G = nx.MultiDiGraph()
# pos = [(1, 1), (2, 2), (3, 3), (1, 5), (2, 2), (6, 2), (6, 6), (2, 6)]
pos = [(0, 0), (1, 1), (2, 2), (1, 3), (0, 4), (-1, 3), (-2, 2), (-1, 1)]
G.add_edges_from(Gs.edges())
# pos = nx.spring_layout(G)
plt.figure(figsize=(8,8))
nx.draw_networkx(G, pos, node_size=1000, width=3, arrowsize=50)
plt.show()

# +
target_graph = np.array([[0, 0, 0, 1, 1, 0, 0, 0],
                                [1, 0, 0, 0, 1, 0, 0, 0],
                                [0, 1, 0, 0, 0, 1, 0, 0],
                               [0, 0, 1, 0, 0, 1, 0, 0],
                               [0, 0, 0, 0, 0, 0, 1, 1],
                               [0, 0, 0, 0, 0, 0, 1, 1],
                               [0, 0, 0, 0, 0, 0, 1, 1],
                               [0, 0, 0, 0, 0, 0, 1, 1]])

Gs = nx.from_numpy_matrix(target_graph)
print(Gs.edges())
G = nx.MultiDiGraph()
# pos = [(1, 1), (2, 2), (3, 3), (1, 5), (2, 2), (6, 2), (6, 6), (2, 6)]
pos = [(0, 0), (3, 3), (2, 9), (-1, 5), (5, 1), (4, 5), (8, 5), (9, 3)]
G.add_edges_from(Gs.edges())
G.add_edges_from([(7, 6)])
# pos = nx.spring_layout(G, k=10)
plt.figure(figsize=(10, 10))
nx.draw_networkx(G, pos, node_size=1000, width=3, arrowsize=50)
plt.show()
# -


