import os
from qiskit.circuit.random import random_circuit
from qiskit.quantum_info.random import random_pauli
from qiskit.primitives import Estimator

from purplecaffeine.core import Trial, LocalStorage

n_qubits = 4
depth = 3
shots = 2000

circuit = random_circuit(n_qubits, depth)
circuit.draw(fold=-1)

obs = random_pauli(n_qubits)
obs

local_storage = LocalStorage("./trials")
#You could also set the API storage or S3 storage by calling : 
#api_storage = ApiStorage(host="http://localhost:8000", username="MY_USERNAME", password="MY_PASS")
#s3_storage = S3Storage("bucket", access_key="MY_KEY", secret_access_key="MY_SECRET")

with Trial("Example trial", storage=local_storage) as trial:
    # track some parameters
    trial.add_parameter("estimator", "qiskit.primitives.Estimator")
    trial.add_parameter("depth", depth)
    trial.add_parameter("n_qubits", n_qubits)
    trial.add_parameter("shots", shots)
    
    # track objects of interest
    trial.add_circuit("circuit", circuit)
    trial.add_operator("obs", obs)

    # run
    exp_value = Estimator().run(circuit, obs, shots=shots).result().values.item()
    
    # track results of run
    trial.add_metric("exp_value", exp_value)

local_storage.list()

trials = local_storage.list()

random_trial = trials[0]
trial_uuid = random_trial.uuid

local_storage.get(trial_uuid)