"""
COBYLA optimization loop for QAOA parameter training.

The objective function appends every (gamma, beta, energy) triple to a history
dict so the UI can replay the variational learning process step-by-step.
"""

from scipy.optimize import minimize
from .simulator import simulate_qaoa


# Initial-parameter heuristics tuned per landscape
_INITIAL_PARAMS = {
    "standard":   [2.0, 1.0],
    "equality":   [0.5, 0.5],
    "inequality": [0.4, 0.4],
}


def run_cobyla(cost_function, num_qubits: int, mode: str = "standard",
               shots: int = 8192, seed: int = 42,
               max_iter: int = 60, tol: float = 1e-4):
    """
    Run COBYLA to minimize the expected QAOA energy.

    Parameters
    ----------
    cost_function : callable
        Classical cost function (see core/costs.py).
    num_qubits : int
        Number of qubits.
    mode : str
        One of "standard", "equality", "inequality" — selects initial params.
    shots : int
        Shots per objective evaluation (high count for a smooth landscape).
    seed : int
        Fixed seed during optimization to make the landscape deterministic.
    max_iter : int
        Maximum COBYLA iterations.
    tol : float
        Convergence tolerance.

    Returns
    -------
    result : OptimizeResult
        scipy result object.
    history : dict
        Keys "gamma", "beta", "energy" — one entry per objective call.
    """
    history = {"gamma": [], "beta": [], "energy": []}

    def objective(params):
        g, b = params
        _, energy, _, _ = simulate_qaoa(g, b, cost_function,
                                        num_qubits=num_qubits,
                                        shots=shots, seed=seed)
        history["gamma"].append(float(g))
        history["beta"].append(float(b))
        history["energy"].append(float(energy))
        return energy

    initial = _INITIAL_PARAMS.get(mode, [0.5, 0.5])
    result = minimize(objective, initial, method="COBYLA",
                      tol=tol, options={"maxiter": max_iter})
    return result, history