import numpy as np
import os

from qiskit import transpile, QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit_aer import Aer
from qiskit.providers.basic_provider import BasicProvider
from qiskit.visualization import circuit_drawer, plot_histogram
from IPython.display import display, Math, Latex

from quantuminspire.credentials import get_authentication
from quantuminspire.api import QuantumInspireAPI
from quantuminspire.qiskit import QI

QI_URL = os.getenv('API_URL', 'https://api.quantum-inspire.com/')

def format_vector(state_vector, decimal_precision=7):
    """ Format the state vector into a LaTeX formatted string.

    Args:
        state_vector (list or array): The state vector with complex
                                      values e.g. [-1, 2j+1].

    Returns:
        str: The LaTeX format.
    """
    result = []
    epsilon = 1/pow(10, decimal_precision)
    state_vector = np.asarray(state_vector)
    bit_length = (len(state_vector) - 1).bit_length()
    for index, complex_value in enumerate(state_vector):
        has_imag_part = np.round(complex_value.imag, decimal_precision) != 0.0
        value = complex_value if has_imag_part else complex_value.real
        value_round = np.round(value, decimal_precision)
        if np.abs(value_round) < epsilon:
            continue

        binary_state = '{0:0{1}b}'.format(index, bit_length)
        result.append(r'{0:+2g}\left\lvert {1}\right\rangle '.format(value_round, binary_state))
    return ''.join(result)

def run_circuit(q_circuit, q_register, number_of_qubits=None, backend_name='statevector_simulator'):
    """ Run a circuit on all base state vectors and show the output.

    Args:
        q_circuit (QuantumCircuit):
        q_register (QuantumRegister)
        number_of_qubits (int or None): The number of qubits.
        backend (str): ...
    """
    if not isinstance(number_of_qubits, int):
        number_of_qubits = q_register.size

    if q_register.size != number_of_qubits:
        warnings.warn('incorrect register size?')

    latex_text = r'\mathrm{running\ circuit\ on\ set\ of\ basis\ states:}'
    display(Math(latex_text))

    base_states = 2 ** number_of_qubits
    backend = Aer.get_backend(backend_name)
    for base_state in range(base_states):
        pre_circuit = QuantumCircuit(q_register)
        state = base_state
        for kk in range(number_of_qubits):
            if state % 2 == 1:
                pre_circuit.x(q[kk])
            state = state // 2

        input_state = r'\left\lvert{0:0{1}b}\right\rangle'.format(base_state, number_of_qubits)
        circuit_total = pre_circuit.compose(q_circuit)

        circuit_total = transpile(circuit_total, backend)
        job = backend.run(circuit_total, shots=10)
        output_state = job.result().get_statevector(circuit_total)

        latex_text = input_state + r'\mathrm{transforms\ to}: ' + format_vector(output_state)
        display(Math(latex_text))

n=3
N=2**n
q = QuantumRegister(n)
qc = QuantumCircuit(q)

if n==1:
    def black_box(qc, q):
        qc.z(q)
elif n==2:
    def black_box(qc, q):
        for i in range(n):
            qc.s(q[i])
        qc.h(q[1])
        qc.cx(q[0], q[1])
        qc.h(q[1])
        for i in range(n):
            qc.s(q[i])
else:
    def black_box(qc, q):
        qc.h(q[2])
        qc.ccx(q[0], q[1], q[2])
        qc.h(q[2])
black_box(qc,q)
cplot=qc.draw(output='mpl')
display(cplot)

print('black box circuit:')
run_circuit(qc, q)

def n_controlled_Z(circuit, controls, target):
    """Implement a Z gate with multiple controls"""
    if (len(controls) > 2):
        raise ValueError('The controlled Z with more than 2 ' +
                         'controls is not implemented')
    elif (len(controls) == 1):
        circuit.h(target)
        circuit.cx(controls[0], target)
        circuit.h(target)
    elif (len(controls) == 2):
        circuit.h(target)
        circuit.ccx(controls[0], controls[1], target)
        circuit.h(target)

def inversion_about_average(circuit, f_in, n):
    """Apply inversion about the average step of Grover's algorithm."""
    # Hadamards everywhere
    if n==1:
        circuit.x(f_in[0])
        return
    for j in range(n):
        circuit.h(f_in[j])
    # D matrix: flips the sign of the state |000> only
    for j in range(n):
        circuit.x(f_in[j])
    n_controlled_Z(circuit, [f_in[j] for j in range(n-1)], f_in[n-1])
    for j in range(n):
        circuit.x(f_in[j])
    # Hadamards everywhere again
    for j in range(n):
        circuit.h(f_in[j])

qInvAvg = QuantumCircuit(q)
inversion_about_average(qInvAvg, q, n)
qInvAvg.draw(output='mpl')

print('inversion average circuit:')
qc = QuantumCircuit(q)
inversion_about_average(qc, q, n)
run_circuit(qc, q, n)

theta = 2*np.arcsin(1/np.sqrt(N))
r=np.pi*np.sqrt(N)/4
display(Math(r'\textrm{Rotation of the winner: } \theta = %.2f \mathrm{\ [deg]}' % (np.rad2deg(theta))) )
print('Optimal number of Grover iterations for n=%d: %.1f' % (n,r) )
T=int(r)

for i in range(int(r+2)):
    p=np.sin((i+1/2)*theta)**2
    print('%d iterations: p %.2f' % (i, p))

"""Grover search implemented in Qiskit.

This module contains the code necessary to run Grover search on 3
qubits, both with a simulator and with a real quantum computing
device. This code is the companion for the paper
"An introduction to quantum computing, without the physics",
Giacomo Nannicini, https://arxiv.org/abs/1708.03684.

"""
def input_state(circuit, f_in, n):
    """(n+1)-qubit input state for Grover search."""
    for j in range(n):
        circuit.h(f_in[j])

q = QuantumRegister(n)
ans = ClassicalRegister(n)
qc = QuantumCircuit(q, ans)

input_state(qc, q, n)

backend = Aer.get_backend('statevector_simulator')
qc = transpile(qc, backend)
job = backend.run(qc, shots=10)
result = job.result()
state_vector = result.get_statevector(qc)
m=display( Math(r'\mathrm{state\ after\ initialization:\ }' + format_vector(state_vector)))

# apply T rounds of oracle and inversion about the average
print('number of iterations T=%d'% T)
for t in range(T):
    for i in range(n):
        qc.barrier(q[i]) # for better visualization
    qc.id(q[0])
    # Apply T full iterations
    black_box(qc, q)
    for i in range(n):
        qc.barrier(q[i])
    qc.id(q[0])
    inversion_about_average(qc, q, n)

# Measure the output register in the computational basis
for j in range(n):
    qc.measure(q[j], ans[j])

# Execute circuit
backend = BasicProvider().get_backend(name="basic_simulator")
qc = transpile(qc, backend=backend)
job = backend.run(qc, shots=10)
result = job.result()

# Get counts and plot histogram
counts = result.get_counts()
plot_histogram(counts)

qc.draw(output='mpl')

authentication = get_authentication()
QI.set_authentication(authentication, QI_URL)

QI.backends()

qi_backend = QI.get_backend('QX single-node simulator')
qc=transpile(qc, backend=qi_backend)
j=qi_backend.run(qc, shots=512)

result = j.result()
print('Generated histogram:')
print(result.get_counts())

plot_histogram(result.get_counts(qc))

