import numpy as np
from math import pi
from qiskit import ClassicalRegister, QuantumCircuit, QuantumRegister, BasicAer, execute, IBMQ
from qiskit.visualization import plot_bloch_multivector, plot_histogram

def pe_circ_init(size:int):
    #quantum register and  single state for storing eigen-state
    qr = QuantumRegister(size)
    psi = QuantumRegister(1,name="psi")
    #classical bits for readout
    cr = ClassicalRegister(size)
    
    qc = QuantumCircuit(qr,psi,cr)
    # apply Hadamards and init the eigenstate as |1>
    for i in range(size):
        qc.h(i)
    qc.x(size)
    
    qc.barrier()
    
    return qc

def C_U(circuit:QuantumCircuit, size:int, phi:float):
    
    rep = 1
    for counting_qubit in range(size):
        for i in range(rep):
            circuit.cu1(2*pi*phi, counting_qubit, size)
        rep *= 2
    
    circuit.barrier()
    
    return 

def qft_dagger(qc:QuantumCircuit,size:int):
    
    for qubit in range(size//2):
        qc.swap(qubit,size-qubit-1)
    for j in range(size):
        for m in range(j):
            qc.cu1(-pi/float(2**(j-m)), m, j)
        qc.h(j)
    qc.barrier()
    return

def readout(qc:QuantumCircuit,size:int):
    for i in range(size):
        qc.measure(i,i)
    return

def phaseEstCircuit(size:int, phi:float):
    qc = pe_circ_init(size)
    C_U(qc, size, phi)
    qft_dagger(qc, size)
    readout(qc,size)
    
    return qc

def counts_decode(counts):
    decoded = {}
    for key in counts:
        decimal_key = 0.0
        i = -1
        for c in key:
            b = int(c)
            if b:
                decimal_key += 2**i
            i -= 1
        decoded[decimal_key] = counts[key]
    return decoded

def expectation(decode_counts):
    res = 0.0
    total = 0
    for key in decode_counts:
        total += decode_counts[key]
        res += key*decode_counts[key]
    return res/total

def choosemax(decode_counts):
    counts = 0
    result = 0.0
    for key in decode_counts:
        if decode_counts[key] > counts:
            counts = decode_counts[key]
            result = key
    return  result

def getEstPhi(PhaseEstQC:QuantumCircuit, back_end, shots, metric):
    # get backend
    result = execute(qc,backend=back_end,shots=shots).result()
    counts = result.get_counts()
    plot_histogram(counts)

    dt = counts_decode(counts)
    est_val = metric(dt)

    return est_val

qubit = 4
nTrail = 20
back_end = BasicAer.get_backend("qasm_simulator")
shots = 4096

Phi = np.random.random((nTrail,))

Phi_est = []
for phi in Phi:
    qc = phaseEstCircuit(qubit, phi)
    single_run_res = getEstPhi(qc, back_end,shots,metric=choosemax)
    Phi_est.append(single_run_res)
Phi_est = np.array(Phi_est)

error = np.abs(Phi_est-Phi)/Phi

print(Phi)
print(Phi_est)
print(error)

qubit = 5
nTrail = 20
back_end = BasicAer.get_backend("qasm_simulator")
shots = 4096

Phi_est = []
for phi in Phi:
    qc = phaseEstCircuit(qubit, phi)
    single_run_res = getEstPhi(qc, back_end,shots,metric=choosemax)
    Phi_est.append(single_run_res)
Phi_est = np.array(Phi_est)

error = np.abs(Phi_est-Phi)/Phi
print(Phi)
print(Phi_est)
print(error)

qubit = 6
nTrail = 20
back_end = BasicAer.get_backend("qasm_simulator")
shots = 4096

Phi_est = []
for phi in Phi:
    qc = phaseEstCircuit(qubit, phi)
    single_run_res = getEstPhi(qc, back_end,shots,metric=choosemax)
    Phi_est.append(single_run_res)
Phi_est = np.array(Phi_est)

error = np.abs(Phi_est-Phi)/Phi
print(Phi)
print(Phi_est)
print(error)

qubit = 7
nTrail = 20
back_end = BasicAer.get_backend("qasm_simulator")
shots = 4096

Phi_est = []
for phi in Phi:
    qc = phaseEstCircuit(qubit, phi)
    single_run_res = getEstPhi(qc, back_end,shots,metric=choosemax)
    Phi_est.append(single_run_res)
Phi_est = np.array(Phi_est)

error = np.abs(Phi_est-Phi)/Phi
print(Phi)
print(Phi_est)
print(error)

qubit = 8
nTrail = 20
back_end = BasicAer.get_backend("qasm_simulator")
shots = 4096

Phi_est = []
for phi in Phi:
    qc = phaseEstCircuit(qubit, phi)
    single_run_res = getEstPhi(qc, back_end,shots,metric=choosemax)
    Phi_est.append(single_run_res)
Phi_est = np.array(Phi_est)

error = np.abs(Phi_est-Phi)/Phi
print(Phi)
print(Phi_est)
print(error)





def get_error(size, Phi, back_end, shots, metric):
    Phi_est = []
    for phi in Phi:
        qc = phaseEstCircuit(qubit, phi)
        qc.draw()
        single_run_res = getEstPhi(qc, back_end,shots,metric)
        Phi_est.append(single_run_res)
    Phi_est = np.array(Phi_est)

    return np.abs(Phi_est-Phi)/Phi

error_4 = get_error(size=4, Phi=Phi, back_end=back_end,shots=2048,metric=choosemax)

error_5 = get_error(size=5, Phi=Phi, back_end=back_end,shots=2048,metric=choosemax)

error_6 = get_error(size=6, Phi=Phi, back_end=back_end,shots=2048,metric=choosemax)

import matplotlib.pyplot as plt

plt.plot(Phi,error_4,'rs',Phi,error_5,'bs',Phi,error_6,'gs')

error_4

error_5

