from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, execute, Aer, IBMQ
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)

digit = 1.

bounds = np.array([-2., 0., 1.5, 3.])
print('Domain boundaries:',bounds)

m = int(np.ceil(np.log2(len(bounds))))
print('Label qubits:',m)

phase = True

nint = qt.get_nint([digit,*bounds])
npres = qt.get_npres([digit,*bounds])
n = nint + npres

if phase:
    n+=1
    
binary_x = qt.my_binary_repr(digit, n, nint=nint, phase=phase)

qx = QuantumRegister(n, 'x')
qtarg = QuantumRegister(1, 'targ')
qans = QuantumRegister(n-1, 'a')
qlab = QuantumRegister(m, 'l')
out_reg = ClassicalRegister(m,'out')
circ = QuantumCircuit(qx, qtarg, qans, qlab, out_reg)

x_gate = qt.input_bits_to_qubits(binary_x, circ, reg=qx, wrap=True)
circ.append(x_gate, qx);

circ = qt.label_gate(circ, qx, qtarg, qans, qlab, bounds=bounds, nint=nint, phase=phase)

circ.measure(qlab, out_reg);

circ.draw('latex')

shots=10
emulator = Aer.get_backend('qasm_simulator')
job = execute(circ, emulator, shots=shots )
hist = job.result().get_counts()

print('Result:')
for label in hist.keys():
    if qt.bin_to_dec(label, nint=None, phase=False)==0.:
        print(digit,'is outwith the bounds of','['+str(bounds[0])+',',str(bounds[-1])+')','with probability',float(hist[label])/shots)
    else:
        print(bounds[int(qt.bin_to_dec(label, nint=None, phase=False))-1],'<=',digit,'<',bounds[int(qt.bin_to_dec(label, nint=None, phase=False))],'with probability',float(hist[label])/shots)

digit = 1.

bounds = np.array([-2., 0., 1.5, 3.])

print('Domain boundaries:',bounds)

m = int(np.ceil(np.log2(len(bounds))))

phase = True

nint = qt.get_nint([digit,*bounds])
npres = qt.get_npres([digit,*bounds])
n = nint + npres

if phase:
    n+=1
    
binary_x = qt.my_binary_repr(digit, n, nint=nint, phase=phase)
print('Label qubits:',m)

A1s = np.array([3., -1., 2.5, -0.5])

nintc = qt.get_nint(A1s)
npresc = qt.get_npres(A1s)
nc = nintc + npresc

if phase:
    nc+=1
    
n = np.max([n,nc])

qx = QuantumRegister(n, 'x')
qcoff = QuantumRegister(n, 'c')
qlab = QuantumRegister(m, 'l')
out_reg = ClassicalRegister(n,'out')
circ = QuantumCircuit(qx, qcoff, qlab, out_reg)

qtarg = qcoff[0] 
qans = qcoff[1:]

x_gate = qt.input_bits_to_qubits(binary_x, circ, reg=qx, wrap=True)
circ.append(x_gate, qx);

l_gate = qt.label_gate(circ, qx, qtarg, qans, qlab, bounds=bounds, nint=nint, phase=phase, wrap=True)
circ.append(l_gate, [*qx, qtarg, *qans, *qlab]);

circ = qt.cin_gate(circ, qcoff, qlab, A1s, nint=nintc, phase=phase)

circ.measure(qcoff, out_reg);

circ.draw('latex')

shots=10
emulator = Aer.get_backend('qasm_simulator')
job = execute(circ, emulator, shots=shots )
hist = job.result().get_counts()

print('Target:')
ind = np.argwhere((np.array([bounds[-1],*bounds[:-1]])<=digit)*(bounds>digit)).flatten()[0]
print(A1s[ind],'as',bounds[ind-1],'<=',digit,'<',bounds[ind])

print('Result:')
for label in hist.keys():
    print(qt.bin_to_dec(label, nint=nintc, phase=phase),'->',label,'with probability',float(hist[label])/shots)

qx = QuantumRegister(n, 'x')
qcoff = QuantumRegister(n, 'c')
qlab = QuantumRegister(m, 'l')
out_reg = ClassicalRegister(n,'out')
circ = QuantumCircuit(qx, qcoff, qlab, out_reg)

qtarg = qcoff[0] 
qans = qcoff[1:]

x_gate = qt.input_bits_to_qubits(binary_x, circ, reg=qx, wrap=True)
circ.append(x_gate, qx);

l_gate = qt.label_gate(circ, qx, qtarg, qans, qlab, bounds=bounds, nint=nint, phase=phase, wrap=True)
circ.append(l_gate, [*qx, qtarg, *qans, *qlab]);

