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

digit = 10.25
k = 2
target = np.power(digit, k)

print(digit,'^'+str(k),'=',target)

nintx = qt.get_nint(digit)
npresx = qt.get_npres(digit)
nx = nintx + npresx + 1
print(nx, nintx, npresx)

nint = qt.get_nint(target)
npres = qt.get_npres(target)
n = nint + npres + 1
print(n, nint)

phase = True

binary = qt.my_binary_repr(digit, nx, nint=nintx, phase=phase)
binary_t = qt.my_binary_repr(target, n, nint=nint, phase=phase)

qx = QuantumRegister(nx, 'x')
qtarg = QuantumRegister(n,'targ')
c_out = ClassicalRegister(n, 'c_out')
circ = QuantumCircuit(qx, qtarg, c_out)

in_gate = qt.input_bits_to_qubits(binary, circ, reg=qx, wrap=True)
circ.append(in_gate, qx);

power_gate = qt.QFTPowerN(circ, qx, qtarg, k, wrap=True, nintx=nintx, nint=nint)
circ.append(power_gate, [*qx, *qtarg]);

circ.measure(qtarg, c_out);

circ.draw('latex')

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

print('Target:')
print(digit,'^'+str(k),'=',target,'->',binary_t)

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

digit = 1.25
k = 5
target = np.power(digit, k)

print(digit,'^'+str(k),'=',target)

nintx = qt.get_nint(digit)
npresx = qt.get_npres(digit)
nx = nintx + npresx + 1
print(nx, nintx, npresx)

nint = qt.get_nint(target)
npres = qt.get_npres(target)
n = nint + npres + 1
print(n, nint)

phase = True

binary = qt.my_binary_repr(digit, nx, nint=nintx, phase=phase)
binary_t = qt.my_binary_repr(target, n, nint=nint, phase=phase)

qx = QuantumRegister(nx, 'x')
qtarg = QuantumRegister(n,'targ')
c_out = ClassicalRegister(n, 'c_out')
circ = QuantumCircuit(qx, qtarg, c_out)

in_gate = qt.input_bits_to_qubits(binary, circ, reg=qx, wrap=True)
circ.append(in_gate, qx);

power_gate = qt.QFTPowerN(circ, qx, qtarg, k, wrap=True, nintx=nintx, nint=nint)
circ.append(power_gate, [*qx, *qtarg]);

circ.measure(qtarg, c_out);

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

print('Target:')
print(digit,'^'+str(k),'=',target,'->',binary_t)

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