from qiskit import *
import numpy as np
from functools import partial
from qiskit.aqua.components.optimizers import COBYLA
from scipy.optimize import minimize, fmin_slsqp
import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline

theta = np.array([0.5, np.pi*2, 1.4, 1/3, np.pi/2, 1.3, 1.4, 0.22, 1.2, 0.99, 1.0, 0.23, 0.14, 0.90, np.pi])

def run_quantum_circ(n, theta):
    '''
    This creates a quantum circuit with n qubits. 
    '''
    qr = QuantumRegister(n)
    cr = ClassicalRegister(n)
    circuit = QuantumCircuit(qr, cr)
    
    #Applies Hadamard Gate to first 
    circuit.h(0)
    
    #Entangles the qubits
    for i in range(n-1):
        circuit.cx(i, i+1)
        
    # For ease of understanding 
    circuit.barrier()
    
    # Applying U3 gates -- later
    i = 0
    for i in range(n):
        circuit.u3(theta[3*i], theta[(3*i)+1], theta[(3*i)+2], i)
    
    circuit.barrier()
    
    # Unentangling qubits 
    for i in reversed(range(n)):
        if i-1==-1:
            break
        else:
            circuit.cx(i-1, i)
            
    circuit.h(0)   
    circuit.barrier()
            
    # Measuring the values 
    for i in range(n-1):
        circuit.measure(i, i)
    
    backend = Aer.get_backend('qasm_simulator')
    result = execute(circuit, backend=backend, shots=10000).result()
    return result.get_counts()

theta_c = np.array([np.pi*2, np.pi/2, 0.22, 1.0, 0.90])

def run_classical_circ(n, theta):
    '''
    This creates a quantum circuit with n qubits. 
    '''
    qr = QuantumRegister(n)
    cr = ClassicalRegister(n)
    circuit = QuantumCircuit(qr, cr)
    
        
    # For ease of understanding 
    circuit.barrier()
    
    # Applying U3 gates -- later
    i = 0
    for i in range(n):
        circuit.u2(theta[i], 0, i)
      
    circuit.barrier()
            
    # Measuring the values 
    for i in range(n-1):
        circuit.measure(i, i)
    
    backend = Aer.get_backend('qasm_simulator')
    result = execute(circuit, backend=backend, shots=10000).result()
    return result.get_counts()

# Translate from dictionary to sorted values. 
def dic_to_vec(dic):
    '''
    Converts to dictionary outputted by the quantum/classical circuit to a 2^n dimensional (dim = 32) vector. 
    '''
    scores = np.zeros(32)
    num_shots = 10000
    for key in dic.keys():
        i = int(key,2)
        scores[i] = dic[key]/num_shots
    return scores

# The Payoff matrix -- made it here at random
v1 = np.random.rand(5,1)*np.random.rand(1)
v2 = np.random.rand(1,32)
payoff_mat = np.matmul(v1,v2)

def payoff_finder(result, payoff):
    '''
    Input: results from quantum/classical calc (32x1), payoff matrix.  
    Output: payoff matrix for that particular 
    '''
    return np.matmul(payoff, result)

def quantum_cost_function(n, rand_payoff_mat, theta):
    '''
    Input: number of qubits, initial values of theta, the payoff matrix you've generated
    Output: a singular number that represents how large the difference between the standard deviation
            and L1 norm. 
    '''
    
    # Runs the quantum circuit and returns a dictionary of values. 
    dic = run_quantum_circ(n, theta)
    vec = dic_to_vec(dic)
    payoff_vec = payoff_finder(vec, rand_payoff_mat)
    val = np.linalg.norm(payoff_vec, ord = 1) - np.std(payoff_vec)
    return -val

#quantum_cost_function(5,theta,payoff_mat)

def classical_cost_function(n, rand_payoff_mat, theta):
    '''
    Input: number of qubits, initial values of theta, the payoff matrix you've generated
    Output: a singular number that represents how large the difference between the standard deviation
            and L1 norm. 
    '''
    dic = run_classical_circ(n, theta)
    vec = dic_to_vec(dic)
    payoff_vec = payoff_finder(vec, rand_payoff_mat)
    val = np.linalg.norm(payoff_vec, ord = 1) - np.std(payoff_vec)
    return -val

vals = []
def callback(res):
    params.append(res)

def callback_classic(res):
    params_c.append(res)

#backup_params = params

params = []
qcost = partial(quantum_cost_function, 5, payoff_mat)
results = minimize(qcost, theta, method='Nelder-Mead', callback = callback)#, tol=1e-9)
results

params_c = []
ccost = partial(classical_cost_function, 5, payoff_mat)
results_c = minimize(ccost, theta_c, method='Nelder-Mead', callback = callback_classic)
results_c

vals = []
for param in params:
    vals.append(0-quantum_cost_function(5, payoff_mat, param))
vals_c = []
for param in params_c:
    vals_c.append(0-classical_cost_function(5, payoff_mat, param))


max_c = np.asarray(vals_c).max()
max_q = np.asarray(vals).max()

with plt.xkcd():
    plt.plot(list(range(len(vals))), vals, label="quantum strategy")
    plt.plot(list(range(len(vals_c))), vals_c, color='r', label = "classical strategy")
    #plt.axhline(max_c, color='r', ls='--', label = 'Maximum quantum value reached')
    #plt.axhline(max_q, color='b', ls='--', label = 'Maximum classical value reached')
    plt.xlabel('Iterations')
    plt.ylabel('Maximization Function')
    plt.legend()

