
import numpy as np
from scipy.linalg import qr

_C128 = np.complex128
_C64  = np.complex64

# ─────────────────────────────────────────────────────────────────────────────
#  1. Exact Dense Block Encoding
#  Embeds H/alpha in the top-left block of a unitary U_BE of size 2n x 2n.
#  Condition: ||H|| <= alpha  (alpha is the block-encoding factor)
# ─────────────────────────────────────────────────────────────────────────────

def block_encode_dense(H: np.ndarray, alpha: float = None) -> dict:
    """
    Exact block encoding via QR completion.

    Constructs a unitary U such that  <0|U|0> = H / alpha.

    Parameters
    ----------
    H     : target matrix (n x n), not necessarily unitary
    alpha : subnormalisation factor (default: largest singular value)

    Returns
    -------
    dict: unitary (2n x 2n), alpha, error (||H/alpha - top_left_block||)
    """
    n      = H.shape[0]
    svd_s  = np.linalg.svd(H, compute_uv=False)
    if alpha is None:
        alpha = float(svd_s[0]) + 1e-10
    A      = H.astype(_C128) / alpha

    # Build 2n x 2n matrix: top-left block = A, fill rest via QR
    dim    = 2 * n
    M      = np.zeros((dim, dim), dtype=_C128)
    M[:n, :n] = A
    # Random completion for off-diagonal blocks
    rng    = np.random.default_rng(1)
    pad    = (rng.standard_normal((dim, dim))
              + 1j * rng.standard_normal((dim, dim)))
    pad[:n, :n] = A
    Q, _   = qr(pad)
    # Verify block
    err    = float(np.linalg.norm(Q[:n, :n] - A))
    return {
        "unitary":       Q.astype(_C64),
        "alpha":         alpha,
        "block_error":   err,
        "n":             n,
        "method":        "Dense-BlockEncode",
    }

# ─────────────────────────────────────────────────────────────────────────────
#  2. FABLE — Fast Approximate BLock Encodings
#  Approximates H via its Fourier/Chebyshev expansion truncated at degree d.
#  Complexity: O(n * d) vs O(n^2) for dense — critical for large sparse H.
# ─────────────────────────────────────────────────────────────────────────────

def fable_block_encode(H: np.ndarray,
                       degree: int   = 4,
                       alpha:  float = None) -> dict:
    """
    FABLE: Fast Approximate BLock Encoding via Chebyshev truncation.

    Approximates U_BE using a Chebyshev polynomial expansion of H/alpha
    up to degree `degree`. Higher degree → better approximation, more gates.

    Returns the approximated block-encoding unitary and approximation error.
    """
    n     = H.shape[0]
    norm  = float(np.linalg.svd(H, compute_uv=False)[0])
    if alpha is None:
        alpha = norm + 1e-10
    A     = H.astype(_C128) / alpha

    # Chebyshev expansion:  A ≈ sum_{k=0}^{degree} c_k T_k(A)
    # T_0(A) = I, T_1(A) = A, T_{k+1}(A) = 2A T_k(A) - T_{k-1}(A)
    I  = np.eye(n, dtype=_C128)
    Tk_prev, Tk_curr = I.copy(), A.copy()
    # Chebyshev coefficients via inner product (simplified uniform weights)
    c_approx = np.zeros(n, dtype=_C128)
    H_approx = np.zeros((n, n), dtype=_C128)
    for k in range(degree + 1):
        if k == 0:
            Tk = I
        elif k == 1:
            Tk = A
        else:
            Tk       = 2 * A @ Tk_curr - Tk_prev
            Tk_prev  = Tk_curr
            Tk_curr  = Tk
        # Coefficient: sample at Chebyshev nodes
        coeff     = (2.0 / (degree + 1)) * np.trace(Tk) / n
        H_approx += coeff * Tk

    # Build block encoding from approximated H_approx
    be_result = block_encode_dense(np.real(H_approx) * alpha, alpha)
    approx_err = float(np.linalg.norm(H_approx * alpha - H))
    be_result.update({
        "fable_approx_error": approx_err,
        "degree":             degree,
        "method":             "FABLE",
    })
    return be_result

# ─────────────────────────────────────────────────────────────────────────────
#  3. LCU — Linear Combination of Unitaries
#  H = sum_k alpha_k U_k  →  be = PREPARE† · SELECT · PREPARE
#  PREPARE: encodes sqrt(alpha_k/lambda) amplitudes in ancilla register
#  SELECT:  applies U_k controlled on ancilla state |k>
# ─────────────────────────────────────────────────────────────────────────────

