import numpy as np
from matplotlib import pyplot as plt
import os
import cv2
from qiskit import *

def rgb2gray(rgb):
    return np.dot(rgb[...,:3], [0.2989, 0.5870, 0.1140])

def get_Coeff(Image, watermark=None):
    """ Function that takes in an Greyscale Image matrix of any size and and optional
        watermark image and shape and returns qubit coefficients in the computational
        basis. 
        TO DO: build in a check that ensures the sizes of the Image and watermark are the same
    """
    
    Flatten_Image = np.reshape(Image, np.shape(Image)[0]*np.shape(Image)[1])
    
    # Get the theta angles of the Image and watermark from the greyscale values
    Theta_I = 2*np.arccos(np.sqrt(Flatten_Image/255))
                          
    if watermark is None:
        Theta_W = np.zeros(len(Theta_I))
    else:  
        Flatten_watermark = np.reshape(watermark, np.shape(watermark)[0]*np.shape(watermark)[1])
        Theta_W = 2*np.arccos(np.sqrt(Flatten_watermark/255))
        
        
    # Compute the qubit Coeff and return them in a normalized qubit state vector. 
    Coeff = np.stack((np.cos(Theta_I/2)*np.cos(Theta_W/2),
                       np.cos(Theta_I/2)*np.sin(Theta_W/2),
                       np.sin(Theta_I/2)*np.cos(Theta_W/2),
                       np.sin(Theta_I/2)*np.sin(Theta_W/2)))
    
    
    Coeff = np.ndarray.flatten(np.transpose(Coeff))
    
    return Coeff/np.linalg.norm(Coeff)

def initialise_state(desired_initial_state):
    """ Returns an initialized quantum circuit for a given input state"""
    n = int(round(np.log2(desired_initial_state.size))) #number of qubits
    
    qc_init = QuantumCircuit(n)

    qc_init.initialize(desired_initial_state, range(n))
    return qc_init

def get_Images(probability_vector):
    """ Takes in an probability vector obtained from measuring all the qubits in the quantum circuit and returns
        the corresponding matrices of the greyscale values of the measured images.
        TO DO: vectorize this function
    """
    
    n = int(np.round(len(probability_vector)/4)) # get sizes back of classical images
    c_Image = np.zeros(n)
    w_Image = np.zeros(n)
    
    # Sum and normalize the probabilities to get the greyscale pixel value out
    for i in range(n):
        c_Image[i] = (probability_vector[4*i]+probability_vector[4*i+1])*255/np.sum(probability_vector[4*i:(4*i)+4])
        w_Image[i] = (probability_vector[4*i]+probability_vector[4*i+2])*255/np.sum(probability_vector[4*i:(4*i)+4])
    
    c_Image = np.reshape(c_Image,[np.int(np.round(np.sqrt(n))),np.int(np.round(np.sqrt(n)))])
    w_Image = np.reshape(w_Image,[np.int(np.round(np.sqrt(n))),np.int(np.round(np.sqrt(n)))])
    
    return c_Image, w_Image

N = 16 #size in pixels


# Carrier Image
img_file = os.path.expanduser("cat.png")
img = cv2.imread(img_file)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
res = cv2.resize(img, dsize=(N, N), interpolation=cv2.INTER_CUBIC)
C_Image = rgb2gray(res)

# Watermark
img_file2 = os.path.expanduser("IBM-logo.png")
img2 = cv2.imread(img_file2)
img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
res2 = cv2.resize(img2, dsize=(N, N), interpolation=cv2.INTER_CUBIC)
W_Image = rgb2gray(res2)

plt.rcParams['figure.figsize'] = [9, 7]
plt.rcParams['figure.dpi'] = 100

fig, axs = plt.subplots(1,2)

axs[0].imshow(C_Image/255, cmap=plt.get_cmap('gray'), vmin=0, vmax=1)
axs[0].title.set_text('Carrier Image')

axs[1].imshow(W_Image/255, cmap=plt.get_cmap('gray'), vmin=0, vmax=1)
axs[1].title.set_text('Watermark Image')

State_vector = get_Coeff(C_Image, W_Image)
qc_init = initialise_state(State_vector)
#qc_init.draw('mpl') 

