import math
import numpy as np
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, Aer, execute
from qiskit.quantum_info.operators import Operator
from qiskit.circuit.library import CU1Gate
from qiskit.visualization import plot_histogram
from fractions import Fraction

# Provided helper functions and oracles

def make_permutation_matrix(n, permutation):
    r = np.zeros((n,n), dtype=int)
    for i in range(n):
        r[permutation(i), i] = 1
    return r

def mult_mat(x,k,N):
    n = math.ceil(math.log(N, 2))
    return make_permutation_matrix(
    2**n,
    permutation=lambda y:  y*pow(x,k) % N if y<N else y)

def c_mult_mat(x,k,N):
    n = math.ceil(math.log(N, 2))
    permutation = lambda y: y if y%2==0 or (y-1)/2 >= N else 2*(int((y-1)/2)*pow(x,k) % N) + 1
    return make_permutation_matrix(2*(2**n), permutation )

def mult_op(x,k,N):
    return Operator(mult_mat(x,k,N))


#controlled-U oracle 
def c_mult_op(x,k,N):
    return Operator(c_mult_mat(x,k,N))


# QFT and IQFT
def qft(circ, q, n):
    """n-qubit QFT on q in circ."""
    for j in range(n):
        circ.h(q[j])
        for k in range(j+1,n):
            circ.cu1(math.pi/float(2**(k-j)), q[k], q[j])
    for i in range(n//2):
        circ.swap(q[i], q[n-i-1])
        
def iqft(circ,q,n):
    for j in range(n):
        circ.h(q[j])
        for k in range(j+1,n):
            #circ.cu1(-math.pi/float(2**(k-j)), q[k], q[j])
             gate = CU1Gate(-np.pi/float(2**(k-j)) )
             circ.append(gate, [q[k],q[j]])
    for i in range(n//2):
        circ.swap(q[i], q[n-i-1])

# Constants
N = 15
x = 7

# Quantum registers
n_control = math.ceil(math.log(N, 2)) * 2  # Number of qubits in the control register
control_qubits = QuantumRegister(n_control)
target_qubits = QuantumRegister(math.ceil(math.log(N, 2)) + 1) # Extra qubit for modulo multiplication
classical_bits = ClassicalRegister(n_control)

# Quantum circuit
circuit = QuantumCircuit(control_qubits, target_qubits, classical_bits)

# Initialize control qubits in superposition
circuit.h(control_qubits)

# Apply controlled unitary operations
for qubit in range(n_control):
    k = 2 ** qubit
    circuit.append(c_mult_op(x, k, N), [control_qubits[qubit]] + target_qubits[:])

# Apply inverse QFT
iqft(circuit, control_qubits, n_control)

# Measure control qubits
circuit.measure(control_qubits, classical_bits)

# Simulate the circuit
simulator = Aer.get_backend('qasm_simulator')
result = execute(circuit, simulator, shots=1024).result()
counts = result.get_counts(circuit)

# Display the histogram
plot_histogram(counts)

# Determine the order r
def determine_order(counts):
    # Find the output with the highest probability
    most_common_outcome = max(counts, key=counts.get)
    phase = int(most_common_outcome, 2) / (2**n_control)
    frac = Fraction(phase).limit_denominator(N)
    return frac.denominator

r = determine_order(counts)
print(f"Order r: {r}")

# Post-processing: Find factors of N
factor_found = False
if r % 2 == 0:
    plus = math.gcd(int(pow(x, r//2) + 1), N)
    minus = math.gcd(int(pow(x, r//2) - 1), N)
    if plus != 1 and plus != N:
        print(f"One factor of {N} is {plus}")
        factor_found = True
    if minus != 1 and minus != N:
        print(f"One factor of {N} is {minus}")
        factor_found = True

if not factor_found:
    print("No factors found. Try a different random x.")

