import numpy as np
import random
import copy
import sys
import networkx as nx
import matplotlib.pyplot as plt

from numpy import pi
from tqdm import trange
from grakel import datasets, Graph
from sklearn.model_selection import KFold

sys.path.append("../")
from qwgc.QWGC import QWGC

# MUTAGを取ってくる
Data = datasets.fetch_dataset('MUTAG', verbose=False)
data_x, data_y = np.array(Data.data), np.array(Data.target)

# visualization of data
subtract = 0
lens = []
for d, l in zip(data_x[:10], data_y[:10]):
    print(l)
    plt.figure(figsize=(10, 10))
    G = nx.DiGraph()
    connection = d[0]
    nodesize = [(i+1)**800 for i in d[1].values()]
    edge_weight = d[2]
    lens.append(len([i for i in d[1].values()]))
    adjacency = Graph(connection).get_adjacency_matrix()
    nodes = np.array([str(i+1) for i, _ in enumerate(adjacency)])
    edges = []
    weight = []
    for i, v in edge_weight.items():
        ed = [str(st-subtract) for st in list(i)]
        ed.append(v+1)
        edges.append(tuple(ed))
    subtract = max(d[1].keys())
    G.add_nodes_from(nodes)
    G.add_weighted_edges_from(edges)
    pos = nx.kamada_kawai_layout(G)
    edge_labels = {(i, j): w['weight'] for i, j, w in G.edges(data=True)}
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
    nx.draw_networkx(G, pos, with_labels=True, alpha=0.8)
    plt.show()

from qwgc.preprocess.qwalk import QuantumWalk
# 対象のデータ
data = data_x[0]
label = data_y[0]
# 隣接行列
adjacency = Graph(data[0]).get_adjacency_matrix()
count = np.count_nonzero(adjacency)//2
# 量子ウォークのハイパーパラメータ
step = 1
# 次数が2の場合のコインのパラメータ(今回は簡単のために全てアダマールコイン)
coin = np.kron(np.identity(count), 1/np.sqrt(2)*np.array([[1, 1], [1, -1]]))
# 初期状態 (0からスタート)
initial = None 
# 量子ウォーカーが測定される確率
qwalk = QuantumWalk(initial, coin, adjacency)
qwalk.n_steps(step)
probs = qwalk.calc_probs()

# 描画
plt.figure(figsize=(10, 10))
G = nx.DiGraph()
# ノードの大きさで確率を表す
connection = data[0]
nodesize = [(i+0.1)*800 for i in probs]
edge_weight = data[2]
nodes = np.array([str(i+1) for i, _ in enumerate(adjacency)])
edges = []
weight = []
for i, v in edge_weight.items():
    ed = [str(st) for st in list(i)]
    ed.append(v+1)
    edges.append(tuple(ed))
subtract = max(data[1].keys())
G.add_nodes_from(nodes)
G.add_weighted_edges_from(edges)
pos = nx.spring_layout(G)
edge_labels = {(i, j): w['weight'] for i, j, w in G.edges(data=True)}
nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
nx.draw_networkx(G, pos, with_labels=True, alpha=0.8, node_size=nodesize)
plt.show()

def draw(nodesize):
    plt.figure(figsize=(10, 10))
    G = nx.DiGraph()
    # ノードの大きさで確率を表す
    connection = data[0]
    nodesize = [(i)*800 for i in probs]
    edge_weight = data[2]
    nodes = np.array([str(i+1) for i, _ in enumerate(adjacency)])
    edges = []
    weight = []
    for i, v in edge_weight.items():
        ed = [str(st) for st in list(i)]
        ed.append(v+1)
        edges.append(tuple(ed))
    subtract = max(data[1].keys())
    G.add_nodes_from(nodes)
    G.add_weighted_edges_from(edges)
    pos = nx.spring_layout(G)
    edge_labels = {(i, j): w['weight'] for i, j, w in G.edges(data=True)}
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
    nx.draw_networkx(G, pos, with_labels=True, alpha=0.8, node_size=nodesize)
    plt.show()

data = data_x[0]
label = data_y[0]
# 隣接行列
adjacency = Graph(data[0]).get_adjacency_matrix()
count = np.count_nonzero(adjacency)//2
# 量子ウォークのハイパーパラメータ
for step in range(1, 10):
    # 次数が2の場合のコインのパラメータ(今回は簡単のために全てアダマールコイン)
    coin = np.kron(np.identity(count), 1/np.sqrt(2)*np.array([[1, 1], [1, -1]]))
    # 初期状態 (0からスタート)
    initial = None 
    # 量子ウォーカーが測定される確率
    qwalk = QuantumWalk(initial, coin, adjacency)
    qwalk.n_steps(step)
    probs = qwalk.calc_probs()
    draw(probs)

from qwgc.preprocess.qwfilter import QWfilter
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit import Aer, execute, transpile
data = data_x[0]
label = data_y[0]
# filter
u3param = [pi, 0, pi/2]
# step
step = 3
# initial 重ね合わせを初期状態として利用
initial = "super"
qwfilter = QWfilter(u3param, step, initial)
# 今回は測定をせずに、振幅をそのまま量子回路にマッピングを行います。
amplitude = qwfilter.single_amplitude(data)

la = len(amplitude)
new_amplitude = list(amplitude) + [0 for i in range(64 - 38)]
print(len(new_amplitude))

nq = 6
#　量子レジスタの定義
qr = QuantumRegister(nq, name="quantum walk")
ancilla = QuantumRegister(2, name="ancilla")
# 古典レジスタの定義
cr = ClassicalRegister(2)
# 量子回路の定義
qc = QuantumCircuit(qr, ancilla, cr)
qc.draw(output='mpl')

qc.initialize(new_amplitude, qr)

# 回転角を初期化
theta = [np.random.uniform(-pi, pi) for i in range(nq)]
for ith, th in enumerate(theta):
    qc.cry(th, qr[ith], ancilla[ith%2])
qc.draw(output="mpl")

qc.measure(ancilla, cr)

backend = Aer.get_backend("qasm_simulator")
shots = 1024
job = execute(qc, backend=backend, shots=shots)
counts = job.result().get_counts(qc)
dinom = counts.get('01', 0) + counts.get('10', 0) + 1e-10
print("クラス-1である確率:", counts.get('01', 0)/dinom, "クラス1である確率:", counts.get('10', 0)/dinom)
if counts.get('01', 0)/dinom > counts.get('10', 0)/dinom:
    answer = -1
else:
    answer = 1
print("このグラフはクラス ", answer, "です.")

print("正解は", data_y[1], "です。")