
import numpy as np
from scipy.optimize import minimize
from scipy.linalg   import expm

# ── dtype budget (12 GB Colab) ────────────────────────────────────────────────
_C = np.complex64
_R = np.float32

# ─────────────────────────────────────────────────────────────────────────────
#  Low-level statevector helpers  (all in-place, float32)
# ─────────────────────────────────────────────────────────────────────────────

def _zero_state(n: int) -> np.ndarray:
    s = np.zeros(2**n, dtype=_C)
    s[0] = 1.0
    return s

def _apply_ry(state: np.ndarray, q: int, theta: float) -> np.ndarray:
    c, s_ = np.cos(theta / 2), np.sin(theta / 2)
    step   = 1 << q
    dim    = len(state)
    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] = _C( c * state[a] - s_ * state[b])
            new[b] = _C(s_ * state[a] +  c * state[b])
    return new

def _apply_rz(state: np.ndarray, q: int, theta: float) -> np.ndarray:
    step = 1 << q
    dim  = len(state)
    new  = state.copy()
    ep   = np.exp( 1j * theta / 2, dtype=_C)
    em   = np.exp(-1j * theta / 2, dtype=_C)
    for i in range(0, dim, 2 * step):
        for j in range(step):
            new[i + j]        = em * state[i + j]
            new[i + j + step] = ep * state[i + j + step]
    return new

def _apply_cnot(state: np.ndarray, ctrl: int, tgt: int) -> np.ndarray:
    dim = len(state)
    new = state.copy()
    for s in range(dim):
        if (s >> ctrl) & 1:
            new[s] = state[s ^ (1 << tgt)]
    return new

def _apply_cz(state: np.ndarray, q0: int, q1: int) -> np.ndarray:
    dim = len(state)
    new = state.copy()
    for s in range(dim):
        if ((s >> q0) & 1) and ((s >> q1) & 1):
            new[s] = -state[s]
    return new

def _expectation(state: np.ndarray, H: np.ndarray) -> float:
    return float(np.real(np.conj(state) @ (H @ state)))

# ─────────────────────────────────────────────────────────────────────────────
#  Hardware-Efficient Ansatz (brick-layer Ry-Rz + CNOT ring)
# ─────────────────────────────────────────────────────────────────────────────

def _he_ansatz(params: np.ndarray, n: int, layers: int) -> np.ndarray:
    """
    Brick-layer ansatz:
      for each layer: Ry(θ), Rz(φ) on every qubit, then CNOT ring.
    params shape: (layers * 2 * n,)
    """
    state = _zero_state(n)
    idx   = 0
    for _ in range(layers):
        for q in range(n):
            state = _apply_ry(state, q, float(params[idx]));   idx += 1
            state = _apply_rz(state, q, float(params[idx]));   idx += 1
        for q in range(n - 1):
            state = _apply_cnot(state, q, q + 1)
        state = _apply_cnot(state, n - 1, 0)   # ring closure
    return state

def n_params_he(n: int, layers: int) -> int:
    return layers * 2 * n

# ─────────────────────────────────────────────────────────────────────────────
#  HI-VQE — Hardware-Informed VQE
#  Uses CZ (native on many superconducting devices) + asymmetric entanglement
#  + noise-aware parameter initialisation (small angles near identity)
# ─────────────────────────────────────────────────────────────────────────────

def _hi_ansatz(params: np.ndarray, n: int, layers: int) -> np.ndarray:
    """
    HI-VQE ansatz:
      Ry on odd qubits, Rz on even qubits per sublayer,
      CZ on (0,1), (2,3), … entanglement pairs.
    params shape: (layers * n,)
    """
    state = _zero_state(n)
    idx   = 0
    for _ in range(layers):
        for q in range(n):
            if q % 2 == 0:
                state = _apply_ry(state, q, float(params[idx]))
            else:
                state = _apply_rz(state, q, float(params[idx]))
            idx += 1
        for q in range(0, n - 1, 2):
            state = _apply_cz(state, q, q + 1)
    return state

def n_params_hi(n: int, layers: int) -> int:
    return layers * n

