import cirq
import random
import matplotlib.pyplot as plt
import numpy as np
import math

def qft_dagger_cirq(qc, qubits, n):
    for qubit in range(n//2):
        qc.append(cirq.SWAP(qubits[qubit], qubits[n-qubit-1]))
    for j in range(n):
        for m in range(j):
            qc.append((cirq.CZ**(-1/2**(j-m)))(qubits[m],qubits[j]))
        qc.append(cirq.H(qubits[j]))

def generalised_qpe_cirq(amt_estimation_qubits, angle, shots = 10000):
    # Create and set up circuit
    qubits = cirq.LineQubit.range(amt_estimation_qubits+1)
    circuit = cirq.Circuit()
    
    # Apply H-Gates to counting qubits:
    for qubit in range(amt_estimation_qubits):
        circuit.append(cirq.H(qubits[qubit]))

    #print(circuit)    
        
    # Prepare our eigenstate |psi>:
    circuit.append(cirq.X(qubits[amt_estimation_qubits]))

    #print(circuit) 
    
    repetitions = 1
    for counting_qubit in range(amt_estimation_qubits):
        for i in range(repetitions):
            circuit.append((cirq.CZ**((angle/math.pi)))(qubits[counting_qubit],qubits[amt_estimation_qubits]));
        repetitions *= 2
        
    #print(circuit) 

    # Do the inverse QFT:
    qft_dagger_cirq(circuit, qubits, amt_estimation_qubits)
    
    # Measure of course!
    circuit.append(cirq.measure(*qubits[:-1], key='m'))
        
    simulator = cirq.Simulator()
    results = simulator.run(circuit , repetitions = shots)
    theta_estimates = np.sum(2 ** np.arange(amt_estimation_qubits) * results.measurements['m'], axis=1) / 2**amt_estimation_qubits
    
    print(theta_estimates)
    counts = dict()
    for i in theta_estimates:
        counts[i] = counts.get(i, 0) + 1
    print(counts)
    
    unique,pos = np.unique(theta_estimates,return_inverse=True)
    counts = np.bincount(pos)
    maxpos = counts.argmax()
    print(unique[maxpos],counts[maxpos])
    
#generalised_qpe_cirq(7,(13*math.pi/9))
generalised_qpe_cirq(7,(13*math.pi/9))

