import numpy as np
import qtda_module as qtda
from qiskit.visualization import plot_histogram
import matplotlib.pyplot as plt
from scipy.spatial import distance_matrix
from qiskit import ClassicalRegister, Aer, execute

def generate_circles():
    phi_values = np.linspace(0, 2*np.pi, 10)
    small_circle = [[np.cos(phi), np.sin(phi)] for phi in phi_values]
    phi_values = np.linspace(0, 2*np.pi, 12)
    large_circle = [[3 + 2*np.cos(phi), 2*np.sin(phi)] for phi in phi_values]
    
    return np.array(small_circle + large_circle)
    
point_data = generate_circles()
plt.scatter(point_data[:,0], point_data[:,1])

# alternatively a distance matrix (here generated from the point data for exemplification)
# can be used
dist_mat = distance_matrix(point_data, point_data)

filtration = qtda.DataFiltration(
    data=point_data, 
#     distance_matrix=dist_mat, 
    max_dimension=4, 
    max_edge_length=10
    )

filtration.plot_persistence_diagram()

n_vertices = 4 # number of qubits to represent the simplicial complex
num_eval_qubits = 5 # number of numerical evaluation qubits for the QPE algorithm

S0 = [(0,0,0,1),(0,0,1,0), (0,1,0,0),(1,0,0,0)] # 0-simplex: points in the graph
S1 = [(0,0,1,1),(0,1,1,0),(1,1,0,0),(1,0,0,1),(1,0,1,0)] # 1-simplex: connection lines between points
S2 = [(1,0,1,1)] # 2-simplex: filled triangle
S3 = []

state_dict = {0: S0, 1: S1, 2: S2, 3: S3}
k = 1 # order of the combinatorial Laplacian
qc = qtda.QTDAalgorithm(num_eval_qubits, k, state_dict)
qc.draw('mpl')

qc.add_register(ClassicalRegister(num_eval_qubits)) #add classical register to measure evaluation qubits
for q in qc.eval_qubits:
    qc.measure(q,q)

shots = 1000
backend = Aer.get_backend('qasm_simulator')
job = execute(qc, backend, shots=shots)
counts = job.result().get_counts(qc)

dim = len(state_dict[1]) # dimension of 1-simplex space
prob = counts.get("0"*num_eval_qubits)/shots # probability of eigenvalue 0

print('The number of 1-holes is', dim * prob)
plot_histogram(counts)

shots = 1000
num_eval_qubits = 10

data = qtda.Q_top_spectra(state_dict, num_eval_qubits, shots)
spectra = data.get_spectra()

for top_order in spectra.keys():
    print()
    print('Topological order: ', top_order)
    print('Number of holes: ', spectra[top_order][0.0])
    print('Dimension of the k-simplex subspace: ', len(data.state_dict[top_order]))
    print('Eigenvalues of the combinatorial laplacian with dimension of corresponding eigenspaces: ')
    print(spectra[top_order])
    display(plot_histogram(data.get_counts()[top_order]))

n_vertices = 4

S0 = [(0,0,0,1),(0,0,1,0), (0,1,0,0),(1,0,0,0)]
S1 = [(0,0,1,1),(0,1,1,0),(1,1,0,0),(1,0,0,1),(1,0,1,0)]
S2 = []
S3 = []

state_dict = {0: S0, 1: S1, 2: S2, 3: S3}

shots = 1000
num_eval_qubits = 10

data = qtda.Q_top_spectra(state_dict, num_eval_qubits, shots)
spectra = data.get_spectra()

for top_order in spectra.keys():
    print()
    print('Topological order: ', top_order)
    print('Number of holes: ', spectra[top_order][0.0])
    print('Dimension of the k-simplex subspace: ', len(data.state_dict[top_order]))
    print('Eigenvalues of the combinatorial laplacian with dimension of corresponding eigenspaces: ')
    print(spectra[top_order])
    display(plot_histogram(data.get_counts()[top_order]))

import qiskit
qiskit.__qiskit_version__

