
import numpy as np
from scipy.optimize import minimize

def _maxcut_cost_matrix(edges, n_nodes):
    dim  = 2 ** n_nodes
    cost = np.zeros(dim)
    for (u, v) in edges:
        for s in range(dim):
            zu = 1 if (s >> u) & 1 else -1
            zv = 1 if (s >> v) & 1 else -1
            cost[s] += 0.5 * (1 - zu * zv)
    return cost

def _apply_cost_layer(state, cost_diag, gamma):
    return np.exp(-1j * gamma * cost_diag) * state

def _apply_mixer_layer(state, beta, n_qubits):
    dim = 2 ** n_qubits
    rx  = np.array([[np.cos(beta), -1j * np.sin(beta)],
                    [-1j * np.sin(beta), np.cos(beta)]])
    mixed = state.copy()
    for q in range(n_qubits):
        step = 2 ** q
        new  = np.zeros(dim, dtype=complex)
        for i in range(0, dim, 2 * step):
            for j in range(step):
                a, b = i + j, i + j + step
                new[a] = rx[0,0]*mixed[a] + rx[0,1]*mixed[b]
                new[b] = rx[1,0]*mixed[a] + rx[1,1]*mixed[b]
        mixed = new
    return mixed

def build_qaoa_circuit(problem, p_layers, params):
    n         = problem["n_nodes"]
    dim       = 2 ** n
    cost_diag = _maxcut_cost_matrix(problem["edges"], n)
    state     = np.ones(dim, dtype=complex) / np.sqrt(dim)
    for p in range(p_layers):
        gamma = params[2*p]
        beta  = params[2*p + 1]
        state = _apply_cost_layer(state, cost_diag, gamma)
        state = _apply_mixer_layer(state, beta, n)
    return state, cost_diag

def optimize_qaoa(problem, p_layers=1, init_params=None, method="COBYLA"):
    n_params = 2 * p_layers
    if init_params is None:
        init_params = np.random.uniform(0, np.pi, n_params)
    cost_diag = _maxcut_cost_matrix(problem["edges"], problem["n_nodes"])
    def objective(params):
        state, cd = build_qaoa_circuit(problem, p_layers, params)
        return -float(np.real(np.conj(state) @ (cd * state)))
    res = minimize(objective, init_params, method=method,
                   options={"maxiter": 300, "rhobeg": 0.5})
    best_state, cd = build_qaoa_circuit(problem, p_layers, res.x)
    probs    = np.abs(best_state) ** 2
    best_cut = int(np.argmax(probs))
    return {
        "optimal_params":  res.x,
        "optimal_cost":   -res.fun,
        "probabilities":   probs,
        "best_bitstring":  format(best_cut, f"0{problem['n_nodes']}b"),
        "converged":       res.success,
    }

def run_qaoa(problem, p_layers=1):
    return optimize_qaoa(problem, p_layers)
