
import numpy as np
from scipy.linalg import expm

_C = np.complex64

def _dominant_eigenphase(U: np.ndarray, psi: np.ndarray):
    """Return the phase of U closest to the eigenvalue dominating psi."""
    evals, evecs = np.linalg.eig(U.astype(np.complex128))
    overlaps     = np.abs(evecs.T.conj() @ psi.astype(np.complex128))
    idx          = int(np.argmax(overlaps))
    phase        = float(np.angle(evals[idx])) / (2 * np.pi) % 1.0
    return phase, evals[idx]

# ─────────────────────────────────────────────────────────────────────────────
#  1. Textbook QPE  (exact ancilla simulation)
# ─────────────────────────────────────────────────────────────────────────────

def textbook_qpe(U:         np.ndarray,
                 eigenstate: np.ndarray,
                 n_ancilla:  int = 6) -> dict:
    """
    Textbook QPE via exact simulation of the ancilla register.

    Constructs the ancilla state:
      sum_{k=0}^{2^t-1}  e^{2pi i k phi} |k>  (QFT basis)
    then applies inverse QFT to read out phi.
    """
    true_phase, dominant_eval = _dominant_eigenphase(U, eigenstate)
    t   = n_ancilla
    dim = 2 ** t

    # Ancilla state: |+>^t evolved under controlled-U^(2^k)
    anc = np.zeros(dim, dtype=np.complex128)
    for k in range(dim):
        phase_k  = (k * true_phase) % 1.0
        anc[k]   = np.exp(2j * np.pi * phase_k) / np.sqrt(dim)

    # Inverse QFT on ancilla
    N    = dim
    iqft = np.conj(np.exp(2j * np.pi * np.outer(np.arange(N), np.arange(N)) / N)).T / np.sqrt(N)
    out  = iqft @ anc
    probs = np.abs(out) ** 2

    meas      = int(np.argmax(probs))
    est_phase = meas / dim
    return {
        "estimated_phase":  est_phase,
        "true_phase":       true_phase,
        "error":            abs(est_phase - true_phase),
        "probabilities":    probs.tolist(),
        "n_ancilla":        n_ancilla,
        "method":           "Textbook-QPE",
    }

# ─────────────────────────────────────────────────────────────────────────────
#  2. Iterative QPE  (1 ancilla, bit-by-bit, Kitaev protocol)
# ─────────────────────────────────────────────────────────────────────────────

def iterative_qpe(U:         np.ndarray,
                  eigenstate: np.ndarray,
                  n_bits:     int = 8) -> dict:
    """
    Iterative (Kitaev) QPE — resolves one bit per circuit execution.
    Uses only 1 ancilla qubit; scales to large U without full ancilla register.
    """
    true_phase, _ = _dominant_eigenphase(U, eigenstate)
    bits    = []
    omega   = 0.0
    for k in range(n_bits - 1, -1, -1):
        M_k        = 2 ** k
        eff_phase  = (M_k * true_phase - omega) % 1.0
        prob_one   = np.sin(np.pi * eff_phase) ** 2
        # Simulate single-shot measurement
        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,
        "method":           "Iterative-QPE",
    }

# ─────────────────────────────────────────────────────────────────────────────
#  3. Hadamard-Test QPE
#  Estimates Re(<psi|U^k|psi>) and Im(<psi|U^k|psi>) for increasing k,
#  then reconstructs the phase via classical post-processing (Hadamard test).
# ─────────────────────────────────────────────────────────────────────────────

def hadamard_test_qpe(U:         np.ndarray,
                      eigenstate: np.ndarray,
                      n_samples:  int = 32,
                      shots:      int = 1000) -> dict:
    """
    Hadamard-Test QPE.

    Measures Re/Im expectation values for U^k, k=0..n_samples-1,
    then fits phase via discrete Fourier analysis.
    """
    psi  = eigenstate.astype(np.complex128)
    psi /= np.linalg.norm(psi)

    re_vals, im_vals = [], []
    U_k = np.eye(U.shape[0], dtype=np.complex128)
    for _ in range(n_samples):
        expect = np.conj(psi) @ (U_k @ psi)
        # Shot noise simulation
        re_shot = float(np.real(expect)) + np.random.normal(0, 1/np.sqrt(shots))
        im_shot = float(np.imag(expect)) + np.random.normal(0, 1/np.sqrt(shots))
        re_vals.append(re_shot)
        im_vals.append(im_shot)
        U_k = U_k @ U

    # DFT to find dominant frequency
    signal   = np.array(re_vals) + 1j * np.array(im_vals)
    spectrum = np.fft.fft(signal)
    freqs    = np.fft.fftfreq(n_samples)
    dominant = int(np.argmax(np.abs(spectrum)))
    est_phase = float(freqs[dominant]) % 1.0

    true_phase, _ = _dominant_eigenphase(U, eigenstate)
    return {
        "estimated_phase":  est_phase,
        "true_phase":       true_phase,
        "error":            min(abs(est_phase - true_phase),
                                1 - abs(est_phase - true_phase)),
        "re_expectation":   re_vals,
        "im_expectation":   im_vals,
        "n_samples":        n_samples,
        "method":           "Hadamard-Test-QPE",
    }

# ─────────────────────────────────────────────────────────────────────────────
#  4. Robust QPE  (median-of-means, noise-tolerant)
#  Runs QPE n_rounds times, returns median — robust to readout errors.
# ─────────────────────────────────────────────────────────────────────────────

