import numpy as np
import random

from qiskit.circuit.random import random_circuit
from purplecaffeine.core import BaseStorage, LocalStorage, Trial
from purplecaffeine.widget import Widget

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

local_storage = LocalStorage("./trials/")
n_qubits = 4
depth = 10
shots = 2000

for i in range(0, 3):
    with Trial("Example trial " + str(i), storage=local_storage) as trial:
        # track parameter
        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 circuits
        circuit = random_circuit(n_qubits, depth)
        trial.add_circuit("circuit", circuit)
        trial.add_circuit("another circuit", random_circuit(n_qubits, depth))

        # track operators
        obs = random_pauli(n_qubits)
        trial.add_operator("obs", obs)
        for idx in range(5):
            trial.add_operator(f"obs {idx}", random_pauli(n_qubits))
        
        # track tags
        trial.add_tag(f"tag {i}")

        # track texts
        trial.add_text("text one", "This text will be displayed in text tab")
        trial.add_text("text two", "This text will be displayed in text tab as well")
    
        # run
        exp_value = Estimator().run(circuit, obs, shots=shots).result().values.item()
        
        # track history data as metric
        for _ in range(100):
            trial.add_metric("history", random.randint(0, 10))
            
        # track results as metric
        trial.add_metric("exp_value", exp_value)

Widget(local_storage).show()