import numpy as np
# Importing standard Qiskit libraries
from qiskit import QuantumCircuit, transpile, Aer, IBMQ
from qiskit.tools.jupyter import *
from qiskit.visualization import *
from ibm_quantum_widgets import *

# Loading your IBM Quantum account(s)
provider = IBMQ.load_account()

import qiskit

from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister
from qiskit import execute
from qiskit import Aer
from qiskit import IBMQ
from qiskit.compiler import transpile
from time import perf_counter
from qiskit.tools.visualization import plot_histogram

import numpy as np
import math

'''
params
---------------
picture: square 2d array of integers representing grayscale values
assume the length (n) is some power of 2

return
---------------
a flattened representation of picture using bitstrings (boolean arrays)
'''
def convert_to_bits (picture):
    n = len(picture)
    ret = []
    for i in range(n):
        for j in range(n):
            value = picture[i][j]
            bitstring = bin(value)[2:]
            ret.append([0 for i in range(8 - len(bitstring))] + [1 if c=='1' else 0 for c in bitstring])
    return ret

'''
params
----------------
bitStr: a representation of an image using bitstrings to represent grayscale values

return
----------------
A quantum circuit containing the NEQR representation of the image
'''
def neqr(bitStr): 
    newBitStr = bitStr

    #print(newBitStr)
    #print("\n")

    # Pixel position
    idx = QuantumRegister(math.ceil(math.log2(len(newBitStr))), 'idx')

    # Pixel intensity values
    intensity = QuantumRegister(8, 'intensity')

    # Classical Register
    creg = ClassicalRegister(10, 'creg')

    # Quantum Image Representation as a quantum circuit
    # with Pixel Position and Intensity registers
    quantumImage = QuantumCircuit(intensity, idx, creg)

    numOfQubits = quantumImage.num_qubits
    print("\n>> Initial Number of Qubits:", numOfQubits)

    # -----------------------------------
    # Drawing the Quantum Circuit
    # -----------------------------------
    lengthIntensity = intensity.size
    lengthIdx = idx.size

    start = perf_counter()

    quantumImage.i([intensity[lengthIntensity-1-i] for i in range(lengthIntensity)])
    quantumImage.h([idx[lengthIdx-1-i] for i in range(lengthIdx)])

    numOfPixels = len(newBitStr)

    for i in range(numOfPixels):
        bin_ind = bin(i)[2:]
        bin_ind = (lengthIdx - len(bin_ind)) * '0' + bin_ind
        bin_ind = bin_ind[::-1]

        # X-gate (enabling zero-controlled nature)
        for j in range(len(bin_ind)):
            if bin_ind[j] == '0':
                quantumImage.x(idx[j])

        # Configuring Multi-Qubit Controlled-NOT (mcx) gates with control and target qubits
        for j in range(len(newBitStr[i])):
            if newBitStr[i][j] == 1:
                quantumImage.mcx(idx, intensity[lengthIntensity-1-j])
        
        # X-gate (reversing the Negated state of the qubits)
        for j in range(len(bin_ind)):
            if bin_ind[j] == '0':
                quantumImage.x(idx[j])
        
        quantumImage.barrier()

    quantumImage.measure(range(10), range(10))

    end = perf_counter()
    print(f">> Circuit construction took {(end-start)} seconds.")

    return (quantumImage, intensity)

if __name__ == '__main__':
    test_picture_2x2 = [[0, 100], [200, 255]]
    test_picture_3x3 = [[25, 50, 75], [100, 125, 150], [175, 200, 225]]
    test_picture_4x4 = [[0, 100, 143, 83], [200, 255, 43, 22], [12, 234, 23, 5], [112, 113, 117, 125]]
    test_picture_5x5 = [[0, 100, 212, 12, 32], [0, 100, 212, 12, 32], [0, 100, 212, 12, 32], [0, 100, 212, 12, 32], [0, 100, 212, 12, 32]]
    
    arr = convert_to_bits(test_picture_2x2)
    arr1 = convert_to_bits(test_picture_3x3)
    arr2 = convert_to_bits(test_picture_4x4)
    arr3 = convert_to_bits(test_picture_5x5)
    
    print("2x2: ", arr, "\n") 
    #print("3x3: ", arr1, "\n") 
    #print("4x4: ", arr2, "\n") 
    #print("5x5: ", arr3, "\n") 
    
    qc_image, _ = neqr(arr)

qc_image.draw()

# Circuit Dimensions  
print('>> Circuit dimensions')
print('>> Circuit depth (length of critical path): ', qc_image.depth())
print('>> Total Number of Gate Operations: ', qc_image.size())