X_gate = qt.cin_gate(circ, qcoff, qlab, A1s, nint=nintc, phase=phase, wrap=True, label='X')
X_gate_inv = qt.cin_gate(circ, qcoff, qlab, A1s, nint=nintc, phase=phase, wrap=True, inverse=True, label='X')
circ.append(X_gate, [*qcoff, *qlab]);
circ.append(X_gate_inv, [*qcoff, *qlab]);

circ.measure(qcoff, out_reg);

circ.draw('latex')

shots=10
emulator = Aer.get_backend('qasm_simulator')
job = execute(circ, emulator, shots=shots )
hist = job.result().get_counts()

print('Result:')
for label in hist.keys():
    print(qt.bin_to_dec(label, nint=nintc, phase=phase),'->',label,'with probability',float(hist[label])/shots)

def f_x(x):
    return np.arctan(x)

phase = True

nintx = 3
nx = 5
npresx = nx - nintx
if phase:
    npresx-=1

xmin = -2.**(nintx)
xmax = 2.**(nintx) - 2.**(-npresx)

xs = np.linspace(xmin,xmax,2**nx)

npres0, npres1 = 4, 4

norder = 1
m = 3

n, n0, nint, nintcs, coeffs, bounds = qt.optimize_coeffs_qubits(f_x, xs, m, npres0, npres1, norder=norder, phase=phase, label_swap=True)

print('Qubits:', nx, n, n0, m)
print('Integer qubits:', nintx, nint, nintcs[0,0], nintcs[0,1])
print('Memory:', 16*(2**(n+n0+nx+m))/2**20)

_, _, _, _, coeffs_, _ = qt.optimize_coeffs_qubits(f_x, xs, m, npres0, npres1, norder=norder, phase=phase, label_swap=False)
ys_rnd = qt.piecewise_poly(xs, coeffs_.T, bounds)

qx = QuantumRegister(nx, 'x')
qout = QuantumRegister(n, 'o')
qlab = QuantumRegister(m, 'l')
qcoff = QuantumRegister(n0, 'c')
circ = QuantumCircuit(qx, qout, qlab, qcoff)

# 1. Initialize the x register (let as produce a superposition across all x values)
for i in np.arange(nx):
    circ.h(qx[i]);

# We can use the coefficient register as the ancillary register for the labelling step as before
qtarg = qcoff[0] 
qans = [*qcoff[1:], *qout][:nx-1]

# 2. Label the data with the labelling operation
l_gate = qt.label_gate(circ, qx, qtarg, qans, qlab, bounds=bounds, nint=nintx, phase=phase, wrap=True)
circ.append(l_gate, [*qx, qtarg, *qans, *qlab]);

# 3. Load A1 into the coefficient register conditioned on the label register
X1_gate = qt.cin_gate(circ, qcoff, qlab, coeffs[0], nint=nintcs[0,0], phase=phase, wrap=True)
circ.append(X1_gate, [*qcoff, *qlab]);

# 4. Multiply the coefficient and x registers and save the output on the output register
mul_gate = qt.QFTMultPhase(circ, qcoff, qx, qout, wrap=True, nint1=nintcs[0,0], nint2=nintx, nint3=nint)
circ.append(mul_gate, [*qcoff, *qx, *qout]);

# 5. Unload A1 from the coefficient register with the inverse load operation
X1_gate_inv = qt.cin_gate(circ, qcoff, qlab, coeffs[0], nint=nintcs[0,0], phase=phase, wrap=True, inverse=True)
circ.append(X1_gate_inv, [*qcoff, *qlab]);

# 6. Load A0 into the coefficient register conditioned on the label register
X0_gate = qt.cin_gate(circ, qcoff, qlab, coeffs[1], nint=nintcs[0,1], phase=phase, wrap=True)
circ.append(X0_gate, [*qcoff, *qlab]);

# 7. Add the coefficient register to the output register
add_gate = qt.QFTAddition(circ, qcoff, qout, nint1=nintcs[0,1], nint2=nint, wrap=True, phase=phase)
circ.append(add_gate, [*qcoff, *qout]);

# Unload A0 to clear the coefficient register
X0_gate_inv = qt.cin_gate(circ, qcoff, qlab, coeffs[1], nint=nintcs[0,1], phase=phase, wrap=True, inverse=True)
circ.append(X0_gate_inv, [*qcoff, *qlab]);