def run_hi_vqe(hamiltonian: np.ndarray,
               n_qubits:    int,
               layers:      int  = 2,
               shots:       int  = 0,
               init_params: np.ndarray = None,
               method:      str  = "COBYLA",
               maxiter:     int  = 600) -> dict:
    """
    HI-VQE (Hardware-Informed VQE).

    Parameters
    ----------
    hamiltonian : (2^n, 2^n) complex array
    n_qubits    : number of qubits
    layers      : ansatz depth  (default 2)
    shots       : >0 → shot-noise simulation; 0 → exact
    method      : scipy optimizer name
    maxiter     : max optimizer iterations

    Returns
    -------
    dict: ground_energy, optimal_params, optimal_state, converged, history
    """
    np_hi = n_params_hi(n_qubits, layers)
    if init_params is None:
        # Hardware-informed init: small noise near zero (reduces barren plateaus)
        init_params = np.random.normal(0, 0.05, np_hi).astype(_R)

    history = []

    def cost(p):
        state = _hi_ansatz(p, n_qubits, layers)
        if shots > 0:
            probs = np.abs(state).astype(np.float64) ** 2
            probs /= probs.sum()
            samples = np.random.choice(len(state), size=shots, p=probs)
            diag_H  = np.real(np.diag(hamiltonian)).astype(np.float64)
            val     = float(np.mean(diag_H[samples]))
        else:
            val = _expectation(state, hamiltonian)
        history.append(val)
        return val

    res   = minimize(cost, init_params, method=method,
                     options={"maxiter": maxiter})
    state = _hi_ansatz(res.x, n_qubits, layers)
    return {
        "ground_energy":   res.fun,
        "optimal_params":  res.x,
        "optimal_state":   state,
        "converged":       res.success,
        "history":         history,
        "algorithm":       "HI-VQE",
        "layers":          layers,
        "n_qubits":        n_qubits,
    }

# ─────────────────────────────────────────────────────────────────────────────
#  Standard HE-VQE  (brick-layer, uses _he_ansatz)
# ─────────────────────────────────────────────────────────────────────────────

def run_he_vqe(hamiltonian: np.ndarray,
               n_qubits:    int,
               layers:      int  = 2,
               init_params: np.ndarray = None,
               method:      str  = "COBYLA",
               maxiter:     int  = 600) -> dict:
    """Standard Hardware-Efficient VQE with Ry/Rz brick-layer ansatz."""
    np_he = n_params_he(n_qubits, layers)
    if init_params is None:
        init_params = np.random.uniform(-np.pi, np.pi, np_he).astype(_R)

    history = []

    def cost(p):
        state = _he_ansatz(p, n_qubits, layers)
        val   = _expectation(state, hamiltonian)
        history.append(val)
        return val

    res   = minimize(cost, init_params, method=method,
                     options={"maxiter": maxiter})
    state = _he_ansatz(res.x, n_qubits, layers)
    return {
        "ground_energy":   res.fun,
        "optimal_params":  res.x,
        "optimal_state":   state,
        "converged":       res.success,
        "history":         history,
        "algorithm":       "HE-VQE",
        "layers":          layers,
        "n_qubits":        n_qubits,
    }

# ─────────────────────────────────────────────────────────────────────────────
#  SPSA Optimizer  (Simultaneous Perturbation Stochastic Approximation)
#  Colab-friendly: only 2 circuit evaluations per step (vs 2*n for param-shift)
# ─────────────────────────────────────────────────────────────────────────────

def spsa_optimize(cost_fn, init_params: np.ndarray,
                  n_iter:  int   = 200,
                  a:       float = 0.1,
                  c:       float = 0.1,
                  A:       float = 10.0,
                  alpha:   float = 0.602,
                  gamma:   float = 0.101) -> dict:
    """
    SPSA optimizer — ideal for noisy quantum cost functions.
    Uses Robbins-Monro schedule: a_k = a/(k+A)^alpha, c_k = c/k^gamma.
    """
    p       = init_params.copy().astype(float)
    history = []
    for k in range(1, n_iter + 1):
        ak  = a / (k + A) ** alpha
        ck  = c / k ** gamma
        delta = 2 * np.random.randint(0, 2, size=len(p)) - 1   # Rademacher
        fp  = cost_fn(p + ck * delta)
        fm  = cost_fn(p - ck * delta)
        gk  = (fp - fm) / (2 * ck * delta)
        p  -= ak * gk
        history.append(float(cost_fn(p)))
    return {
        "optimal_params": p,
        "optimal_value":  history[-1],
        "history":        history,
        "n_iter":         n_iter,
        "optimizer":      "SPSA",
    }

