import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import sklearn
from sklearn import datasets
from qiskit import QuantumRegister, ClassicalRegister
from qiskit import QuantumRegister
from qiskit import QuantumCircuit
from qiskit import Aer, execute
from math import pi,log
from qiskit import *   
from qiskit.extensions import XGate, UnitaryGate
import tensorflow as tf
import cv2
#imported stuff from Weiwen's code
import torch
import torchvision
from torchvision import datasets
import torchvision.transforms as transforms
import torch.nn as nn
import shutil
import os
import time
import sys
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
from skcuda.linalg import svd
from matplotlib import cm
import functools
%matplotlib inline  
backend = Aer.get_backend('qasm_simulator')
print = functools.partial(print, flush=True)

interest_num = [3,6]
ori_img_size = 28 #original image is 28x28
img_size = 4 #size to downsize to (set to 28 if no downsizing)
# number of subprocesses to use for data loading
num_workers = 0
# how many samples per batch to load
batch_size = 1
inference_batch_size = 1

# Weiwen: modify the target classes starting from 0. Say, [3,6] -> [0,1]
def modify_target(target):
    for j in range(len(target)):
        for idx in range(len(interest_num)):
            if target[j] == interest_num[idx]:
                target[j] = idx
                break
    new_target = torch.zeros(target.shape[0],2)
    for i in range(target.shape[0]):        
        if target[i].item() == 0:            
            new_target[i] = torch.tensor([1,0]).clone()     
        else:
            new_target[i] = torch.tensor([0,1]).clone()      
    return target,new_target

def select_num(dataset,labels,interest_num):
    valid_labels = np.array([])
    for num in interest_num:
        if valid_labels.size > 0:
            valid_labels += labels== num
        else:
            valid_labels = labels == num
    dataset = dataset[valid_labels]
    labels = labels[valid_labels]
    return (dataset,labels)


def nomalizatied(ori,fig):
    row = fig['row']
    col = fig['col']
    
    transformed = []
    for i in range(len(ori)):
        transform = np.resize(ori[i],(row,col))
        transformed.append(transform)
    
    return transformed  

################ Weiwen on 12-30-2020 ################
# Using torch to load MNIST data
######################################################

# convert data to torch.FloatTensor
transform = transforms.Compose([transforms.Resize((ori_img_size,ori_img_size)),
                                transforms.ToTensor()])
# Path to MNIST Dataset
(train_data,train_labels),(test_data,test_labels) = tf.keras.datasets.mnist.load_data()

(train_data,train_labels) = select_num(train_data,train_labels,interest_num)
(test_data,test_labels) =  select_num(test_data,test_labels,interest_num)

#train_data = nomalizatied(train_data,{'row':4,'col':4})
#test_data = nomalizatied(test_data,{'row':4,'col':4})

train_data = train_data[:1000]
train_labels = train_labels[:1000]

train_data = torch.Tensor(train_data) # transform to torch tensor
train_labels = torch.Tensor(train_labels)
test_data = torch.Tensor(test_data)
test_labels = torch.Tensor(test_labels)

train = torch.utils.data.TensorDataset(train_data,train_labels) # create your datset
test =torch.utils.data.TensorDataset(test_data,test_labels)



u, s, v = svd(np.array(train_data[0:10]/255))

# prepare data loaders
train_loader = torch.utils.data.DataLoader(train, batch_size=batch_size,
    num_workers=num_workers, shuffle=False, drop_last=True)
test_loader = torch.utils.data.DataLoader(test, batch_size=inference_batch_size, 
    num_workers=num_workers, shuffle=False, drop_last=True)

class ToQuantumData(object):
    def __call__(self, input_vec):
        vec_len = len(input_vec)
        input_matrix = np.zeros((vec_len, vec_len))
        input_matrix[0] = input_vec
        input_matrix = input_matrix.transpose()
        u, s, v = np.linalg.svd(input_matrix)
        output_matrix = np.dot(u, v)
        output_data = output_matrix[:, 0].view()
        return output_data

class ToQuantumMatrix(object):
    def __call__(self, input_vec):
        vec_len = len(input_vec)
        input_matrix = np.zeros((vec_len, vec_len))
        input_matrix[0] = input_vec
        #print('new input vec')
        input_matrix = input_matrix.transpose()
        #print(input_vec)
        u, s, v = np.linalg.svd(input_matrix)
        output_matrix = np.dot(u, v)
        return output_matrix  


