from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, execute, Aer, IBMQ
from qiskit.circuit.library import PhaseGate
from qiskit.extensions import HamiltonianGate
from qiskit.quantum_info import random_hermitian
import qiskit_tools as qt
import numpy as np
import matplotlib.pyplot as plt
import matplotlib

matplotlib.rcParams['mathtext.fontset'] = 'stix'
matplotlib.rcParams['font.family'] = 'STIXGeneral'
width=0.75
color='black'
fontsize=28
ticksize=22
figsize=(10,8)

nx = 1

A = random_hermitian(2**nx).to_matrix()
print(A)

p_suc = 0.95

phase = True

n = nx + int(np.ceil(np.log2(2. + 1./(2.*(1-p_suc)))))

print('The number of output qubits is set to',n)

lambda_min = np.real(np.min(np.linalg.eigvals(A)))
lambda_max = np.real(np.max(np.linalg.eigvals(A)))
kappa = np.abs(lambda_max/lambda_min)

print('Minimum eigenvalue, maximum eigenvalue and conditional number of A:', lambda_min, lambda_max, kappa)

# Calculate the new scale of the relative magnitudes of maximum and minimum eigenvalues
rescale = int((2**(n-phase)-1)/kappa)
# Make sure this scaling is not vanishingly small
if rescale < 1.e-7:
    rescale = 1
    
# determine the number of integer qubits required
nint = qt.get_nint(rescale)
# Assert that the number of integer qubits is equal to or greater than the number of ancillary qubits assigned
if nint < n - phase:
    nint = n - phase

t0 = qt.bin_to_dec(qt.my_binary_repr(rescale, nint=nint, n=nint, phase=False), nint=0, phase=False) / (lambda_min * (2**phase))
t = 2 * np.pi * t0

qreg = QuantumRegister(nx, 'x')
qout = QuantumRegister(n, 'out')
circ = QuantumCircuit(qreg, qout)

A_gate = HamiltonianGate(A, t)

circ.h(qreg[0]);
circ = qt.PhaseEst(circ, qreg, qout, A_gate)

circ.draw('latex')

backend = Aer.get_backend('statevector_simulator')
job = execute(circ, backend)
result = job.result()
state_vector = np.asarray(result.get_statevector())

state_v = np.asarray(state_vector).reshape((2**n,2**nx)).T
state_v = np.sum(state_v, axis=0)
state_v = state_v/np.sqrt(np.sum(np.abs(state_v)**2))

xs = np.fft.fftfreq(2**n,2.**(-n))/(t0*2**(n))

fig = plt.figure(figsize=np.array(figsize)*[1,2])
ax = fig.add_subplot(211)

plt.scatter(xs, state_v.real, marker='.', color='black')
plt.scatter(xs, state_v.imag, marker='.', color='grey')
ax.set_ylabel('Amplitude', fontsize=fontsize);
ax.tick_params(axis='both', labelsize=ticksize);

ax.axvline(lambda_min, ls='--', color='black')
ax.axvline(lambda_max, ls=':', color='black')

ax = fig.add_subplot(212)

plt.scatter(xs, np.abs(state_v)**2, marker='.', color='black')
ax.set_xlabel(r'$\theta$', fontsize=fontsize);
ax.set_ylabel('Probability', fontsize=fontsize);
ax.tick_params(axis='both', labelsize=ticksize);

ax.axvline(lambda_min, ls='--', color='black')
ax.axvline(lambda_max, ls=':', color='black')

plt.show()