def robust_qpe(U:          np.ndarray,
               eigenstate:  np.ndarray,
               n_ancilla:   int   = 6,
               n_rounds:    int   = 7,
               noise_level: float = 0.0) -> dict:
    """
    Robust QPE via median-of-means over n_rounds independent QPE calls.

    noise_level : bit-flip probability on each ancilla measurement (0 = perfect)
    """
    estimates = []
    for _ in range(n_rounds):
        r = textbook_qpe(U, eigenstate, n_ancilla)
        est = r["estimated_phase"]
        # Simulate readout noise: random bit-flip on binary representation
        if noise_level > 0:
            bits = list(format(int(est * 2**n_ancilla), f"0{n_ancilla}b"))
            bits = [str(1-int(b)) if np.random.random() < noise_level else b
                    for b in bits]
            est  = int("".join(bits), 2) / 2**n_ancilla
        estimates.append(est)

    true_phase, _ = _dominant_eigenphase(U, eigenstate)
    median_est    = float(np.median(estimates))
    return {
        "estimated_phase":  median_est,
        "true_phase":       true_phase,
        "error":            abs(median_est - true_phase),
        "all_estimates":    estimates,
        "n_rounds":         n_rounds,
        "noise_level":      noise_level,
        "method":           "Robust-QPE",
    }

# ─────────────────────────────────────────────────────────────────────────────
#  5. Scaled QPE for Large Hamiltonians
#  Uses sparse eigenvalue solver (power iteration) instead of dense expm.
#  Targets the largest / smallest eigenvalue of H without storing full matrix.
# ─────────────────────────────────────────────────────────────────────────────

def scaled_qpe_large(H:          np.ndarray,
                     n_ancilla:   int   = 8,
                     t:           float = 1.0,
                     target:      str   = "ground",
                     power_iters: int   = 50) -> dict:
    """
    Scaled QPE for large (sparse or structured) Hamiltonians.

    Uses power iteration to approximate the dominant eigenstate,
    then runs textbook QPE on U = exp(-iHt) restricted to that state.

    target : "ground"  → lowest eigenvalue
             "excited" → highest eigenvalue
    """
    dim = H.shape[0]
    # Power iteration to find dominant eigenvector
    rng = np.random.default_rng(42)
    v   = rng.standard_normal(dim).astype(np.complex128)
    v  /= np.linalg.norm(v)

    if target == "ground":
        # Shift: A = lambda_max * I - H  → dominant eigenvector is ground state
        lmax = float(np.real(np.linalg.eigvalsh(H)[-1]))
        A    = lmax * np.eye(dim, dtype=np.complex128) - H.astype(np.complex128)
    else:
        A    = H.astype(np.complex128)

    for _ in range(power_iters):
        v  = A @ v
        nv = np.linalg.norm(v)
        if nv > 1e-12:
            v /= nv

    # Build U = exp(-iHt) (dense, memory guarded)
    assert dim <= 2**15, f"H too large for dense expm ({dim}x{dim}); use sparse methods."
    U = expm(-1j * H.astype(np.complex128) * t).astype(_C)

    result = textbook_qpe(U.astype(np.complex128), v, n_ancilla)
    # Convert phase to energy: E = 2*pi*phase / t
    result["estimated_energy"] = result["estimated_phase"] * 2 * np.pi / t
    evals_true = np.sort(np.real(np.linalg.eigvalsh(H)))
    result["true_ground_energy"] = float(evals_true[0])
    result["true_excited_energy"] = float(evals_true[-1])
    result["method"] = "Scaled-QPE"
    return result

# ─────────────────────────────────────────────────────────────────────────────
#  6. HamiltonianPhaseEstimation  (Qiskit-style wrapper API)
#  Mirrors: qiskit_algorithms.HamiltonianPhaseEstimation interface
# ─────────────────────────────────────────════════════════════════════════════

class HamiltonianPhaseEstimation:
    """
    PKTron wrapper that mirrors the Qiskit HamiltonianPhaseEstimation API.

    Usage
    -----
    hpe = HamiltonianPhaseEstimation(n_ancilla=6)
    result = hpe.estimate(H, eigenstate, t=1.0, mode="iterative")
    print(result["estimated_energy"])
    """

    def __init__(self, n_ancilla: int = 6):
        self.n_ancilla = n_ancilla

    def estimate(self,
                 H:          np.ndarray,
                 eigenstate:  np.ndarray,
                 t:           float = 1.0,
                 mode:        str   = "textbook") -> dict:
        """
        Parameters
        ----------
        H          : Hamiltonian matrix
        eigenstate : initial state (will be auto-normalized)
        t          : evolution time
        mode       : "textbook" | "iterative" | "hadamard" | "robust" | "scaled"

        Returns
        -------
        dict including estimated_energy, true_ground_energy, and phase results
        """
        psi = eigenstate.astype(np.complex128)
        psi /= np.linalg.norm(psi)
        U   = expm(-1j * H.astype(np.complex128) * t).astype(np.complex128)

        if mode == "iterative":
            r = iterative_qpe(U, psi, self.n_ancilla)
        elif mode == "hadamard":
            r = hadamard_test_qpe(U, psi)
        elif mode == "robust":
            r = robust_qpe(U, psi, self.n_ancilla)
        elif mode == "scaled":
            r = scaled_qpe_large(H, self.n_ancilla, t)
        else:
            r = textbook_qpe(U, psi, self.n_ancilla)

        evals = np.sort(np.real(np.linalg.eigvalsh(H)))
        r["estimated_energy"]    = r["estimated_phase"] * 2 * np.pi / t
        r["true_ground_energy"]  = float(evals[0])
        r["hpe_mode"]            = mode
        return r
