from qiskit_algorithms import QAOA, SamplingVQE
from qiskit_algorithms.optimizers import COBYLA
from qiskit.primitives import Sampler
from qiskit_optimization.algorithms import MinimumEigenOptimizer
from qiskit.circuit.library import TwoLocal
from qiskit_optimization import QuadraticProgram

from purplecaffeine import Trial, LocalStorage

qubo = QuadraticProgram(name="qubo_trial")
qubo.binary_var("x")
qubo.binary_var("y")
qubo.binary_var("z")
qubo.minimize(linear=[1, -2, 3], quadratic={("x", "y"): 1, ("x", "z"): -1, ("y", "z"): 2})
print(qubo.prettyprint())

local_storage = LocalStorage("./trials")

with Trial("QAOA trial", storage=local_storage) as trial:
    # track some parameters
    trial.add_parameter("algo", "qaoa")
    trial.add_parameter("sampler", "qiskit.primitives.Sampler")
    trial.add_parameter("optimizer", "qiskit.algorithms.optimizers.COBYLA")
    
    # track usefull data
    trial.add_text("qubo", qubo.export_as_lp_string())
    qaoa_mes = QAOA(
        sampler=Sampler(),
        optimizer=COBYLA(),
        callback=lambda idx, params, mean, std: trial.add_metric("qaoa_history", mean)
    )
    
    # run
    qaoa = MinimumEigenOptimizer(qaoa_mes)
    qaoa_result = qaoa.solve(qubo)
    
    # track results of run
    trial.add_text("qaoa_result", qaoa_result.prettyprint())

ansatz = TwoLocal(rotation_blocks="ry", entanglement_blocks="cz")

with Trial("VQE trial", storage=local_storage) as trial:
    # track some parameters
    trial.add_parameter("algo", "vqe")
    trial.add_parameter("sampler", "qiskit.primitives.Sampler")
    trial.add_parameter("optimizer", "qiskit.algorithms.optimizers.COBYLA")
    
    # track usefull data
    trial.add_text("qubo", qubo.export_as_lp_string())
    
    # track some objects
    trial.add_circuit("ansatz", ansatz)
    
    vqe_mes = SamplingVQE(
        sampler=Sampler(),
        ansatz=ansatz,
        optimizer=COBYLA(),
        callback=lambda idx, params, mean, std: trial.add_metric("vqe_history", mean)
    )
    
    # run
    vqe = MinimumEigenOptimizer(vqe_mes)
    vqe_result = vqe.solve(qubo)
    
    # track results of run
    trial.add_text("vqe_result", vqe_result.prettyprint())

local_storage.list()