import random
import numpy  as np

from qiskit import QuantumCircuit, AncillaRegister, QuantumRegister, execute, Aer

def normalize_state(a = 1, b = 0):
    
    norm = np.sqrt(a * np.conjugate(a) + b * np.conjugate(b))
    
    return [a / norm, b / norm]

def modulus_and_phase(a):
    
    return (np.absolute(a), np.angle(a))

def int_to_binary(value, n_bits = 17, reverse = True):
    
    value = int(value)
    if value != 0:
        bound = int(np.ceil(np.log2(value)))
        if n_bits < bound: n_bits = bound
    
    binary_format = '{0:0' + str(n_bits) + 'b}'
    binary = binary_format.format(value)
    
    if reverse: return binary[ : : -1]
    else: return binary

def binary_to_int(value, reverse = True):
    
    if reverse: value = value[ : : -1]

    return int(value, 2)

def get_state(circuit, precision = 10, fancy_print = False):
    
    # Let's simulate our circuit in order to get the final state vector!
    svsim = Aer.get_backend('statevector_simulator')

    # Do the simulation, return the result and get the state vector
    result = execute(circuit, svsim).result().get_statevector()
    
    if int(precision) < 8: precision = 8
    
    result = result.round(precision)
    
    # Returns non-zero values
    state_indices = result.nonzero()[0]
    
    states = np.empty(0)
    for index in state_indices:
        binary_state = int_to_binary(index)
        states = np.append(states , [result[index], binary_state])
    
    if fancy_print:
        string = ''
        for index in range(0, len(states), 2):
            coef, bstate = states[index], states[index + 1]
            string = string + f'{coef} |{bstate}>   '
        print(string)
    
    return states

def random_noise(gate_label = 'Noise', reveal_error = False):
    
    # Our circuit has 9-qubits
    gate_circuit = QuantumCircuit(9, name = gate_label)
    
    # Define randomly the number of errors to apply from 0 (no errors) up to 2
    number_of_errors  = random.choice(range(3))
    types_of_errors   = ['X', 'Y', 'Z']
    qubits_wo_errors  = list(range(9))
    qpacks_wo_errors  = list(range(9))
    
    global gate_reveal_label
    gate_reveal_label = 'I' * 9
    
    
    def str_assignment(string, index, char):
        
        string = list(string)
        string[index] = str(char)
        string = ''.join(string)
        
        return string

    
    while number_of_errors > 0:
        qubit_assigned = random.choice(qpacks_wo_errors)
        type_error     = random.choice(types_of_errors)
        
        number_of_errors -= 1
        qubits_wo_errors.remove(qubit_assigned)

        gate_reveal_label = str_assignment(gate_reveal_label, qubit_assigned, type_error)

        # Remove 3-qubits packs after assign error to qubit in it
        if qubit_assigned in list(range(3)):
            qpacks_wo_errors = [elem for elem in qpacks_wo_errors if elem > 2]
        elif qubit_assigned in list(range(3, 6)):
            qpacks_wo_errors = [elem for elem in qpacks_wo_errors if (elem < 3 or elem > 5)]
        else:
            qpacks_wo_errors = [elem for elem in qpacks_wo_errors if elem < 6]

        # Can't repeat quantum error gate applied
        if type_error == 'Y':
            gate_circuit.y(qubit_assigned)
            break
        elif type_error == 'X':
            gate_circuit.x(qubit_assigned)
            try: # Maybe we don't have this errors to remove
                types_of_errors.remove('X')
                types_of_errors.remove('Y')
            except:
                pass
        else:
            gate_circuit.z(qubit_assigned)
            try: # Maybe we don't have this errors to remove
                types_of_errors.remove('Y')
                types_of_errors.remove('Z')
            except:
                pass
    
    # Assign identity gate for the rest of qubits
    for qubit in qubits_wo_errors:
        gate_circuit.i(qubit)
    
    gate = gate_circuit.to_gate()
    
    if reveal_error: gate.label = 'Noise:\n' + gate_reveal_label
    
    return gate

def check_error_to_ancilla(state, error = None):
    
    if error is None: error = gate_reveal_label
    
    binary = state[1]
    
    error_types    = ['X0', 'X1', 'X2', 'Z ']
    state_ancillas = [binary[3 : 5], binary[8 : 10], binary[13 : 15], binary[-2 : ]]
    
    ancillas = {}
    for key, anc in zip(error_types, state_ancillas): ancillas[key] = anc  
    
    first_pack  = list(range(0, 7, 3))
    second_pack = list(range(1, 8, 3))
    third_pack  = list(range(2, 9, 3))
    
    
    def ancilla_label(value):
        
        if value == 0: return '10'
        elif value == 1: return '11'
        else: return '01'
    
    
    expected = {}
    for key in error_types: expected[key] = '00'
    
    for qubit, gate in enumerate(list(error)):
        
        n_pack = qubit // 3
        key    = 'X' + str(n_pack)
        n_xerr = qubit % 3
        
        if gate == 'I': continue
        elif gate == 'X':
            if qubit in first_pack:
                expected[key] = ancilla_label(n_xerr)
            elif qubit in second_pack:
                expected[key] = ancilla_label(n_xerr)
            else:
                expected[key] = ancilla_label(n_xerr)
        elif gate == 'Z':
            expected['Z '] = ancilla_label(n_pack)
        else:
            expected['Z '] = ancilla_label(n_pack)
            if qubit in first_pack:
                expected[key] = ancilla_label(n_xerr)
            elif qubit in second_pack:
                expected[key] = ancilla_label(n_xerr)
            else:
                expected[key] = ancilla_label(n_xerr)
    
    print('Error  Anc. expected  Anc. state  Equal?\n' + '-' * 40)
    for key in error_types:
        print(f' {key}     {expected[key]}             {ancillas[key]}          {expected[key] == ancillas[key]}')

def cnotnot(gate_label = 'CNOTNOT'):
    
    gate_circuit = QuantumCircuit(3, name = gate_label)
    gate_circuit.cnot(0, 1)
    gate_circuit.cnot(0, 2)
    
    gate = gate_circuit.to_gate()
    
    return gate

def czz(gate_label = 'CZZ'):
    
    gate_circuit = QuantumCircuit(3, name = gate_label)
    gate_circuit.cz(0, 1)
    gate_circuit.cz(0, 2)
    
    gate = gate_circuit.to_gate()
    
    return gate

def invccnot(gate_label = 'invCCNOT'):
    
    gate_circuit = QuantumCircuit(3, name = gate_label)
    gate_circuit.x(0)
    gate_circuit.ccx(0, 1, 2)
    gate_circuit.x(0)
    
    gate = gate_circuit.to_gate()
    
    return gate

def c6x(gate_label = 'C6X'):
    
    gate_circuit = QuantumCircuit(7, name = gate_label)
    for index in range(1, 7):
        gate_circuit.cnot(0, index)
    
    gate = gate_circuit.to_gate()
    
    return gate

def invccz(gate_label = 'invCCZ'):
    
    gate_circuit = QuantumCircuit(3, name = gate_label)
    gate_circuit.x(0)
    gate_circuit.h(2)
    gate_circuit.ccx(0, 1, 2)
    gate_circuit.h(2)
    gate_circuit.x(0)
    
    gate = gate_circuit.to_gate()
    
    return gate

def ccz(gate_label = 'CCZ'):
    
    gate_circuit = QuantumCircuit(3, name = gate_label)
    gate_circuit.h(2)
    gate_circuit.ccx(0, 1, 2)
    gate_circuit.h(2)
    
    gate = gate_circuit.to_gate()
    
    return gate

