
import numpy as np
from scipy.optimize import minimize

def create_ansatz(n_qubits, n_layers=1):
    """Return initial random parameters for a hardware-efficient ansatz."""
    return np.random.uniform(-np.pi, np.pi, n_qubits * n_layers * 3)

def _apply_ansatz(params, n_qubits, n_layers=1):
    dim   = 2 ** n_qubits
    state = np.zeros(dim, dtype=complex)
    state[0] = 1.0
    idx = 0
    for _ in range(n_layers):
        # Ry + Rz on each qubit
        for q in range(n_qubits):
            ry_angle = params[idx];   idx += 1
            rz_angle = params[idx];   idx += 1
            ry = np.array([[np.cos(ry_angle/2), -np.sin(ry_angle/2)],
                           [np.sin(ry_angle/2),  np.cos(ry_angle/2)]])
            rz = np.array([[np.exp(-1j*rz_angle/2), 0],
                           [0, np.exp(1j*rz_angle/2)]])
            gate = rz @ ry
            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] = gate[0,0]*state[a] + gate[0,1]*state[b]
                    new[b] = gate[1,0]*state[a] + gate[1,1]*state[b]
            state = new
        # Entangling CNOT layer
        for q in range(n_qubits - 1):
            ctrl, tgt = q, q + 1
            new = state.copy()
            for s in range(dim):
                if (s >> ctrl) & 1:
                    flipped = s ^ (1 << tgt)
                    new[s], new[flipped] = state[flipped], state[s]
            state = new
    return state

def compute_expectation(params, hamiltonian, n_qubits, n_layers=1):
    state = _apply_ansatz(params, n_qubits, n_layers)
    return float(np.real(np.conj(state) @ (hamiltonian @ state)))

def run_vqe(hamiltonian, n_qubits, n_layers=1, init_params=None):
    if init_params is None:
        init_params = create_ansatz(n_qubits, n_layers)
    res = minimize(compute_expectation, init_params,
                   args=(hamiltonian, n_qubits, n_layers),
                   method="COBYLA", options={"maxiter": 500})
    state = _apply_ansatz(res.x, n_qubits, n_layers)
    return {
        "ground_energy":    res.fun,
        "optimal_params":   res.x,
        "optimal_state":    state,
        "converged":        res.success,
        "algorithm":        "VQE",
    }

def run_adapt_vqe(hamiltonian, n_qubits, n_layers=2):
    """AdaptVQE: greedily adds layers until convergence (simplified)."""
    results = []
    best    = None
    for layers in range(1, n_layers + 1):
        r = run_vqe(hamiltonian, n_qubits, n_layers=layers)
        r["algorithm"] = f"AdaptVQE (layers={layers})"
        results.append(r)
        if best is None or r["ground_energy"] < best["ground_energy"]:
            best = r
    best["adapt_history"] = [r["ground_energy"] for r in results]
    return best

def run_sampling_vqe(hamiltonian, n_qubits, n_shots=1024, n_layers=1):
    """SamplingVQE: estimates expectation via shot-based sampling."""
    res   = run_vqe(hamiltonian, n_qubits, n_layers)
    probs = np.abs(res["optimal_state"]) ** 2
    dim   = 2 ** n_qubits
    samples = np.random.choice(dim, size=n_shots, p=probs)
    diag_H  = np.real(np.diag(hamiltonian))
    sampled_energy = float(np.mean(diag_H[samples]))
    res["sampled_energy"] = sampled_energy
    res["n_shots"]        = n_shots
    res["algorithm"]      = "SamplingVQE"
    return res
