
import numpy as np

def _qft_matrix(n):
    N   = 2 ** n
    w   = np.exp(2j * np.pi / N)
    idx = np.arange(N)
    return np.power(w, np.outer(idx, idx)) / np.sqrt(N)

def qpe(unitary: np.ndarray, eigenstate: np.ndarray, n_ancilla: int) -> dict:
    """
    Standard Quantum Phase Estimation.
    Estimates phase phi such that U|psi> = e^{2pi i phi}|psi>.
    """
    n_anc = n_ancilla
    dim   = 2 ** n_anc
    # Build controlled-U^(2^k) phases for each ancilla
    eigs  = np.linalg.eigvals(unitary)
    # Pick eigenvalue closest to eigenstate
    evecs = np.linalg.eig(unitary)[1]
    overlaps = np.abs(evecs.T @ eigenstate)
    best_eig  = eigs[np.argmax(overlaps)]
    true_phase = np.angle(best_eig) / (2 * np.pi) % 1.0

    # Simulate QPE register via phase kickback
    register = np.zeros(dim, dtype=complex)
    register[0] = 1.0
    for k in range(n_anc):
        phase_k = (2 ** k) * true_phase
        register[1 << k] = np.exp(2j * np.pi * phase_k)
    register /= np.linalg.norm(register)

    # Inverse QFT
    iqft   = np.conj(_qft_matrix(n_anc)).T
    out    = iqft @ register
    probs  = np.abs(out) ** 2
    meas   = int(np.argmax(probs))
    est    = meas / dim

    return {
        "estimated_phase": est,
        "true_phase":      true_phase,
        "error":           abs(est - true_phase),
        "probabilities":   probs,
        "n_ancilla":       n_anc,
    }

def iterative_qpe(unitary: np.ndarray, eigenstate: np.ndarray, n_bits: int) -> dict:
    """Iterative QPE — one ancilla qubit, bit-by-bit phase readout."""
    eigs  = np.linalg.eigvals(unitary)
    evecs = np.linalg.eig(unitary)[1]
    overlaps  = np.abs(evecs.T @ eigenstate)
    best_eig  = eigs[np.argmax(overlaps)]
    true_phase = np.angle(best_eig) / (2 * np.pi) % 1.0

    bits   = []
    omega  = 0.0
    for k in range(n_bits - 1, -1, -1):
        phase_shift = -2 * np.pi * omega
        eff_phase   = (2 ** k) * true_phase + phase_shift / (2 * np.pi)
        prob_one    = np.sin(np.pi * eff_phase) ** 2
        bit         = 1 if np.random.random() < prob_one else 0
        bits.append(bit)
        omega       = (omega + bit * 2 ** (-(n_bits - k))) % 1.0

    estimated = sum(b * 2 ** (-(i + 1)) for i, b in enumerate(bits))
    return {
        "estimated_phase": estimated,
        "true_phase":      true_phase,
        "error":           abs(estimated - true_phase),
        "bits":            bits,
        "n_bits":          n_bits,
    }

def hamiltonian_qpe(hamiltonian: np.ndarray, eigenstate: np.ndarray,
                    n_ancilla: int, t: float = 1.0) -> dict:
    """HPE via matrix exponentiation: U = exp(-iHt)."""
    U = np.linalg.matrix_power(
        np.eye(hamiltonian.shape[0], dtype=complex),1)
    from scipy.linalg import expm
    U = expm(-1j * hamiltonian * t)
    result = qpe(U, eigenstate, n_ancilla)
    evals  = np.linalg.eigvalsh(hamiltonian)
    result["estimated_energy"] = result["estimated_phase"] * 2 * np.pi / t
    result["true_ground_energy"] = float(evals[0])
    return result
