
import numpy as np
from scipy.optimize import minimize

def parameter_shift_gradient(cost_fn, params: np.ndarray,
                              shift: float = np.pi / 2) -> np.ndarray:
    """
    Parameter-shift rule gradient: d/d(p_k) C = [C(p+shift) - C(p-shift)] / 2
    Works for any quantum cost function with parameterized Pauli rotations.
    """
    grad = np.zeros_like(params)
    for k in range(len(params)):
        p_plus          = params.copy(); p_plus[k]  += shift
        p_minus         = params.copy(); p_minus[k] -= shift
        grad[k]         = (cost_fn(p_plus) - cost_fn(p_minus)) / 2.0
    return grad

def optimize_parameters(cost_fn, init_params: np.ndarray,
                         method: str = "BFGS",
                         use_grad: bool = True,
                         max_iter: int = 500) -> dict:
    """
    Classical optimizer loop with optional parameter-shift gradients.
    Supports: BFGS (with grad), COBYLA, Nelder-Mead (gradient-free).
    """
    history = []
    def tracked_cost(p):
        val = float(cost_fn(p))
        history.append(val)
        return val

    if use_grad and method == "BFGS":
        def jac(p):
            return parameter_shift_gradient(cost_fn, p)
        res = minimize(tracked_cost, init_params, method="BFGS",
                       jac=jac, options={"maxiter": max_iter})
    else:
        res = minimize(tracked_cost, init_params, method=method,
                       options={"maxiter": max_iter})

    return {
        "optimal_params":  res.x,
        "optimal_value":   res.fun,
        "converged":       res.success,
        "n_evaluations":   len(history),
        "history":         history,
        "method":          method,
    }

def variational_loop(cost_fn, init_params: np.ndarray,
                     n_iter: int = 100, lr: float = 0.05) -> dict:
    """
    Custom gradient-descent loop using parameter-shift rule.
    Good for debugging and visualizing convergence.
    """
    params  = init_params.copy()
    history = []
    for i in range(n_iter):
        val    = float(cost_fn(params))
        grad   = parameter_shift_gradient(cost_fn, params)
        params = params - lr * grad
        history.append(val)
    return {
        "optimal_params":  params,
        "optimal_value":   float(cost_fn(params)),
        "history":         history,
        "n_iter":          n_iter,
        "lr":              lr,
        "method":          "GradientDescent+ParameterShift",
    }
