from qiskit.providers.ibmq.runtime import UserMessenger
from qiskit import Aer, IBMQ
import numpy as np
from runtime_test import main
import matplotlib.pyplot as plt

provider = IBMQ.load_account()
provider = IBMQ.get_provider(group='deployed')

def interim_result_callback(job_id, interim_result):
    print(interim_result)

import uuid

meta = {
  "name": #"grad_paper-{}".format(uuid.uuid4()),
    "grad_paper-lima",
  "description": "A sample gradient program.",
  "max_execution_time": 10000,
}

meta["parameters"] = [
    {"name": "shots", "description": "Number of shots", "type": "int", "required": True},
    {"name": "init_point", "description": "Initial parameters", "type": "ndarray", "required": True},
  ]

meta['return_values'] = [
    {"name": "result", "description": "Final result in list.", "type": "list"}
  ]

meta["interim_results"] = [
    {"name": "vals", "description": "Iteration, value at current optimization step", "type": "list"},
  ]

#backend = provider.backend.ibmq_qasm_simulator
backend = provider.backend.ibmq_lima
#backend = provider.backend.ibmq_manila
options = {'backend_name': backend.name()}
program_id = provider.runtime.upload_program(data='runtime_test.py', metadata=meta)

initial_point = np.random.uniform(0, 2 * np.pi, 5)

inputs = {}
inputs['init_point'] = initial_point

shots = [10, 100, 1000]
fig, (ax1, ax2) = plt.subplots(1, 2)
fig.tight_layout()
for s in shots:
    print(s)
    inputs['shots'] = s
    job = provider.runtime.run(program_id, options=options, inputs=inputs, callback=interim_result_callback)
    gd_loss, x_values = job.result()
    cir_evals = np.array([i for i in range(len(gd_loss))])
    cir_evals *= (2 * len(initial_point) * s)
    ax1.plot(cir_evals, gd_loss, label='Shots %d'%s)
    ax2.plot(x_values[:,0], x_values[:,1], label='Shots %d'%s)

# Make the graph look nice
#ax1.axhline(-1, ls='--', c='tab:red', label='Target')
ax1.set_ylabel('Cost')
ax1.set_xlabel('Circuit Evaluations')
ax1.set_xscale('log')
ax1.legend()
ax2.set_ylim(0, np.pi)
ax2.set_xlim(-np.pi/2, np.pi/2)
ax2.set_ylabel('Theta 2')
ax2.set_xlabel('Theta 1')
#ax2.legend()
plt.show()

