# Define a new benchmark here.
import iqm.benchmarks as qcvv
from iqm.benchmarks import Benchmark
import xarray as xr
from iqm.benchmarks.benchmark_definition import add_counts_to_dataset
from qiskit import QuantumCircuit, transpile
import matplotlib.pyplot as plt

def generate_readout_test_circuit(backend, N_qubits):
    circuits=[]
    qb_to_measure = range(N_qubits)

    qc = QuantumCircuit(N_qubits,N_qubits)
    for qubit in range(N_qubits):
        pass
    qc.barrier()
    qc.measure(qb_to_measure,qb_to_measure)
    circuits.append(qc)

    qc = QuantumCircuit(N_qubits,N_qubits)
    for qubit in range(N_qubits):
        qc.x(qubit)
    qc.barrier()
    qc.measure(qb_to_measure,qb_to_measure)
    circuits.append(qc)
    qc_transpiled = transpile(circuits, backend, optimization_level=0)
    return qc_transpiled


def plot_histogram(
    dataset: xr.Dataset,
):
    fig = plt.figure()
    for ii in range(50):
        try:
            plt.bar(dataset[f"state_{ii}"], dataset[f"counts_{ii}"], label=f"circuit={ii}")
        except:
           break
    plt.xticks(rotation=90)
    plt.legend()
    plt.close()
    return fig


def readout_analysis(run):
    dataset = run.dataset
    plots={}
    plots["histogram"] = plot_histogram(dataset)
    return qcvv.AnalysisResult(dataset=dataset, plots=plots)

class ReadoutTest(Benchmark):
    default_options = dict(N_shots=2048)
    analysis_function = staticmethod(readout_analysis)
    name: str = "readout_test"

    def execute(self, backend)->xr.Dataset:
        """
        Executes the benchmark.
        """
        dataset = xr.Dataset()
        circuits = generate_readout_test_circuit(backend, self.options["N_qubits"])
        job = backend.run(circuits, shots=self.options["N_shots"])
        add_counts_to_dataset(job.result(), dataset)
        return dataset

from iqm.qiskit_iqm import IQMProvider

iqm_server_url ="https://example-station.qc.iqm.fi/cocos/"

benchmark = ReadoutTest(N_qubits=4)
provider = IQMProvider(iqm_server_url)
backend = provider.get_backend()
benchmark.run(backend)

result = benchmark.analyze()
result.plot_all()