
import numpy as np
import math

# ─────────────────────────────────────────────────────────────────
#  PKTron Quantum PCA (qPCA)
#
#  What is qPCA?
#  ─────────────
#  Classical PCA finds the principal components (directions of
#  maximum variance) of a dataset by computing eigenvectors of
#  the covariance matrix.
#
#  Quantum PCA (Lloyd, Mohseni & Rebentrost 2014) speeds this up
#  by encoding the density matrix rho = X X^T / tr(X X^T) as a
#  quantum state and using quantum phase estimation to read out
#  the eigenvalues (principal components) exponentially faster
#  for certain data structures.
#
#  Pipeline (5 stages)
#  ────────────────────
#  1. DATA ENCODING      encode data matrix X as quantum state rho
#  2. DENSITY MATRIX     build rho = X X^T / tr(X X^T)
#  3. QUANTUM PE         simulate QPE on rho to get eigenvalues
#  4. EIGENVECTORS       recover principal components
#  5. PROJECTION         project data onto top-k components
#
#  Memory safety
#  ─────────────
#  All arrays are at most (n_features x n_features).
#  Hard cap: n_features <= 8, n_samples <= 64.
#  Peak RAM across all functions: well under 1 MB.
# ─────────────────────────────────────────────────────────────────


# ═══════════════════════════════════════════════════════════════
#  Result containers
# ═══════════════════════════════════════════════════════════════

class QPCAResult:
    """
    Result from qpca_fit().

    Attributes
    ----------
    eigenvalues       : np.ndarray  sorted descending eigenvalues of rho
    eigenvectors      : np.ndarray  columns are principal components
    explained_ratio   : np.ndarray  fraction of variance per component
    components        : np.ndarray  top-k eigenvectors (shape k x n_features)
    X_projected       : np.ndarray  data projected onto top-k components
    X_reconstructed   : np.ndarray  reconstruction from top-k components
    reconstruction_err: float       ||X - X_reconstructed||_F
    circuit           : object      symbolic QPE QuantumCircuit
    steps             : list        human-readable log
    """
    def __init__(self, eigenvalues, eigenvectors, explained_ratio,
                 components, X_proj, X_recon, recon_err, circuit, steps):
        self.eigenvalues        = eigenvalues
        self.eigenvectors       = eigenvectors
        self.explained_ratio    = explained_ratio
        self.components         = components
        self.X_projected        = X_proj
        self.X_reconstructed    = X_recon
        self.reconstruction_err = recon_err
        self.circuit            = circuit
        self.steps              = steps

    def summary(self):
        sep = "=" * 56
        lines = [sep, "  Quantum PCA — Results", sep]
        for s in self.steps:
            lines.append("  " + s)
        lines.append("")
        lines.append(f"  Top eigenvalues    : "
                     f"{np.round(self.eigenvalues[:4], 6).tolist()}")
        lines.append(f"  Explained variance : "
                     f"{np.round(self.explained_ratio[:4] * 100, 2).tolist()} %")
        lines.append(f"  Projected shape    : {self.X_projected.shape}")
        lines.append(f"  Reconstruction err : {self.reconstruction_err:.6e}")
        lines.append(sep)
        return "\n".join(lines)

    def __repr__(self):
        return (f"<QPCAResult components={self.components.shape} "
                f"recon_err={self.reconstruction_err:.2e}>")


class QKernelPCAResult:
    """Result from quantum_kernel_pca()."""
    def __init__(self, K, eigenvalues, X_projected, steps):
        self.K           = K
        self.eigenvalues = eigenvalues
        self.X_projected = X_projected
        self.steps       = steps

    def summary(self):
        sep = "=" * 56
        lines = [sep, "  Quantum Kernel PCA — Results", sep]
        for s in self.steps:
            lines.append("  " + s)
        lines.append("")
        lines.append(f"  Kernel matrix shape : {self.K.shape}")
        lines.append(f"  Top eigenvalues     : "
                     f"{np.round(self.eigenvalues[:4], 4).tolist()}")
        lines.append(f"  Projected shape     : {self.X_projected.shape}")
        lines.append(sep)
        return "\n".join(lines)


# ═══════════════════════════════════════════════════════════════
#  Internal helpers
# ═══════════════════════════════════════════════════════════════

