#!/usr/bin/env python3
import matplotlib
matplotlib.use('Agg')
from utils.plot_utils import wash_benchmark_data, parse_data, plot_absolute, plot_relative, image_path
import matplotlib.pyplot as plt

gate_data = parse_data(
    packages=['projectq', 'qiskit', 'cirq', 'quest', 'qulacs', 'yao'],
    labels=['X', 'H', 'T', 'CNOT', 'Toffoli']
)

circuit_data = parse_data(
    packages=['projectq', 'cirq', 'quest', 'pennylane', 'qulacs', 'yao'],
    labels=['QCBM']
)

batch_circuit_data = parse_data(
    packages=['yao'],
    labels=['QCBM (batch)']
)

# plot single gate absolute time
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(10, 8), sharey=True)
ls, labels = plot_absolute(ax1, gate_data, 'X')
plot_absolute(ax2, gate_data, 'H')
plot_absolute(ax3, gate_data, 'CNOT')
plot_absolute(ax4, gate_data, 'Toffoli')

ax1.set(title="X gate")
ax2.set(title="H gate")
ax3.set(title="CNOT gate")
ax4.set(title="Toffoli gate")

lgd = fig.legend(
    ls,
    labels=labels,
    loc="upper right",
    borderaxespad=0.1,
    bbox_to_anchor=(1.12, 0.9))

plt.tight_layout()
plt.savefig(image_path('gates.png'), bbox_extra_artists=(lgd,), bbox_inches='tight')

# plot single gate relative time
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(10, 8), sharey=True)
ls, labels = plot_relative(ax1, gate_data, 'X')
plot_relative(ax2, gate_data, 'H')
plot_relative(ax3, gate_data, 'CNOT')
plot_relative(ax4, gate_data, 'Toffoli')

ax1.set(title="X gate (relative time)")
ax2.set(title="H gate (relative time)")
ax3.set(title="CNOT gate (relative time)")
ax4.set(title="Toffoli gate (relative time)")

lgd = fig.legend(
    ls,
    labels=labels,
    loc="upper right",
    borderaxespad=0.1,
    bbox_to_anchor=(1.12, 0.9))

plt.tight_layout()
plt.savefig(image_path('gates_relative.png'), bbox_extra_artists=(lgd,), bbox_inches='tight')


# plot parameterized quantum circuit absolute time
fig = plt.figure(figsize=(8, 6))
ax = plt.subplot(111)
ls, labels = plot_absolute(ax, circuit_data, 'QCBM')

ax.set(title="Parameterized Circuit")
lgd = ax.legend(
    ls,
    labels=labels,
    loc="upper left")
plt.savefig(image_path('pcircuit.png'), bbox_extra_artists=(lgd,), bbox_inches='tight')

# parameterized quantum circuit relative time
fig = plt.figure(figsize=(8, 6))
ax = plt.subplot(111)
ls, labels = plot_relative(ax, circuit_data, 'QCBM')

ax.set(title="Parameterized Circuit (relative time)")
lgd = ax.legend(
    ls,
    labels=labels,
    loc="upper left",
    borderaxespad=0.1,
    bbox_to_anchor=(1.02, 0.9))
plt.savefig(image_path('pcircuit_relative.png'), bbox_extra_artists=(lgd,), bbox_inches='tight')


# plot batched parameterized quantum circuit absolute time
fig = plt.figure(figsize=(8, 6))
ax = plt.subplot(111)
ls, labels = plot_absolute(ax, batch_circuit_data, 'QCBM (batch)')
ax.set(title="Batched Parameterized Circuit")
lgd = ax.legend(
    ls,
    labels=labels,
    loc="upper left")
plt.savefig(image_path('pcircuit_batch.png'), bbox_extra_artists=(lgd,), bbox_inches='tight')

# batched parameterized quantum circuit relative time
fig = plt.figure(figsize=(8, 6))
ax = plt.subplot(111)
ls, labels = plot_relative(ax, batch_circuit_data, 'QCBM (batch)', to='yao (cuda)', log=False)
ax.set(title="Batched Parameterized Circuit (relative time)")
lgd = ax.legend(
    ls,
    labels=labels,
    loc="upper left")
plt.savefig(image_path('pcircuit_batch_relative.png'), bbox_extra_artists=(lgd,), bbox_inches='tight')
