
import numpy as np
from scipy.optimize import minimize

def _angle_embed(x, n_qubits):
    """Angle-encode data vector x into a state vector."""
    dim   = 2 ** n_qubits
    state = np.zeros(dim, dtype=complex)
    state[0] = 1.0
    for q in range(min(n_qubits, len(x))):
        angle = float(x[q])
        ry    = np.array([[np.cos(angle/2), -np.sin(angle/2)],
                          [np.sin(angle/2),  np.cos(angle/2)]])
        step  = 2 ** q
        new   = state.copy()
        for i in range(0, dim, 2*step):
            for j in range(step):
                a, b = i+j, i+j+step
                new[a] = ry[0,0]*state[a] + ry[0,1]*state[b]
                new[b] = ry[1,0]*state[a] + ry[1,1]*state[b]
        state = new
    return state

def _variational_layer(state, params, n_qubits):
    dim = 2 ** n_qubits
    for q in range(n_qubits):
        angle = params[q] if q < len(params) else 0.0
        rz    = np.array([[np.exp(-1j*angle/2), 0],
                          [0, np.exp(1j*angle/2)]])
        step  = 2 ** q
        new   = state.copy()
        for i in range(0, dim, 2*step):
            for j in range(step):
                a, b = i+j, i+j+step
                new[a] = rz[0,0]*state[a]
                new[b] = rz[1,1]*state[b]
        state = new
    return state

def build_qnn(n_qubits: int = 2, n_layers: int = 1) -> dict:
    """Return a QNN config dict."""
    n_params = n_qubits * n_layers
    params   = np.random.uniform(-np.pi, np.pi, n_params)
    return {"n_qubits": n_qubits, "n_layers": n_layers, "params": params}

def _qnn_predict(x, params, n_qubits, n_layers):
    state = _angle_embed(x, n_qubits)
    chunk = n_qubits
    for l in range(n_layers):
        p     = params[l*chunk:(l+1)*chunk]
        state = _variational_layer(state, p, n_qubits)
    # Measure Z on qubit 0
    dim    = 2 ** n_qubits
    z_exp  = sum(np.abs(state[s])**2 * (1 if not (s & 1) else -1) for s in range(dim))
    return float(np.real(z_exp))

def train_qnn(data: np.ndarray, labels: np.ndarray,
              n_qubits: int = 2, n_layers: int = 1,
              epochs: int = 50) -> dict:
    """Train a variational QNN classifier."""
    n_params = n_qubits * n_layers
    params   = np.random.uniform(-np.pi, np.pi, n_params)
    def loss(p):
        preds = np.array([_qnn_predict(x, p, n_qubits, n_layers) for x in data])
        return float(np.mean((preds - labels) ** 2))
    res = minimize(loss, params, method="COBYLA",
                   options={"maxiter": epochs * 10})
    preds = np.array([_qnn_predict(x, res.x, n_qubits, n_layers) for x in data])
    acc   = float(np.mean((np.sign(preds) == np.sign(labels))))
    return {
        "optimal_params": res.x,
        "final_loss":     res.fun,
        "accuracy":       acc,
        "predictions":    preds,
        "converged":      res.success,
    }

def quantum_kernel(x1: np.ndarray, x2: np.ndarray, n_qubits: int = 2) -> float:
    """Quantum kernel: |<phi(x1)|phi(x2)>|^2"""
    s1 = _angle_embed(x1, n_qubits)
    s2 = _angle_embed(x2, n_qubits)
    return float(np.abs(np.conj(s1) @ s2) ** 2)

def quantum_kernel_matrix(X: np.ndarray, n_qubits: int = 2) -> np.ndarray:
    n = len(X)
    K = np.zeros((n, n))
    for i in range(n):
        for j in range(i, n):
            K[i,j] = K[j,i] = quantum_kernel(X[i], X[j], n_qubits)
    return K