# Get the number of qubits needed to run the circuit
active_qubits = {}
for op in qc_image.data:
    if op[0].name != "barrier" and op[0].name != "snapshot":
        for qubit in op[1]:
            active_qubits[qubit.index] = True
print(f">> Width: {len(active_qubits)} qubits")
print(f">> Width (total number of qubits and clbits): {qc_image.width()} qubits")

print(f">> Gates used: {qc_image.count_ops()}")
print(f">> Fundamental Gates: {qc_image.decompose().count_ops()}")

# Testing qc_image circuit using Fake Simulator 
# to simulate the error in a real quantum computer
from qiskit.providers.aer import AerSimulator
from qiskit.test.mock import FakeMontreal
device_backend = FakeMontreal()
sim_montreal = AerSimulator.from_backend(device_backend)

#### TRANSPILING
start_transpile = perf_counter()
tcirc = transpile(qc_image, sim_montreal, optimization_level=1) # Optimization Level = 1 (default)
end_transpile = perf_counter()
print(f"\n>> Circuit transpilation took {(end_transpile-start_transpile)} seconds.")

#### SIMULATING
start_sim = perf_counter()
# Run transpiled circuit on FakeMontreal 
result_noise = sim_montreal.run(tcirc).result()
# Get counts for each possible result 
counts_noise = result_noise.get_counts()
end_sim = perf_counter()
print(f"\n>> Circuit simulation took {(end_sim-start_sim)} seconds.")

#plot_histogram(counts_noise, figsize=(30, 15))
#print(counts_noise)

from qiskit.compiler import assemble
aer_sim = Aer.get_backend('aer_simulator')
t_qc_image = transpile(qc_image, aer_sim)

shots = 1024

qobj = assemble(t_qc_image, shots=shots)
job_neqr = aer_sim.run(qobj)
result_neqr = job_neqr.result()
counts_neqr = result_neqr.get_counts()
print('Encoded: 00 = ', arr[0])
print('Encoded: 01 = ', arr[1])
print('Encoded: 10 = ', arr[2])
print('Encoded: 11 = ', arr[3])

print("\nExpected Counts per pixel = ", shots/len(arr))
print("Experimental Counts", counts_neqr)
print("\nExpected probability for each pixel = ", 1.0/(len(arr)))
plot_histogram(counts_neqr)

machine = "ibmq_santiago"
backend = provider.get_backend(machine)
santiago_sim = AerSimulator.from_backend(backend)
start_sim = perf_counter()
result = santiago_sim.run(qc_image).result()
counts = result.get_counts(qc_image)
end_sim = perf_counter()
print("\n", end_sim-start_sim)
print("\nExpected Counts per pixel = ", shots/len(arr), "\n")
for(measured_state, count) in counts.items():
    big_endian_state = measured_state[::-1]
    print(f"Measured {big_endian_state} {count} times.")

machine = "ibmq_belem"
backend = provider.get_backend(machine)
santiago_sim = AerSimulator.from_backend(backend)
start_sim = perf_counter()
result = santiago_sim.run(qc_image).result()
counts = result.get_counts(qc_image)
end_sim = perf_counter()
print(end_sim-start_sim)
print("\nExpected Counts per pixel = ", shots/len(arr), "\n")
for(measured_state, count) in counts.items():
    big_endian_state = measured_state[::-1]
    print(f"Measured {big_endian_state} {count} times.")

# Testing qc_image circuit using Fake Simulator 
# to simulate the error in a real quantum computer
from qiskit.providers.aer import AerSimulator
from qiskit.test.mock import FakeMumbai
device_backend = FakeMumbai()
sim_brooklyn = AerSimulator.from_backend(device_backend)

#### TRANSPILING
start_transpile = perf_counter()
tcirc = transpile(qc_image, sim_brooklyn, optimization_level=1) # Optimization Level = 1 (default)
end_transpile = perf_counter()
print(f"\n>> Circuit transpilation took {(end_transpile-start_transpile)} seconds.")

#### SIMULATING
start_sim = perf_counter()
# Run transpiled circuit on FakeMontreal 
result_noise = sim_brooklyn.run(tcirc).result()
# Get counts for each possible result 
counts_noise = result_noise.get_counts()
end_sim = perf_counter()
print(f"\n>> Circuit simulation took {(end_sim-start_sim)} seconds.")

from qiskit.test.mock import FakeMontreal
fake_athens = FakeMontreal()
t_qc = transpile(qc_image, fake_athens, optimization_level=3)
qobj = assemble(t_qc, shots=4096)
result = fake_athens.run(qobj).result()
counts = result.get_counts(qc_image)
print(counts)
plot_histogram(counts)

plot_histogram(counts, figsize=(30, 15))

