import config
import random
from qiskit import QuantumRegister, ClassicalRegister, QuantumCircuit, IBMQ, Aer, execute
from qiskit.visualization import plot_histogram
from qiskit.providers.ibmq import least_busy

def d_oracle(qc, oracle, choice):
    """
    Add oracle for Deutsch's algoritm, in the form of a gate, to quantum circuit.
    
    Parameters:
    -----------
    qc: QuantumCircuit
        Quantum circuit with 2 qubits.
    oracle: int
        Type of oracle. 0 for balanced and 1 for constant.
    choice: int
        Choice of oracle type. For a balanced oracle, 0 will make f(x)=x and 1 will make f(x)=¬x. 
        For a constant oracle, 0 will make f(x)=0 and 1 will make f(x)=1.
        
    Returns:
    --------
    qc: QuantumCircuit
        Quantum circuit with Deutsch's algorithm oracle added at the end.
    """
    qreg           = QuantumRegister(2, 'q')
    oracle_circuit = QuantumCircuit(qreg, name='$U_f$')
    
    # Balanced oracle
    if oracle == 0:
        if choice == 0:
            oracle_circuit.cx(0, 1)
        elif choice == 1:
            oracle_circuit.cx(0, 1)
            oracle_circuit.x(1)
            
    # Constant oracle
    elif oracle == 1:
        if choice == 0:
            pass # qubit is already initialized to 0
        elif choice == 1:
            oracle_circuit.x(1)
    
    # Turn oracle into single gate and append to qc
    oracle_gate = oracle_circuit.to_gate()
    qc.append(oracle_gate, [0, 1])
    
    return qc

def deutsch(oracle = None, choice = None):
    """
    Create quantum circuit implementing Deutsch's algorithm.
    
    Parameters:
    -----------
    oracle: int
        Type of oracle. 0 for balanced and 1 for constant.
    choice: int
        Choice of oracle type. For a balanced oracle, 0 will make f(x)=x and 1 will make f(x)=¬x. 
        For a constant oracle, 0 will make f(x)=0 and 1 will make f(x)=1.
        
    NOTE: if any parameter is set to none (default), the value will be randomized.    
    
    Returns:
    --------
    qc: QuantumCircuit
        Quantum circuit implementing Deutsch's algorithm.
    """
    # Randomize oracle and choice if not given
    if oracle == None:
        oracle = random.randint(0, 1)
    if choice == None:
        choice = random.randint(0, 1)
    
    qreg = QuantumRegister(2, 'q')
    creg = ClassicalRegister(1, 'c')
    qc   = QuantumCircuit(qreg, creg)
    
    # Prepare initial state
    qc.x(qreg[1])
    
    qc.barrier()
    
    # Put into superpusition
    qc.h(qreg)
    
    qc.barrier()
    
    # Add oracle
    d_oracle(qc, oracle, choice)
    
    qc.barrier()
    
    # Apply interference
    qc.h(qreg[0])
    
    qc.barrier()
    
    # Measure first qubit
    qc.measure(qreg[0], creg)
    
    # Return circuit
    return qc

circ = deutsch()
circ.draw()

circ_balanced = deutsch(0)
circ_constant = deutsch(1) # We will let the choice be random for both circuits

IBMQ.save_account(f"{config.IBM_KEY}", overwrite = True)
IBMQ.load_account()

backend = Aer.get_backend('qasm_simulator')

# Uncomment below to run on hardware
# provider = IBMQ.get_provider(hub='ibm-q')
# backend  = least_busy(provider.backends())

result  = execute(circ_balanced, backend, shots=1024).result()
count_b = result.get_counts(circ_balanced)

plot_histogram(count_b)

backend = Aer.get_backend('qasm_simulator')

# Uncomment below to run on hardware
# provider = IBMQ.get_provider(hub='ibm-q')
# backend  = least_busy(provider.backends())

result  = execute(circ_constant, backend, shots=1024).result()
count_c = result.get_counts(circ_constant)

plot_histogram(count_c)

