from qiskit import QuantumCircuit, transpile
from qiskit_aer import AerSimulator
from qiskit.tools.visualization import plot_histogram

import random

circ = QuantumCircuit(40, 40)

# Initialize with a Hadamard layer
circ.h(range(40))
# Apply some random CNOT and T gates
qubit_indices = [i for i in range(40)]
for i in range(10):
    control, target, t = random.sample(qubit_indices, 3)
    circ.cx(control, target)
    circ.t(t)
circ.measure(range(40), range(40))

# Create statevector method simulator
statevector_simulator = AerSimulator(method='statevector')

# Transpile circuit for backend
tcirc = transpile(circ, statevector_simulator)

# Try and run circuit
statevector_result =  statevector_simulator.run(tcirc, shots=1).result()
print('This succeeded?: {}'.format(statevector_result.success))
print('Why not? {}'.format(statevector_result.status))

# Create extended stabilizer method simulator
extended_stabilizer_simulator = AerSimulator(method='extended_stabilizer')

# Transpile circuit for backend
tcirc = transpile(circ, extended_stabilizer_simulator)

extended_stabilizer_result = extended_stabilizer_simulator.run(tcirc, shots=1).result()
print('This succeeded?: {}'.format(extended_stabilizer_result.success))

small_circ = QuantumCircuit(2, 2)
small_circ.h(0)
small_circ.cx(0, 1)
small_circ.t(0)
small_circ.measure([0, 1], [0, 1])
# This circuit should give 00 or 11 with equal probability...
expected_results ={'00': 50, '11': 50}

tsmall_circ = transpile(small_circ, extended_stabilizer_simulator)
result = extended_stabilizer_simulator.run(
    tsmall_circ, shots=100).result()
counts = result.get_counts(0)
print('100 shots in {}s'.format(result.time_taken))

plot_histogram([expected_results, counts],
               legend=['Expected', 'Extended Stabilizer'])

# Add runtime options for extended stabilizer simulator
opts = {'extended_stabilizer_approximation_error': 0.03}

reduced_error = extended_stabilizer_simulator.run(
    tsmall_circ, shots=100, **opts).result()

reduced_error_counts = reduced_error.get_counts(0)
print('100 shots in {}s'.format(reduced_error.time_taken))
plot_histogram([expected_results, reduced_error_counts],
               legend=['Expected', 'Extended Stabilizer'])

print("The circuit above, with 100 shots at precision 0.03 "
      "and default mixing time, needed {}s".format(int(reduced_error.time_taken)))

opts = {
    'extended_stabilizer_approximation_error': 0.03,
    'extended_stabilizer_mixing_time': 100
}

optimized = extended_stabilizer_simulator.run(
    tsmall_circ, shots=100, **opts).result()

print('Dialing down the mixing time, we completed in just {}s'.format(optimized.time_taken))

# We set these options here only to make the example run more quickly.
opts = {'extended_stabilizer_mixing_time': 100}

multishot = extended_stabilizer_simulator.run(
    tcirc, shots=100, **opts).result()
print("100 shots took {} s".format(multishot.time_taken))

opts = {
    'extended_stabilizer_measure_sampling': True,
    'extended_stabilizer_mixing_time': 100
}

measure_sampling = extended_stabilizer_simulator.run(
    circ, shots=100, **opts).result()
print("With the optimization, 100 shots took {} s".format(result.time_taken))

import qiskit.tools.jupyter
%qiskit_version_table
%qiskit_copyright