# Unlabel the label register
l_gate_inv = qt.label_gate(circ, qx, qtarg, qans, qlab, bounds=bounds, nint=nintx, phase=phase, wrap=True, inverse=True)
circ.append(l_gate_inv, [*qx, qtarg, *qans, *qlab]);

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

state_v = np.asarray(state_vector).reshape((2**n0,2**m,2**n,2**nx)).T
state_v = np.sum(state_v, axis=(2,3))

y_unorm = np.argwhere(np.round(np.abs(state_v)**2,15)>0.)[:,1]
y_out = []
for y in y_unorm:
    y_out.append(qt.bin_to_dec(qt.my_binary_repr(y, n=n, nint=None, phase=False), nint=nint, phase=phase))

xs_comp = []
xs_ = np.linspace(-2**(nx-1),2**(nx-1)-1,2**nx)
for x in xs_:
    binary = qt.my_binary_repr(x, nx, nint=None, phase=True)
    xs_comp.append(int(qt.bin_to_dec(binary, nint=None, phase=False)))

y_out = np.array(y_out)[xs_comp]

fig = plt.figure(figsize=figsize)
ax = fig.gca()

for bound in bounds:
    ax.axvline(bound, ls=':', color='black')
    
ax.plot(xs, f_x(xs), color='black');
ax.scatter(xs, y_out, color='black', lw=2);
ax.plot(xs, ys_rnd, color='black', ls=':', lw=3);
ax.set_xlabel(r'$x$', fontsize=fontsize);
ax.set_ylabel(r'$f(x)$', fontsize=fontsize);
ax.set_xlim(xmin,xmax);
ax.tick_params(axis='both', labelsize=ticksize);
plt.show()

def f_x(x):
    return np.arctan(x)

phase = True
phasex = True

nintx = 3
nx = 5
npresx = nx - nintx

xmin = 0.

if phasex:
    npresx-=1
    xmin = -2.**(nintx)
    
xmax = 2.**(nintx) - 2.**(-npresx)


xs = np.linspace(xmin,xmax,2**nx)

n, n0, nint, nintcs, coeffs, bounds = qt.optimize_coeffs_qubits(f_x, xs, m, npres0, npres1, norder=norder, phase=phase, label_swap=True)

print('Qubits:', nx, n, n0, m)
print('Integer qubits:', nintx, nint, nintcs[0,0], nintcs[0,1])
print('Memory:', 16*(2**(n+n0+nx+m))/2**20)

qx = QuantumRegister(nx, 'x')
qout = QuantumRegister(n, 'o')
qlab = QuantumRegister(m, 'l')
qcoff = QuantumRegister(n0, 'c')
circ = QuantumCircuit(qx, qout, qlab, qcoff)

for i in np.arange(nx):
    circ.h(qx[i]);

lpf_gate = qt.lpf(circ, qx, qout, qlab, qcoff, coeffs, bounds, nintx=nintx, nintcs=nintcs, nint=nint, phase=phase, phasex=phasex, wrap=True)
circ.append(lpf_gate, [*qx, *qout, *qlab, *qcoff]);

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

state_v = np.asarray(state_vector).reshape((2**n0,2**m,2**n,2**nx)).T
state_v = np.sum(state_v, axis=(2,3))

y_unorm = np.argwhere(np.round(np.abs(state_v)**2,15)>0.)[:,1]
y_out = []
for y in y_unorm:
    y_out.append(qt.bin_to_dec(qt.my_binary_repr(y, n=n, nint=None, phase=False), nint=nint, phase=phase))

xs_comp = []
xs_ = np.linspace(-2**(nx-1),2**(nx-1)-1,2**nx)
for x in xs_:
    binary = qt.my_binary_repr(x, nx, nint=None, phase=True)
    xs_comp.append(int(qt.bin_to_dec(binary, nint=None, phase=False)))

y_out = np.array(y_out)[xs_comp]

Ax = qt.piecewise_poly(xs, np.array([coeffs_[0],np.zeros(len(coeffs_[1]))]).T, bounds)

fig = plt.figure(figsize=figsize)
ax = fig.gca()

for bound in bounds:
    ax.axvline(bound, ls=':', color='black')
    
ax.plot(xs, f_x(xs), color='black');
ax.scatter(xs, y_out, color='black', lw=2);
#ax.plot(xs, Ax, color='black', lw=2);
ax.set_xlabel(r'$x$', fontsize=fontsize);
ax.set_ylabel(r'$f(x)$', fontsize=fontsize);
ax.set_xlim(xmin,xmax);
ax.tick_params(axis='both', labelsize=ticksize);
plt.show()