def run_spsa_vqe(hamiltonian: np.ndarray,
                 n_qubits:    int,
                 layers:      int = 2,
                 n_iter:      int = 300) -> dict:
    """HI-VQE ansatz optimised with SPSA (noise-robust, Colab-fast)."""
    np_hi = n_params_hi(n_qubits, layers)
    init  = np.random.normal(0, 0.05, np_hi)

    def cost(p):
        return _expectation(_hi_ansatz(p, n_qubits, layers), hamiltonian)

    res   = spsa_optimize(cost, init, n_iter=n_iter)
    state = _hi_ansatz(res["optimal_params"], n_qubits, layers)
    return {
        "ground_energy":  res["optimal_value"],
        "optimal_params": res["optimal_params"],
        "optimal_state":  state,
        "history":        res["history"],
        "algorithm":      "SPSA-HI-VQE",
        "n_qubits":       n_qubits,
        "layers":         layers,
    }

# ─────────────────────────────────────────────────────────────────────────────
#  AdaptVQE+  —  greedy operator pool with gradient screening
#  Pool: {Ry(q), Rz(q), CNOT(q,q+1)} for all q
#  Stops when max gradient < threshold or max_operators reached
# ─────────────────────────────────────────────────────────────────────────────

def _pool_operator_names(n: int):
    ops = []
    for q in range(n):
        ops.append(("Ry", q))
        ops.append(("Rz", q))
    for q in range(n - 1):
        ops.append(("CNOT", q, q + 1))
    return ops

def _apply_pool_op(state, op):
    name = op[0]
    if name == "Ry":
        return _apply_ry(state, op[1], np.pi / 4)
    elif name == "Rz":
        return _apply_rz(state, op[1], np.pi / 4)
    elif name == "CNOT":
        return _apply_cnot(state, op[1], op[2])
    return state

def run_adapt_vqe_plus(hamiltonian:     np.ndarray,
                       n_qubits:        int,
                       max_operators:   int   = 10,
                       grad_threshold:  float = 1e-3,
                       maxiter_inner:   int   = 200) -> dict:
    """
    AdaptVQE+: greedily selects operators from a pool based on energy gradient.

    Each step:
      1. Compute gradient of each pool operator on current state.
      2. Select operator with largest |gradient|.
      3. Optimise all current parameters jointly.
      4. Repeat until converged or max_operators reached.
    """
    pool     = _pool_operator_names(n_qubits)
    state    = _zero_state(n_qubits)
    selected = []     # list of (op, param)
    history  = []
    energy   = _expectation(state, hamiltonian)

    for step in range(max_operators):
        # Gradient screening
        grads = []
        for op in pool:
            s_plus  = _apply_pool_op(state, op)
            g       = abs(_expectation(s_plus, hamiltonian) - energy)
            grads.append(g)
        best_idx = int(np.argmax(grads))
        if grads[best_idx] < grad_threshold:
            break
        selected.append((pool[best_idx], 0.0))

        # Re-optimise all parameters
        def cost(params):
            s = _zero_state(n_qubits)
            for (op, _), p in zip(selected, params):
                nm = op[0]
                if nm == "Ry":
                    s = _apply_ry(s, op[1], float(p))
                elif nm == "Rz":
                    s = _apply_rz(s, op[1], float(p))
                elif nm == "CNOT":
                    s = _apply_cnot(s, op[1], op[2])
            return _expectation(s, hamiltonian)

        init = np.array([p for (_, p) in selected])
        res  = minimize(cost, init, method="COBYLA",
                        options={"maxiter": maxiter_inner})
        for i, val in enumerate(res.x):
            selected[i] = (selected[i][0], val)

        energy = res.fun
        history.append(energy)
        state  = _zero_state(n_qubits)
        for (op, p) in selected:
            nm = op[0]
            if nm == "Ry":
                state = _apply_ry(state, op[1], float(p))
            elif nm == "Rz":
                state = _apply_rz(state, op[1], float(p))
            elif nm == "CNOT":
                state = _apply_cnot(state, op[1], op[2])

    return {
        "ground_energy":    energy,
        "optimal_state":    state,
        "selected_ops":     [(str(op), float(p)) for op, p in selected],
        "n_operators":      len(selected),
        "history":          history,
        "algorithm":        "AdaptVQE+",
        "n_qubits":         n_qubits,
    }
