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
from qiskit.circuit.library.standard_gates import RYGate

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

def targ_dist(x, bias=0):
    y = (x+bias)**(-7./6)
    return y/np.sqrt(np.sum(np.abs(y)**2))    

nx = 6
nintx = 4

xmin = -2.**(nintx)
xmax = 2.**(nintx) - 2.**(-nx+nintx+1)
xs = np.linspace(xmin, xmax, 2**nx)

amps = targ_dist(xs, bias=2.**(nintx+1))
probs = np.abs(amps)**2

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

ax.plot(xs, probs, color='black')
ax.set_xlabel(r'$x$', fontsize=fontsize);
ax.set_ylabel(r'$\sqrt{p(x)}$', fontsize=fontsize);
ax.set_xlim(xmin,xmax);
ax.tick_params(axis='both', labelsize=ticksize);

def GR_func(j, probs, m, n):
    j = np.array(j).astype(int)
    As = []
    for i in np.arange(2**m):
        As.append(np.sum(probs[i*2**(n-m):(i+1)*2**(n-m)]))
    As1 = []
    for i in np.arange(2**(m+1)):
        As1.append(np.sum(probs[i*2**(n-(m+1)):(i+1)*2**(n-(m+1))]))
    return np.arccos(np.sqrt(np.array(As1)[::2][j]/np.array(As)[j]))

qx = QuantumRegister(nx, 'x')
circ = QuantumCircuit(qx)

coeffs = np.arccos(np.sqrt(np.sum(probs[:2**(nx-1)])))
circ.ry(2*coeffs, qx[nx-1]);

for m in np.arange(1,nx):
    js = np.arange(2**m)
    coeffs = GR_func(js, probs, m, nx)
    
    for j in np.arange(2**m):
        control_bits = qt.my_binary_repr(j, m, nint=None, phase=False)
        if j>0:
            control_bits_ = qt.my_binary_repr(j-1, m, nint=None, phase=False)[::-1]
        else:
            control_bits_ = np.ones(m).astype(int).astype(str)
        
        for i,control_bit in enumerate(control_bits[::-1]):
            if control_bit=='0' and control_bits_[i]=='1':
                circ.x(qx[nx-i-1]);

        R_gate = RYGate(2*coeffs[j]).control(int(m))
        circ.append(R_gate, [*qx[nx-m-1:][::-1]]);

        if j<2**m-1:
            control_bits_ = qt.my_binary_repr(j+1, m, nint=None, phase=False)[::-1]
        else:
            control_bits_ = np.ones(m).astype(int).astype(str)

        for i,control_bit in enumerate(control_bits[::-1]):
            if control_bit=='0' and control_bits_[i]=='1':
                circ.x(qx[nx-i-1]);

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

if np.any(state_vector.imag>0.):
    state_v = np.copy(state_vector)
else:
    state_v = state_vector.real

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

ax.scatter(xs, np.abs(state_v)**2, color='black')
ax.plot(xs, probs, color='black')
ax.set_xlabel(r'$x$', fontsize=fontsize);
ax.set_ylabel(r'$\sqrt{p(x)}$', fontsize=fontsize);
ax.set_xlim(xmin,xmax);
ax.tick_params(axis='both', labelsize=ticksize);

circ.draw('latex')

# Register size for storing the zeta angles
nanc = 8
# Register size for storing the linear coefficients for the piecewise function 
ncoff = 8
# Register size for the label register
nlab = 4

print('Qubits:', nx, nanc, ncoff, nlab,'total:',nanc+ncoff+nx+nlab)
print('Memory:', 16*(2**(nanc+ncoff+nx+nlab))/2**20)

qx = QuantumRegister(nx, 'x')
qanc = QuantumRegister(nanc, 'anc')
qlab = QuantumRegister(nlab, 'l')
qcoff = QuantumRegister(ncoff, 'c')
circ = QuantumCircuit(qx, qanc, qlab, qcoff)

mlow = 4

probs_ = []
for i in np.arange(2**mlow):
    probs_.append(np.sum(probs[i*2**(nx-mlow):(i+1)*2**(nx-mlow)]))
probs_ = np.array(probs_)

GRl_gate = qt.Grover_Rudolph_load(circ, qx[nx-mlow:], probs_, wrap=True)
circ.append(GRl_gate, qx[nx-mlow:]);

for m in np.arange(mlow,nx):
    
    js = np.arange(2**m)
    coeffs = GR_func(js, probs, m, nx)

    f_x = lambda x: GR_func(x, probs, m, nx)
    
    bounds_ = np.linspace(0,2**m,(2**nlab)+1).astype(int)
    coeffs = qt.get_bound_coeffs(f_x, bounds_, 1, reterr=False).T#[::-1]
    bounds = bounds_[1:]

    max_list0 = np.array([coeffs[0], coeffs[1], coeffs[0]*2**nx, (coeffs[0]*2**nx)+coeffs[-1]])
    max_list1 = max_list0
    nintcs = []
    nintcs.append(int(np.ceil(np.log2(np.max(np.abs(max_list0))))))
    nintcs.append(int(np.ceil(np.log2(np.max(np.abs(max_list1))))))
    nint = nintcs[-1]
    nintcs = np.array([nintcs])

    func_gate = qt.linear_piecewise_approx(circ, qx[nx-m:], qanc, qlab, qcoff, coeffs, bounds, nint=nint, nintx=nx, nintcs=nintcs, phase=False, wrap=True)
    circ.append(func_gate, [*qx[nx-m:], *qanc, *qlab, *qcoff]);

    rot_gate = qt.CRotation(circ, qanc, qx[nx-m-1], nint=nint, wrap=True)
    circ.append(rot_gate, [*qanc, qx[nx-m-1]])

    func_gate_inv = qt.linear_piecewise_approx(circ, qx[nx-m:], qanc, qlab, qcoff, coeffs, bounds, nint=nint, nintx=nx, nintcs=nintcs, phase=False, wrap=True, inverse=True)
    circ.append(func_gate_inv, [*qx[nx-m:], *qanc, *qlab, *qcoff]);

circ.draw('latex')

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

if np.any(state_vector.imag>0.):
    state_v = np.copy(state_vector)
else:
    state_v = state_vector.real

state_v = np.asarray(state_vector).reshape((2**ncoff,2**nlab,2**nanc,2**nx)).T
state_v = np.sum(state_v, axis=(1,2,3))

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

ax.scatter(xs, np.abs(state_v)**2, color='black')
ax.plot(xs, probs, color='black')
ax.set_xlabel(r'$x$', fontsize=fontsize);
ax.set_ylabel(r'$\sqrt{p(x)}$', fontsize=fontsize);
ax.set_xlim(xmin,xmax);
ax.tick_params(axis='both', labelsize=ticksize);