# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/Speed_comparison.ipynb (unless otherwise specified).

__all__ = ['ex', 'graph', 'qiskit_time', 'quimb_time', 'qtensor_time', 'gen_time_data', 'plot_xar_data',
           'time_comparison_xarray', 'ql']

# Cell
import numpy as np
import qtensor as qt
from qtensor.tests.qiskit_qaoa_energy import simulate_qiskit_amps
from qtensor.tests.qaoa_quimb import simulate_one_parallel as simulate_quimb_energy
import matplotlib.pyplot as plt
import time

from cartesian_explorer import Explorer

# Cell
ex = Explorer()

# Cell
@ex.provider
def graph(N, d=3, graph_type='random', seed=10):
    return qt.toolbox.random_graph(nodes=N, type=graph_type, degree=d, seed=seed)

# Cell
@ex.provider
def qiskit_time(graph, p):
    gamma, beta = [.1]*p, [.3]*p
    start = time.time()
    try:
        _ = simulate_qiskit_amps(graph, gamma, beta)
    except Exception as e:
        print('Exception:', e)
        return None
    end = time.time()
    return end - start

# Cell
@ex.provider
def quimb_time(graph, p,
               n_processes=2,
               qmb_optimizer_time=0.3,
               qmb_ordering_algo='uniform'
              ):
    start = time.time()
    try:
        _ = simulate_quimb_energy(graph, p,
                                  n_processes=n_processes,
                                  optimizer_time=qmb_optimizer_time,
                                  ordering_algo=qmb_ordering_algo
                                 )
    except Exception as e:
        print('Exception:', e)
        return None
    end = time.time()
    return end - start

# Cell
@ex.provider
def qtensor_time(graph, p,
               n_processes=2,
               qtn_ordering_algo='greedy'
              ):
    gamma, beta = [.1]*p, [.3]*p
    opt = qt.toolbox.get_ordering_algo(qtn_ordering_algo)
    sim = qt.QAOAQtreeSimulator(qt.DefaultQAOAComposer, optimizer=opt)
    start = time.time()
    try:
        _ = sim.energy_expectation_parallel(graph, gamma, beta, n_processes=n_processes)
    except Exception as e:
        print('Exception:', e)
        return None
    end = time.time()
    return end - start

# Cell
def gen_time_data(Ns, ps, **kwargs):
    xar = ex.get_variables_xarray(
        ('qtensor_time', 'quimb_time', 'qiskit_time'),
        p=ps,
        N=Ns,
        **kwargs,
    )
    return xar

# Cell
def plot_xar_data(xar_orig):
    xar = xar_orig.transpose('p', 'varname', 'N')
    def custom_plot(x, y, **kwargs):
        plt.plot(x, y, '*-', **kwargs)
        plt.yscale('log')
        plt.xscale('log', base=2)
    fig = ex.plot_xarray(xar, plot_func=custom_plot)

    axes = fig.axes
    axes[0].get_shared_y_axes().join(*axes)
    axes[0].get_shared_x_axes().join(*axes)
    ps = [x.data for x in xar.coords['p']]


    [ax.autoscale() for ax in axes]
    [ax.set_ylabel('Time, sec') for ax in axes]
    [ax.set_title(x) for ax, x in zip(axes, [f"p = {p}" for p in ps])]
    fig.suptitle('Time for single energy query on a single Skylake node', fontfamily='serif', style='italic')
    fig.set_size_inches(10.5, 4.5)
    fig.tight_layout(rect=[0, 0.03, 1, 0.98])


# Cell
from qtensor_specs import cli, click

# Cell
@cli.command()
@click.option('--n_len', default=4)
@click.option('--p_len', default=2)
@click.option('-O', '--qtn_ordering_algo', default='greedy')
@click.option('-N', '--n_processes', default=2)
@click.argument('filename')
def time_comparison_xarray(filename, n_len, p_len, **kwargs):
    Ns = 2*(.5*np.pi**(np.arange(0, 12)/2)).astype(np.int) + 10
    Ns = Ns[:n_len]
    ps = [2, 3, 4]
    ps = ps[:p_len]
    kwargs = {k:[v] for k, v in kwargs.items()}
    xar = gen_time_data(Ns, ps, **kwargs)
    xar.to_netcdf(filename)


# Cell
import logging
ql = logging.getLogger('qiskit')
ql.setLevel(logging.ERROR)