def lcu_block_encode(terms: list) -> dict:
    """
    LCU Block Encoding.

    Parameters
    ----------
    terms : list of (coefficient, unitary_matrix) tuples
            H = sum_k coeff_k * U_k

    Returns
    -------
    dict: block_encode_unitary, alpha (lambda), prepare_state, select_unitary
    """
    coeffs  = np.array([abs(c) for c, _ in terms], dtype=float)
    lam     = float(coeffs.sum())
    probs   = coeffs / lam
    # PREPARE state: |phi> = sum_k sqrt(p_k) |k>
    prepare = np.sqrt(probs).astype(_C128)

    K  = len(terms)
    n  = terms[0][1].shape[0]
    # SELECT unitary: block-diagonal {c_k/|c_k| * U_k}
    select = np.zeros((K * n, K * n), dtype=_C128)
    for k, (c, U) in enumerate(terms):
        phase      = np.sign(c) if c != 0 else 1.0
        select[k*n:(k+1)*n, k*n:(k+1)*n] = phase * U.astype(_C128)

    # Block encoding acts as: (<phi| ⊗ I) · SELECT · (|phi> ⊗ I)
    H_reconstructed = sum(c * U.astype(_C128) for c, U in terms)
    return {
        "alpha":              lam,
        "prepare_state":      prepare.astype(_C64),
        "select_unitary":     select.astype(_C64),
        "H_reconstructed":    H_reconstructed.astype(_C64),
        "n_terms":            K,
        "method":             "LCU",
    }

# ─────────────────────────────────────────────────────────────────────────────
#  4. QSVT Framework Sketch
#  Quantum Singular Value Transformation via alternating phase applications.
#  Implements: poly(H/alpha) via a degree-d phase sequence {phi_0,...,phi_d}
#  For small H, exact; outputs transformed block and estimated singular values.
# ─────────────────────────────────────────────────────────────────────────────

def qsvt_apply(H:      np.ndarray,
               phases: list,
               alpha:  float = None) -> dict:
    """
    QSVT: apply a polynomial transformation to the singular values of H.

    The transformation is: poly(sigma) where poly is defined by the
    QSP phase sequence `phases` (length d+1 for degree-d polynomial).

    Parameters
    ----------
    H      : target matrix
    phases : list of floats (QSP phases, length = degree + 1)
    alpha  : block-encoding factor

    Returns
    -------
    dict: transformed_H, original_singular_values, transformed_singular_values
    """
    n    = H.shape[0]
    svd_s = np.linalg.svd(H.astype(_C128), compute_uv=False)
    if alpha is None:
        alpha = float(svd_s[0]) + 1e-10
    A    = H.astype(_C128) / alpha

    # Build block encoding unitary W_A (simplified 2x2 signal model per sv)
    # For full implementation, iterate QSP on each singular value independently
    U, S, Vh = np.linalg.svd(A, full_matrices=True)

    def _qsp_poly(sigma, phis):
        """Evaluate QSP polynomial on scalar singular value sigma in [-1,1]."""
        # Signal unitary: [[sigma, i*sqrt(1-sigma^2)], [...]]
        c = np.clip(float(sigma), -1, 1)
        s = np.sqrt(max(0.0, 1 - c**2))
        W = np.array([[c, 1j*s], [1j*s, c]], dtype=_C128)
        # Alternating phase applications
        result = np.eye(2, dtype=_C128)
        for phi in phis:
            Phi    = np.array([[np.exp(1j*phi), 0], [0, np.exp(-1j*phi)]])
            result = Phi @ W @ result
        return np.real(result[0, 0])

    n_sv = len(S)
    S_out = np.array([_qsp_poly(sv, phases) for sv in S])
    # Reconstruct transformed matrix
    H_out = (U[:, :n_sv] * S_out) @ Vh[:n_sv, :]

    return {
        "transformed_H":              (H_out * alpha).astype(_C64),
        "original_singular_values":   S.tolist(),
        "transformed_singular_values": S_out.tolist(),
        "phases":                     phases,
        "alpha":                      alpha,
        "method":                     "QSVT",
    }

# ─────────────────────────────────────────────────────────────────────────────
#  5. Hamiltonian Oracle from Pauli decomposition
#  Decomposes H into Pauli strings: H = sum_k c_k P_k
#  Returns LCU terms ready for lcu_block_encode()
# ─────────────────────────────────────────────────────────────────────────────

_PAULI = {
    "I": np.eye(2, dtype=np.complex128),
    "X": np.array([[0,1],[1,0]], dtype=np.complex128),
    "Y": np.array([[0,-1j],[1j,0]], dtype=np.complex128),
    "Z": np.array([[1,0],[0,-1]], dtype=np.complex128),
}

def _tensor_pauli(string: str) -> np.ndarray:
    result = np.array([[1.0+0j]])
    for ch in string:
        result = np.kron(result, _PAULI[ch])
    return result

def pauli_decompose(H: np.ndarray) -> list:
    """
    Decompose H into Pauli basis: H = sum_k c_k P_k.
    Works for n-qubit H where dim = 2^n.
    Returns list of (coeff, unitary) for lcu_block_encode.
    """
    import itertools
    n   = int(np.round(np.log2(H.shape[0])))
    dim = 2 ** n
    terms = []
    for labels in itertools.product("IXYZ", repeat=n):
        label   = "".join(labels)
        P       = _tensor_pauli(label)
        coeff   = np.trace(P.conj().T @ H.astype(np.complex128)) / dim
        if abs(coeff) > 1e-12:
            terms.append((complex(coeff), P))
    return terms