n = int(round(np.log2(State_vector.size))) # number of qubits
shots0 = 100000

# Create a Quantum Circuit
meas = QuantumCircuit(n, n)
meas.barrier(range(n))

# map the quantum measurement to the classical bits
meas.measure(range(n), range(n))

# The Qiskit circuit object supports composition using
# the addition operator.
qc = qc_init + meas

#drawing the circuit
#qc.draw('mpl')

# Use Aer's qasm_simulator
backend_sim = Aer.get_backend('qasm_simulator')

# Execute the circuit on the qasm simulator.
job_sim = execute(qc, backend_sim, shots=shots0)

# Grab the results from the job.
result_sim = job_sim.result()

#from qiskit import IBMQ
#IBMQ.load_account()
#provider = IBMQ.get_provider("ibm-q")
#backend = provider.get_backend("ibmq_essex")
#job = q.execute(qc, backend=backend, shots=shots0)
#job_monitor(job)

counts = result_sim.get_counts(qc)
probability_vector = np.zeros(2**n)
int_counts = counts.int_outcomes()

for k in range(2**n):
    probability_vector[k] = int_counts.get(k, 0)/shots0

RC_Image, RW_Image = get_Images(probability_vector)

fig, axs = plt.subplots(2,2)

axs[0,0].imshow(RC_Image/255, cmap=plt.get_cmap('gray'), vmin=0, vmax=1)
axs[0,0].title.set_text('Retrieved Carrier Image')

axs[0,1].imshow(C_Image/255, cmap=plt.get_cmap('gray'), vmin=0, vmax=1)
axs[0,1].title.set_text('Original Carrier Image')

axs[1,0].imshow(RW_Image/255, cmap=plt.get_cmap('gray'), vmin=0, vmax=1)
axs[1,0].title.set_text('Retrieved Watermark Image')

axs[1,1].imshow(W_Image/255, cmap=plt.get_cmap('gray'), vmin=0, vmax=1)
axs[1,1].title.set_text('Original Watermark Image')

plt.show()

n = int(round(np.log2(State_vector.size)))
shots0 = 100000

# Create the Quantum Scramble Circuit with controlled U3 rotations 
meas = QuantumCircuit(n, n)
meas.barrier(range(n))
meas.cu3(np.pi/4, 0, 0, 3, 0)
meas.cu3(-np.pi/4, 0, 0, 4, 0) 
meas.cu3(-np.pi/2, 0, 0, 5, 0) 
meas.cu3(-np.pi/3, 0, 0, 6, 0) 
meas.cu3(np.pi/2, 0, 0, 7, 0) 

# map the quantum measurement to the classical bits
meas.measure(range(n), range(n))

# The Qiskit circuit object supports composition using
# the addition operator.
qc = qc_init + meas

# Optional uncomment to draw the circuit
#qc.draw('mpl')

# Use Aer's qasm_simulator
backend_sim = Aer.get_backend('qasm_simulator')

# Execute the circuit on the qasm simulator.
# We've set the number of repeats of the circuit
# to be 1024, which is the default.
job_sim = execute(qc, backend_sim, shots=shots0)

# Grab the results from the job.
result_sim = job_sim.result()

counts = result_sim.get_counts(qc)
probability_vector = np.zeros(2**n)
int_counts = counts.int_outcomes()

for k in range(2**n):
    probability_vector[k] = int_counts.get(k, 0)/shots0

RRC_Image, SCrW_Image = get_Images(probability_vector)

fig, axs = plt.subplots(2,2)

axs[0,0].imshow(RRC_Image/255, cmap=plt.get_cmap('gray'), vmin=0, vmax=1)
axs[0,0].title.set_text('Retrieved Carrier Image')

axs[0,1].imshow(C_Image/255, cmap=plt.get_cmap('gray'), vmin=0, vmax=1)
axs[0,1].title.set_text('Original Carrier Image')


axs[1,0].imshow(SCrW_Image/255, cmap=plt.get_cmap('gray'), vmin=0, vmax=1)
axs[1,0].title.set_text('Scrambled Watermark Image')

axs[1,1].imshow(W_Image/255, cmap=plt.get_cmap('gray'), vmin=0, vmax=1)
axs[1,1].title.set_text('Original Watermark Image')

plt.show()