def dj_oracle(qc, n, oracle, choice, to_gate=True):
    """
    Add oracle for Deutsch-Jozsa algoritm, in the form of a gate, to quantum circuit.
    
    Parameters:
    -----------
    qc: QuantumCircuit
        Quantum circuit with n+1 qubits.
    n: int
        Number of qubits in first register. n-input to Deutsch-Jozsa. 
        Should be one less that the total size of the circuit.
    oracle: int
        Type of oracle. 0 for balanced and 1 for constant.
    choice: int
        Choice of oracle type. For a balanced oracle, this parameter is irrelevant. 
        For a constant oracle, 0 will make f(x)=0 and 1 will make f(x)=1.
    to_gate: bool
        If true, the oracle will be added as a gate to the original circuit. 
        If false, the individual gates will be added to the original circuit.
        
    Returns:
    --------
    qc: QuantumCircuit
        Quantum circuit with Deutsch-Jozsa oracle added at the end.
    """
    qc_len = 0
    for qr in qc.qregs:
        qc_len += len(qr)
    if qc_len != n + 1:
        raise Exception(('Length of oracle and circuit given do not match correctly.'
                         ' Parameter n should be one less than number of qubits in parameter qc.'))
    
    if to_gate:
        qreg_input     = QuantumRegister(n, 'x')
        qreg_output    = QuantumRegister(1, 'y')
        oracle_circuit = QuantumCircuit(qreg_input, qreg_output, name='$U_f$')

        # Balanced oracle
        if oracle == 0:
            for i in range(n):
                oracle_circuit.cx(qreg_input[i], qreg_output)

        # Constant oracle
        elif oracle == 1:
            if choice == 0:
                pass # y is already initialized to 0
            elif choice == 1:
                oracle_circuit.x(qreg_output)
        
        oracle_gate = oracle_circuit.to_gate()
        qc.append(oracle_gate, range(n+1))
    else:        
        # Balanced oracle
        if oracle == 0:
            for i in range(n):
                qc.cx(i, n)

        # Constant oracle
        elif oracle == 1:
            if choice == 0:
                pass # y is already initialized to 0
            elif choice == 1:
                qc.x(n)
        
    return qc

def deutsch_jozsa(n, oracle=None, choice=None, oracle_gate=True):
    """
    Create quantum circuit implementing n-input Deutsch-Jozsa algorithm.
    
    Parameters:
    -----------
    n: int
        Number of qubits in first register. n-input to Deutsch-Jozsa. 
        Should be one less that the total size of the circuit.
    oracle: int
        Type of oracle. 0 for balanced and 1 for constant.
    choice: int
        Choice of oracle type. For a balanced oracle, this parameter is irrelevant. 
        For a constant oracle, 0 will make f(x)=0 and 1 will make f(x)=1.
    oracle_gate: bool
        If true, the oracle will be added as a gate to the original circuit. 
        If false, the individual gates enclosed by barriers will be added to the original circuit.
        
    Returns:
    --------
    qc: QuantumCircuit
        Quantum circuit with Deutsch-Jozsa oracle added at the end.
    """
    # Randomize oracle and choice if not given
    if oracle == None:
        oracle = random.randint(0, 1)
    if choice == None:
        choice = random.randint(0, 1)
    
    qreg_in  = QuantumRegister(n, 'x')
    qreg_out = QuantumRegister(1, 'y')
    creg     = ClassicalRegister(n, 'c')
    qc       = QuantumCircuit(qreg_in, qreg_out, creg)
    
    # Prepare initial state
    qc.x(qreg_in)
    
    qc.barrier()
    
    # Put into superpusition
    qc.h(range(n+1))
    
    qc.barrier()
    
    # Add oracle
    dj_oracle(qc, n, oracle, choice, oracle_gate)
    
    qc.barrier()
    
    # Apply interference
    qc.h(qreg_in)
    
    qc.barrier()
    
    # Measure first qubit
    qc.measure(qreg_in, creg)
    
    return qc

qc = deutsch_jozsa(3, oracle=0, oracle_gate=False)
qc.draw()

backend = Aer.get_backend('qasm_simulator')

# Uncomment below to run on hardware
# provider = IBMQ.get_provider(hub='ibm-q')
# backend  = least_busy(provider.backends())

result = execute(qc, backend, shots=1024).result()
counts = result.get_counts(qc)

plot_histogram(counts)