def data_pre_pro(vector):
    # Converting classical data to quantum data
    trans_to_vector = ToQuantumData()
    trans_to_matrix = ToQuantumMatrix() 
    #print("Classical Data: {}".format(vector))
    #print("Quantum Data: {}".format(trans_to_vector(vector)))
    return trans_to_matrix(vector),trans_to_vector(vector)


def PaddingNeeded(n):
    #check if n is a power of 2    
    num_decimal = np.log2(n)
    num_int = int(-(-num_decimal//1)) #round any decimal number up
    dims_needed = int(2**(num_int))
    return dims_needed

def pad_zeros(data):
    row = len(data)
    col = len(data[0])
    dims = max(row,col)
    dims_needed = PaddingNeeded(dims)
    return np.pad(data,((0,dims_needed-row),(0,dims_needed-col)),'constant', constant_values=(0,0))
    
  
from tqdm import tqdm
dataset = []
labels = []
# Use the first image from test loader as example
for batch_idx, (data, target) in tqdm(enumerate(train_loader)):
    torch.set_printoptions(threshold=sys.maxsize)
    np.set_printoptions(threshold=sys.maxsize)
    data_grid = torchvision.utils.make_grid(data) #read the data from torch
    np_data = data_grid.numpy() #convert tensor to numpy array
    image = np.asarray(np_data[0] * 255, np.uint8)    
    im = Image.fromarray(image,mode="L")
    im = im.resize((4,4),Image.BILINEAR)
    np_image = np.asarray(im)*1/255 #normalize the resized image
    data_pad = pad_zeros(np_image) #pad the data with zeros (if necessary) to make dimensionality a power of 2
    np_1d = data_pad.ravel() #convert the padded data to 1-d array for encoding function
    quantum_matrix,quantum_data = data_pre_pro(np_1d) #convert the data to unitary matrix
    dataset.append(quantum_matrix)
    labels.append(quantum_data)

    

for matrix in dataset:
    print(is_unitary(np.matrix(matrix)))

train_labels_binary = train_labels == 3

train_labels_binary = np.array(train_labels_binary).astype('int') # 3 is Label 1, 6 is label 0

#------------------------------------------------------------------------------------------
#DATA ENCODING SECTION
#------------------------------------------------------------------------------------------

def is_unitary(m):
    return np.allclose(np.eye(m.shape[0]), m.H * m)

def quantum_encode(quantum_matrix):
    dims = len(quantum_matrix[0]) #dimensionality of data
    q=int(np.log2(dims)) # Number of qubits = log2(Dimensionality of data) (could be a decimal number)
    print(dims, 'classical bits encoded onto ', q, ' qubits')  
    c=1 #one classical bit
    #Create quantum register and circuit
    inp = QuantumRegister(q,"in_qbit")
    circ = QuantumCircuit(inp)
    # Add classical register
    c_reg = ClassicalRegister(c,"reg")
    circ.add_register(c_reg)
    # From Listing 3: create auxiliary qubits
    #aux = QuantumRegister(1,"aux_qbit")
    #circ.add_register(aux)
    circ.append(UnitaryGate(quantum_matrix, label="Input"), inp[0:q])
    return circ, q

#encode the classicaL data into a quantum circuit (2^N parameters onto N qubits)
circ, q= quantum_encode(quantum_matrix)
print(circ)    
    
    
# Using StatevectorSimulator from the Aer provider
simulator = Aer.get_backend('statevector_simulator')
result = execute(circ, simulator).result()
statevector = result.get_statevector(circ)

print("Data to be encoded: \n {}\n".format(quantum_data))
print("Data read from the circuit: \n {}".format(statevector))

#All functions needed for the functionality of the circuit simulation
def generate_and_save_images(model, epoch, test_input):
  # Notice `training` is set to False.
  # This is so all layers run in inference mode (batchnorm).
  predictions = model(test_input, training=False)

  fig = plt.figure(figsize=(4,4))

  for i in range(predictions.shape[0]):
      plt.subplot(4, 4, i+1)
      dp = np.array((predictions[i] * 127.5) + 127.5).astype('uint8')
      plt.imshow(dp)
      plt.axis('off')

  #plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
  plt.show()
    

    
def ran_ang():
    #return np.pi/2
    return np.random.rand()*np.pi

def single_qubit_unitary(circ_ident,qubit_index,values):
    circ_ident.ry(values[0],qubit_index)

def dual_qubit_unitary(circ_ident,qubit_1,qubit_2,values):
    circ_ident.ryy(values[0],qubit_1,qubit_2)

def controlled_dual_qubit_unitary(circ_ident,control_qubit,act_qubit,values):
    circ_ident.cry(values[0],control_qubit,act_qubit)
    #circ_ident.cry(values[0],act_qubit,control_qubit)
    
def traditional_learning_layer(circ_ident,num_qubits,values,style="Dual",qubit_start=1,qubit_end=5):
    if style == "Dual":
        for qub in np.arange(qubit_start,qubit_end):
            single_qubit_unitary(circ_ident,qub,values[str(qub)])
        for qub in np.arange(qubit_start,qubit_end-1):
            dual_qubit_unitary(circ_ident,qub,qub+1,values[str(qub)+","+str(qub+1)])
    elif style =="Single":
        for qub in np.arange(qubit_start,qubit_end):
            single_qubit_unitary(circ_ident,qub,values[str(qub)])
    elif style=="Controlled-Dual":
        for qub in np.arange(qubit_start,qubit_end):
            single_qubit_unitary(circ_ident,qub,values[str(qub)])
        for qub in np.arange(qubit_start,qubit_end-1):
            dual_qubit_unitary(circ_ident,qub,qub+1,values[str(qub)+","+str(qub+1)])
        for qub in np.arange(qubit_start,qubit_end-1):
            controlled_dual_qubit_unitary(circ_ident,qub,qub+1,values[str(qub)+"--"+str(qub+1)])
'''
def data_loading_circuit(circ_ident,num_qubits,values,qubit_start=1,qubit_end=5):
    k = 0
    for qub in np.arange(qubit_start,qubit_end):
        circ_ident.ry(values[k],qub)
        k += 1
'''
def data_loading_circuit(circ_ident,quantum_matrix,qubit_start=1,qubit_end=5):
    circ_ident.append(UnitaryGate(quantum_matrix, label="Input"), circ_ident.qubits[qubit_start:qubit_end])

        
def quantum_encode(quantum_matrix):
    dims = len(quantum_matrix[0]) #dimensionality of data
    q=int(np.log2(dims)) # Number of qubits = log2(Dimensionality of data) (could be a decimal number)
    print(dims, 'classical bits encoded onto ', q, ' qubits')  
    c=1 #one classical bit
    #Create quantum register and circuit
    inp = QuantumRegister(q,"in_qbit")
    circ = QuantumCircuit(inp)
    # Add classical register
    c_reg = ClassicalRegister(c,"reg")
    circ.add_register(c_reg)
    # From Listing 3: create auxiliary qubits
    #aux = QuantumRegister(1,"aux_qbit")
    #circ.add_register(aux)
    circ.append(UnitaryGate(quantum_matrix, label="Input"), inp[0:q])
    return circ, q
        
def swap_test(circ_ident,num_qubits):
    num_swap = num_qubits//2
    for i in range(num_swap):
        circ_ident.cswap(0,i+1,i+num_swap+1)
    circ_ident.h(0)
    circ_ident.measure(0,0)
        
def init_random_variables(q,style):
    trainable_variables = {}
    if style=="Single":
        for i in np.arange(1,q+1):
            trainable_variables[str(i)] = [ran_ang()]
    elif style=="Dual":
        for i in np.arange(1,q+1):
            trainable_variables[str(i)] = [ran_ang()]
            if i != q:
                trainable_variables[str(i)+","+str(i+1)] = [ran_ang()]
    elif style=="Controlled-Dual":
        for i in np.arange(1,q+1):
            trainable_variables[str(i)] = [ran_ang()]
            if i != q:
                trainable_variables[str(i)+","+str(i+1)] = [ran_ang()]
                trainable_variables[str(i)+"--"+str(i+1)] = [ran_ang()]
    return trainable_variables
    
def get_probabilities(circ,counts=5000):
    job = execute(circ, backend, shots=counts)
    results = job.result().get_counts(circ)
    try:
        prob = results['0']/(results['1']+results['0'])
        prob = (prob-0.5)
        if prob <= 0.005:
            prob = 0.005
        else:
            prob = prob*2
    except:
        prob = 1
    return prob
        
# Define loss function. SWAP Test returns probability, so minmax probability is logical
def cost_function(p,yreal,trimming):
    if yreal == 0:
        return -np.log(p)
        #return 1-p
    elif yreal == 1:
        return -np.log(1-p)
        #return p
    
def generator_cost_function(p):
    return -np.log(p)

def update_weights(init_value,lr,grad):
    while lr*grad > 2*np.pi:
        lr /= 10
        print("Warning - Gradient taking steps that are very large. Drop learning rate")
    weight_update = lr*grad
    new_value = init_value
    print("Updating with a new value of " + str(weight_update))
    if new_value-weight_update > 2*np.pi:
        new_value = (new_value-weight_update) - 2*np.pi
    elif new_value-weight_update < 0:
        new_value = (new_value-weight_update) + 2*np.pi
    else:
        new_value = new_value - weight_update
    return new_value 
# Define loss function. SWAP Test returns probability, so minmax probability is logical
def cost_function(p,yreal,trimming):
    if yreal == 0:
        return -np.log(p)
        #return 1-p
    elif yreal == 1:
        return -np.log(1-p)
        #return p
    
def generator_cost_function(p):
    return -np.log(p)

def update_weights(init_value,lr,grad):
    while lr*grad > 2*np.pi:
        lr /= 10
        print("Warning - Gradient taking steps that are very large. Drop learning rate")
    weight_update = lr*grad
    new_value = init_value
    print("Updating with a new value of " + str(weight_update))
    if new_value-weight_update > 2*np.pi:
        new_value = (new_value-weight_update) - 2*np.pi
    elif new_value-weight_update < 0:
        new_value = (new_value-weight_update) + 2*np.pi
    else:
        new_value = new_value - weight_update
    return new_value 



QuantumCircuit(5,1).qubits

# ------------------------------------------------------------------------------------
# We treat the first n qubits are the discriminators state. n is always defined as the
# integer division floor of the qubit count.
# This is due to the fact that a state will always be k qubits, therefore the 
# number of total qubits must be 2k+1. 2k as we need k for the disc, and k to represent
# either the other learned quantum state, or k to represent a data point
# then +1 to perform the SWAP test. Therefore, we know that we will always end up
# with an odd number of qubits. We take the floor to solve for k. 1st k represents 
# disc, 2nd k represents the "loaded" state be it gen or real data
# ------------------------------------------------------------------------------------
# Use different function calls to represent training a GENERATOR or training a DISCRIMINATOR
# ------------------------------------------------------------------------------------
# THIS SECTION IS FOR THE ONLINE GENERATION OF QUANTUM CIRCUITS

def disc_fake_training_circuit(trainable_variables,key,key_value,diff=False,fwd_diff = False,Sample=False):
    if Sample:
        z = q//2
        circ = QuantumCircuit(q,z)
    else:
        circ = QuantumCircuit(q,c)
    circ.h(0)
    if diff == True and fwd_diff == True:
        trainable_variables[key][key_value] += par_shift
    if diff == True and fwd_diff == False:
        trainable_variables[key][key_value] -= par_shift
    traditional_learning_layer(circ,q,trainable_variables,style=layer_style,qubit_start=1,qubit_end=q//2 +1)
    traditional_learning_layer(circ,q,trainable_variables,style=layer_style,qubit_start=q//2 +1,qubit_end=q)
    if Sample:
        for qub in range(q//2):
            circ.measure(q//2 + 1 + qub,qub)
    else:
        swap_test(circ,q)
    if diff == True and fwd_diff == True:
        trainable_variables[key][key_value] -= par_shift
    if diff == True and fwd_diff == False:
        trainable_variables[key][key_value] += par_shift
    return circ

def disc_real_training_circuit(training_variables,data,key,key_value,diff,fwd_diff):
    circ = QuantumCircuit(q,c)
    circ.h(0)
    if diff == True & fwd_diff == True:
        training_variables[key][key_value] += par_shift
    if diff == True & fwd_diff == False:
        training_variables[key][key_value] -= par_shift
    traditional_learning_layer(circ,q,training_variables,style=layer_style,qubit_start=1,qubit_end=q//2 +1)
    data_loading_circuit(circ,data,qubit_start=q//2+1,qubit_end=q)
    if diff == True & fwd_diff == True:
        training_variables[key][key_value] -= par_shift
    if diff == True & fwd_diff == False:
        training_variables[key][key_value] += par_shift
    swap_test(circ,q)
    return circ

def generate_kl_divergence_hist(actual_data, epoch_results_data):
    plt.clf() # clears current figure
    sns.set()
    kl_div_vec = []
    for kl_dim in range(actual_data.shape[1]):
        kl_div = kl_divergence(actual_data[:,kl_dim],epoch_results_data[:,kl_dim])
        kl_div_vec.append(kl_div)
    return kl_div_vec

def bin_data(dataset):
    bins = np.zeros(10)
    for point in dataset:
        indx = int(str(point).split('.')[-1][0]) # The shittest way imaginable to extract the first val aft decimal
        bins[indx] +=1  
    bins /= sum(bins)
    return bins

def kl_divergence(p_dist, q_dist):
    p = bin_data(p_dist)
    q = bin_data(q_dist)
    kldiv = 0
    for p_point,q_point in zip(p,q):
        kldiv += (np.sqrt(p_point) - np.sqrt(q_point))**2
    kldiv = (1/np.sqrt(2))*kldiv**0.5 
    return kldiv
    #return sum(p[i] * log2(p[i]/q[i]) for i in range(len(p)))  # ?... are we confident in this... 

    
# Checkpointing code
def save_variables(train_var,epoch,number):
    with open(f"Epoch-{epoch}-Variables-numbers-{number}",'w') as file:
        file.write(str(train_var))
        
def load_variables(epoch,number):
    with open(f"Epoch-{epoch}-Variables-numbers-{number}",'r') as file:
        texts = file.read()
    return eval(texts)

train_var_1 = load_variables(11,6)
train_var_0 = load_variables(49,3)

train_var_1

z = disc_real_training_circuit(train_var,quantum_matrix,'1',0,diff=False,fwd_diff=False)

#z.append(UnitaryGate(quantum_matrix, label="Input"), z.qubits[10:20])

quantum_matrix

dataset = np.array(dataset)
labels = np.array(labels)

train_labels_binary = train_labels_binary[:dataset.shape[0]]

dataset = dataset[:5000]
train_labels_binary = train_labels_binary[:5000]

  
q=9# Number of qubits = Dimensionality of data = round up to even number = num qubits
c=1
circ = QuantumCircuit(q,c)
circ.h(0)
layer_style = "Controlled-Dual"
train_var_0 = init_random_variables(q//2,layer_style)
train_var_1 = init_random_variables(q//2,layer_style)
#train_var_2 = init_random_variables(q//2,layer_style)

tracked_d_loss = []
tracked_d_loss1 = []
gradients = []
learning_rate=0.015
train_iter = 50
corr = 0
wrong= 0 
loss_d_to_real = 0
print('Starting Training')
print('-'*20)
print("train_var_0 training")

for epoch in range(train_iter):
    start = time.time()
    loss = [0,0]
    par_shift = 0.5*np.pi/((1+epoch)**0.5)
    for index,point in enumerate(dataset[train_labels_binary==0]):
        for key,value in train_var_0.items():
            if str(q//2 + 1 ) in key:
                break
            for key_value in range(len(value)):
                forward_diff = -np.log(get_probabilities(disc_real_training_circuit(train_var_0,point,key,key_value,diff=True,fwd_diff=True)))
                backward_diff = -np.log(get_probabilities(disc_real_training_circuit(train_var_0,point,key,key_value,diff=True,fwd_diff=False)))
                df = 0.5*(forward_diff-backward_diff)
                train_var_0[key][key_value] -= df*learning_rate
                #print(f"Updated Variable {key}-({key_value}) by {-df*learning_rate}")
    print('Time for Epoch {} is {} sec'.format(epoch+1,time.time()-start))
    print("-"*20)
    save_variables(train_var_0,epoch,3)

for epoch in range(train_iter):
    start = time.time()
    loss = [0,0]
    par_shift = 0.5*np.pi/((1+epoch)**0.5)
    for index,point in enumerate(dataset[train_labels_binary==1]):
        for key,value in train_var_1.items():
            if str(q//2 + 1 ) in key:
                break
            for key_value in range(len(value)):
                forward_diff = -np.log(get_probabilities(disc_real_training_circuit(train_var_1,point,key,key_value,diff=True,fwd_diff=True)))
                backward_diff = -np.log(get_probabilities(disc_real_training_circuit(train_var_1,point,key,key_value,diff=True,fwd_diff=False)))
                df = 0.5*(forward_diff-backward_diff)
                train_var_1[key][key_value] -= df*learning_rate
    print('Time for Epoch {} is {} sec'.format(epoch+1,time.time()-start))
    print("-"*20)
    save_variables(train_var_1,epoch,6)


def is_unitary(m):
    return np.allclose(np.eye(m.shape[0]), m.H * m)

is_unitary(np.matrix(quantum_matrix[0]))

corr = 0
wron = 0
for index,point in enumerate(dataset[train_labels_binary==1]):
    p1 = get_probabilities(disc_real_training_circuit(train_var_0,point,key,key_value,diff=False,fwd_diff=False))
    p2 = get_probabilities(disc_real_training_circuit(train_var_1,point,key,key_value,diff=False,fwd_diff=False))
    if p1<p2:
        corr +=1
    else:
        wron+=1
        
acc = corr/(corr+wron)

print(acc)

corr

wron

829