%matplotlib widget
import logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)

from slam.basis import CircuitTemplate

basis = CircuitTemplate()
basis.build(3)
basis.spanning_range = range(3, 4)
basis.circuit.draw()

from slam.cost_function import BasicCost, MakhlinFunctionalCost

objective1 = BasicCost()
objective2 = MakhlinFunctionalCost()

from slam.sampler import GateSample
from qiskit.circuit.library import SwapGate

sampler = GateSample(gate=SwapGate())

from slam.optimizer import TemplateOptimizer

# by default use BFGS
optimizer1 = TemplateOptimizer(
    basis=basis, objective=objective1, override_fail=1, use_callback=True
)
optimizer2 = TemplateOptimizer(
    basis=basis, objective=objective2, override_fail=1, use_callback=True
)
##
optimizer3 = TemplateOptimizer(
    basis=basis,
    objective=objective1,
    override_fail=1,
    use_callback=True,
    override_method="Nelder-Mead",
)
optimizer4 = TemplateOptimizer(
    basis=basis,
    objective=objective2,
    override_fail=1,
    use_callback=True,
    override_method="Nelder-Mead",
)

from slam.utils.visualize import optimizer_training_plot

_ret1 = optimizer1.approximate_from_distribution(sampler)
optimizer_training_plot(_ret1[0], _ret1[1]);

_ret2 = optimizer2.approximate_from_distribution(sampler)
optimizer_training_plot(_ret2[0], _ret2[1])

_ret3 = optimizer3.approximate_from_distribution(sampler)
optimizer_training_plot(_ret3[0], _ret3[1]);

_ret4 = optimizer4.approximate_from_distribution(sampler)
optimizer_training_plot(_ret4[0], _ret4[1]);