def _validate_data(X, max_features=8, max_samples=64):
    """Validate and clip data matrix X (n_samples x n_features)."""
    X = np.array(X, dtype=float)
    assert X.ndim == 2, "X must be 2-D (n_samples x n_features)"
    n_samples, n_features = X.shape
    assert n_features <= max_features, (
        f"n_features={n_features} exceeds max {max_features} "
        f"(12 GB RAM guard)"
    )
    assert n_samples <= max_samples, (
        f"n_samples={n_samples} exceeds max {max_samples} "
        f"(12 GB RAM guard)"
    )
    return X


def _centre(X):
    """Mean-centre X column-wise."""
    return X - X.mean(axis=0)


def _build_qpe_circuit(n_features, n_clock):
    """
    Build a symbolic QPE circuit for qPCA.
    Hard cap: total qubits <= 10 (16 KB statevector).

    Layout
    ------
    [0 .. n_clock-1]  clock register
    [n_clock ..]      density-matrix register (log2(n_features) qubits)
    """
    from .quantum import QuantumCircuit
    n_dm    = max(1, math.ceil(math.log2(max(n_features, 2))))
    n_clock = min(n_clock, 4)
    total   = min(n_clock + n_dm, 10)
    qc      = QuantumCircuit(total)

    # Clock register — Hadamard
    for q in range(n_clock):
        qc.h(q)

    # Encode density matrix rho (symbolic)
    qc.gate_history.append("Encode rho = X X^T / tr(X X^T) on DM register")

    # Controlled e^{i rho t 2^k} gates
    for k in range(n_clock):
        qc.gate_history.append(
            f"Controlled-e^{{i*rho*2^{k}*t}} | ctrl={k} dm={n_clock}..{total-1}"
        )

    # Inverse QFT on clock register
    for j in range(n_clock // 2):
        qc.swap(j, n_clock - 1 - j)
    for j in range(n_clock):
        qc.h(j)
        for kk in range(j + 1, n_clock):
            angle = -math.pi / (2 ** (kk - j))
            qc.rz(angle, j)

    qc.gate_history.append("Measure clock register -> eigenvalue readout")
    return qc


def _density_matrix(X_centred):
    """
    Build normalised density matrix rho = X^T X / tr(X^T X).
    Shape: (n_features x n_features). Always tiny.
    """
    cov = X_centred.T @ X_centred          # (n_features x n_features)
    tr  = np.trace(cov)
    if tr < 1e-12:
        tr = 1.0
    return cov / tr, tr


# ═══════════════════════════════════════════════════════════════
#  Core: qpca_fit
# ═══════════════════════════════════════════════════════════════

def qpca_fit(X, n_components=2, n_clock_bits=3, verbose=True):
    """
    Fit Quantum PCA on dataset X.

    Parameters
    ----------
    X            : array (n_samples x n_features), max 64 x 8
    n_components : int  — number of principal components to keep
    n_clock_bits : int  — QPE precision (default 3, capped at 4)
    verbose      : bool

    Returns
    -------
    QPCAResult
    """
    steps = []
    def log(msg):
        steps.append(msg)
        if verbose: print("  " + msg)

    X = _validate_data(X)
    n_samples, n_features = X.shape
    n_components = min(n_components, n_features)
    log(f"Data: {n_samples} samples x {n_features} features")
    log(f"Keeping top {n_components} components")

    # ── Stage 1: centre data ─────────────────────────────────
    Xc = _centre(X)
    log("Data mean-centred")

    # ── Stage 2: build density matrix rho ───────────────────
    rho, tr = _density_matrix(Xc)
    log(f"Density matrix rho built ({n_features}x{n_features}), "
        f"tr(X^T X) = {tr:.4f}")

    # ── Stage 3: QPE — eigendecomposition of rho ─────────────
    # On real hardware: QPE reads eigenvalues via phase kickback.
    # Here: exact numpy eigendecomposition (simulates QPE output).
    eigenvalues, eigenvectors = np.linalg.eigh(rho)
    # eigh returns ascending order — reverse to descending
    idx          = np.argsort(eigenvalues)[::-1]
    eigenvalues  = eigenvalues[idx]
    eigenvectors = eigenvectors[:, idx]
    log(f"QPE complete — top eigenvalues: "
        f"{np.round(eigenvalues[:n_components], 6).tolist()}")

    # ── Stage 4: explained variance ratio ───────────────────
    total_var      = eigenvalues.sum()
    explained      = eigenvalues / (total_var + 1e-12)
    cumulative     = np.cumsum(explained)
    log(f"Explained variance (top {n_components}): "
        f"{np.round(explained[:n_components] * 100, 2).tolist()} %")
    log(f"Cumulative variance (top {n_components}): "
        f"{round(float(cumulative[n_components-1]) * 100, 2)} %")

    # ── Stage 5: project & reconstruct ──────────────────────
    components  = eigenvectors[:, :n_components].T   # (k x n_features)
    X_projected = Xc @ eigenvectors[:, :n_components]  # (n_samples x k)
    X_recon     = X_projected @ components + X.mean(axis=0)
    recon_err   = float(np.linalg.norm(X - X_recon, "fro"))
    log(f"Projected to {X_projected.shape}  "
        f"reconstruction error = {recon_err:.4e}")

    # ── Stage 6: build symbolic circuit ─────────────────────
    circuit = _build_qpe_circuit(n_features, n_clock_bits)
    log(f"QPE circuit: {circuit.num_qubits} qubits, "
        f"{len(circuit.gate_history)} gates (symbolic)")

    return QPCAResult(
        eigenvalues, eigenvectors, explained,
        components, X_projected, X_recon,
        recon_err, circuit, steps
    )


# ═══════════════════════════════════════════════════════════════
#  Incremental / streaming qPCA
# ═══════════════════════════════════════════════════════════════

def incremental_qpca(X_batches, n_components=2, verbose=True):
    """
    Incremental Quantum PCA — processes data in small batches.
    Useful when full data does not fit in memory at once.

    Each batch updates a running covariance estimate using
    Welford\'s online algorithm, then applies qPCA at the end.

    Parameters
    ----------
    X_batches    : list of arrays, each (batch_size x n_features)
                   max n_features=8, each batch max 64 rows
    n_components : int
    verbose      : bool

    Returns
    -------
    QPCAResult  (fitted on aggregated covariance)
    """
    steps = []
    def log(msg):
        steps.append(msg)
        if verbose: print("  " + msg)

    assert len(X_batches) > 0, "Need at least one batch"
    n_features = X_batches[0].shape[1]
    assert n_features <= 8, "n_features must be <= 8"

    # Welford online mean + covariance
    n_total = 0
    mean    = np.zeros(n_features)
    M2      = np.zeros((n_features, n_features))   # sum of squared deviations

    for batch_idx, Xb in enumerate(X_batches):
        Xb = np.array(Xb, dtype=float)
        assert Xb.shape[1] == n_features, "All batches must have same n_features"
        for row in Xb:
            n_total += 1
            delta    = row - mean
            mean    += delta / n_total
            delta2   = row - mean
            M2      += np.outer(delta, delta2)
        log(f"Batch {batch_idx+1}/{len(X_batches)}: "
            f"{Xb.shape[0]} rows processed (total {n_total})")

    cov = M2 / max(n_total - 1, 1)
    log(f"Aggregated covariance ({n_features}x{n_features}) from {n_total} rows")

    # Assemble a synthetic centred dataset from the covariance
    # (Cholesky factor gives us a matrix whose covariance matches)
    try:
        L      = np.linalg.cholesky(cov + 1e-9 * np.eye(n_features))
        X_synth = (L @ np.random.default_rng(0).standard_normal(
                        (n_features, n_total))).T
    except np.linalg.LinAlgError:
        X_synth = np.random.default_rng(0).standard_normal((n_total, n_features))

    return qpca_fit(X_synth, n_components=n_components,
                    verbose=False)


# ═══════════════════════════════════════════════════════════════
#  Quantum Kernel PCA
# ═══════════════════════════════════════════════════════════════

def quantum_kernel_pca(X, n_components=2,
                       kernel="rbf", gamma=1.0, verbose=True):
    """
    Quantum Kernel PCA.

    Replaces the linear covariance with a quantum-inspired kernel
    matrix K[i,j] = |<phi(x_i)|phi(x_j)>|^2, then applies
    kernel PCA in that feature space.

    Supported kernels
    -----------------
    "rbf"      : K[i,j] = exp(-gamma * ||x_i - x_j||^2)
    "quantum"  : K[i,j] = cos^2(gamma * ||x_i - x_j||)
                 (mimics overlap of two Z-rotation feature maps)
    "poly"     : K[i,j] = (x_i . x_j + 1)^2

    Parameters
    ----------
    X            : array (n_samples x n_features), max 64 x 8
    n_components : int
    kernel       : str  ("rbf", "quantum", "poly")
    gamma        : float
    verbose      : bool

    Returns
    -------
    QKernelPCAResult
    """
    steps = []
    def log(msg):
        steps.append(msg)
        if verbose: print("  " + msg)

    X = _validate_data(X)
    n_samples, n_features = X.shape
    n_components = min(n_components, n_samples)
    log(f"Quantum Kernel PCA: {n_samples} samples, {n_features} features")
    log(f"Kernel: {kernel}  gamma={gamma}")

    # ── Build kernel matrix ──────────────────────────────────
    K = np.zeros((n_samples, n_samples))
    for i in range(n_samples):
        for j in range(i, n_samples):
            diff = X[i] - X[j]
            d2   = diff @ diff
            if kernel == "rbf":
                v = math.exp(-gamma * d2)
            elif kernel == "quantum":
                v = math.cos(gamma * math.sqrt(d2 + 1e-12)) ** 2
            elif kernel == "poly":
                v = (X[i] @ X[j] + 1.0) ** 2
            else:
                raise ValueError(f"Unknown kernel: {kernel}")
            K[i, j] = K[j, i] = v
    log(f"Kernel matrix built ({n_samples}x{n_samples})")

    # ── Centre kernel matrix ─────────────────────────────────
    one_n = np.ones((n_samples, n_samples)) / n_samples
    Kc    = K - one_n @ K - K @ one_n + one_n @ K @ one_n

    # ── Eigendecomposition ───────────────────────────────────
    eigenvalues, eigenvectors = np.linalg.eigh(Kc)
    idx         = np.argsort(eigenvalues)[::-1]
    eigenvalues = eigenvalues[idx]
    eigenvectors= eigenvectors[:, idx]

    # Keep only positive eigenvalues
    pos_mask    = eigenvalues > 1e-10
    eigenvalues = eigenvalues[pos_mask]
    eigenvectors= eigenvectors[:, pos_mask]

    log(f"Top eigenvalues: {np.round(eigenvalues[:n_components], 4).tolist()}")

    # ── Project ──────────────────────────────────────────────
    k = min(n_components, eigenvectors.shape[1])
    X_projected = eigenvectors[:, :k] * np.sqrt(
        np.maximum(eigenvalues[:k], 0)
    )
    log(f"Projected to shape {X_projected.shape}")

    return QKernelPCAResult(K, eigenvalues, X_projected, steps)


# ═══════════════════════════════════════════════════════════════
#  Quantum Randomised SVD  (memory-safe sketch-based SVD)
# ═══════════════════════════════════════════════════════════════

def quantum_randomised_svd(X, rank=2, n_oversamples=2, verbose=True):
    """
    Quantum-inspired Randomised SVD.

    Uses random Gaussian sketching (as in quantum-inspired
    low-rank approximation algorithms) to compute a rank-k
    truncated SVD without forming the full SVD.

    Parameters
    ----------
    X           : array (n_samples x n_features), max 64 x 8
    rank        : int — target rank
    n_oversamples: int — extra samples for accuracy (default 2)
    verbose     : bool

    Returns
    -------
    dict with keys: U, S, Vt, X_approx, error
    """
    steps = []
    def log(msg):
        steps.append(msg)
        if verbose: print("  " + msg)

    X = _validate_data(X)
    n_samples, n_features = X.shape
    k = min(rank + n_oversamples, n_features, n_samples)
    log(f"Randomised SVD: {n_samples}x{n_features}, target rank={rank}")

    # Stage 1: random sketch — Omega is a Gaussian random matrix
    rng   = np.random.default_rng(42)
    Omega = rng.standard_normal((n_features, k))
    Y     = X @ Omega                             # (n_samples x k)
    log(f"Random sketch: {Y.shape}")

    # Stage 2: QR decomposition of sketch
    Q, _ = np.linalg.qr(Y)                       # (n_samples x k)
    log(f"QR decomposition: Q shape {Q.shape}")

    # Stage 3: project X into low-dimensional space
    B = Q.T @ X                                   # (k x n_features)
    log(f"Projected matrix B: {B.shape}")

    # Stage 4: SVD of small matrix B
    U_hat, S, Vt = np.linalg.svd(B, full_matrices=False)
    U = Q @ U_hat                                 # (n_samples x k)

    # Truncate to target rank
    r    = min(rank, len(S))
    U, S, Vt = U[:, :r], S[:r], Vt[:r, :]
    log(f"Singular values (top {r}): {np.round(S, 4).tolist()}")

    # Stage 5: low-rank approximation
    X_approx = U @ np.diag(S) @ Vt
    error    = float(np.linalg.norm(X - X_approx, "fro"))
    log(f"Approximation error ||X - X_approx||_F = {error:.4e}")

    return {
        "U"       : U,
        "S"       : S,
        "Vt"      : Vt,
        "X_approx": X_approx,
        "error"   : error,
        "steps"   : steps,
    }
