
import math
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron HHL Algorithm & Quantum Linear System Solvers
#  Memory safety: all matrices max 8x8, circuits max 10 qubits.
# ─────────────────────────────────────────────────────────────────


# ═══════════════════════════════════════════════════════════════
#  Validation
# ═══════════════════════════════════════════════════════════════

def _validate(A, b):
    A = np.array(A, dtype=complex)
    b = np.array(b, dtype=complex)
    assert A.ndim == 2,                        "A must be 2-D"
    assert A.shape[0] == A.shape[1],           "A must be square"
    assert b.ndim == 1,                        "b must be 1-D"
    assert b.shape[0] == A.shape[0],           "b length must match A dimension"
    assert A.shape[0] <= 8,                    f"A is {A.shape[0]}x{A.shape[0]} — max 8x8 for 12 GB safety"
    assert np.allclose(A, A.conj().T, atol=1e-8), "A must be Hermitian"
    assert not np.allclose(b, 0),              "b must be non-zero"
    return A, b


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

class HHLResult:
    def __init__(self, A, b, x_q, x_c, residual, rel_err,
                 eigenvalues, eigenvectors, kappa, C,
                 b_coeffs, circuit, n_clock, sp_angles, steps):
        self.A = A;               self.b = b
        self.x_quantum = x_q;    self.x_classical = x_c
        self.residual = residual; self.relative_error = rel_err
        self.eigenvalues = eigenvalues; self.eigenvectors = eigenvectors
        self.condition_number = kappa;  self.C = C
        self.b_coeffs = b_coeffs;      self.circuit = circuit
        self.n_clock = n_clock;         self.state_prep_angles = sp_angles
        self.steps = steps;             self.success = residual < 1e-6

    def summary(self):
        sep = "=" * 60
        lines = [sep, "  HHL — Quantum Linear System Solver", sep]
        n = self.A.shape[0]
        lines += [
            f"  System size      : {n} x {n}",
            f"  Condition number : {self.condition_number:.4f}",
            f"  Eigenvalues      : {np.round(self.eigenvalues, 6).tolist()}",
            f"  Normalisation C  : {self.C:.6f}",
            f"  Clock qubits     : {self.n_clock}",
            f"  Circuit qubits   : {self.circuit.num_qubits}",
            f"  Circuit gates    : {len(self.circuit.gate_history)}",
            "",
            f"  b (input)        : {np.round(self.b.real, 6).tolist()}",
            f"  x_classical      : {np.round(self.x_classical, 6).tolist()}",
            f"  x_quantum        : {np.round(self.x_quantum, 6).tolist()}",
            f"  Residual ||Ax-b||: {self.residual:.4e}",
            f"  Relative error   : {self.relative_error:.4e}",
            f"  Status           : {'PASSED' if self.success else 'HIGH RESIDUAL'}",
        ]
        lines.append(sep)
        return "\n".join(lines)

    def __repr__(self):
        return (f"<HHLResult {self.A.shape[0]}x{self.A.shape[0]} "
                f"residual={self.residual:.2e} success={self.success}>")


class QLSResult:
    def __init__(self, method, A, b, x_quantum, x_classical,
                 residual, rel_err, steps):
        self.method = method;           self.A = A; self.b = b
        self.x_quantum = x_quantum;    self.x_classical = x_classical
        self.residual = residual;       self.relative_error = rel_err
        self.steps = steps
        # VQLS is iterative/approximate — looser threshold
        tol = 0.05 if method == "VQLS" else 1e-6
        self.success = residual < tol

    def summary(self):
        sep = "=" * 60
        lines = [sep, f"  QLS Solver: {self.method}", sep]
        lines += [
            f"  b (input)      : {np.round(self.b.real, 6).tolist()}",
            f"  x_classical    : {np.round(self.x_classical, 6).tolist()}",
            f"  x_quantum      : {np.round(self.x_quantum, 6).tolist()}",
            f"  Residual       : {self.residual:.4e}",
            f"  Relative error : {self.relative_error:.4e}",
            f"  Status         : {'PASSED' if self.success else 'FAILED'}",
        ]
        lines.append(sep)
        return "\n".join(lines)

    def __repr__(self):
        return (f"<QLSResult method={self.method} "
                f"residual={self.residual:.2e} success={self.success}>")


# ═══════════════════════════════════════════════════════════════
#  State-prep angles
# ═══════════════════════════════════════════════════════════════

def _state_prep_angles(b_norm):
    n = len(b_norm)
    angles = []
    cumulative = np.cumsum(np.abs(b_norm) ** 2)
    for i in range(n - 1):
        denom = 1.0 - (cumulative[i - 1] if i > 0 else 0.0)
        if denom < 1e-12:
            angles.append(0.0)
        else:
            ratio = min(abs(b_norm[i]) ** 2 / denom, 1.0)
            angles.append(2 * math.acos(math.sqrt(max(1 - ratio, 0))))
    angles.append(0.0)
    return np.array(angles)


# ═══════════════════════════════════════════════════════════════
#  Circuit builder (symbolic, max 10 qubits)
# ═══════════════════════════════════════════════════════════════

def _build_hhl_circuit(A, b_state, eigenvalues, C, n_clock):
    from .quantum import QuantumCircuit
    n       = A.shape[0]
    n_state = max(1, math.ceil(math.log2(max(n, 2))))
    total   = min(n_clock + 1 + n_state, 10)
    ancilla = n_clock
    qc      = QuantumCircuit(total)

    # Stage 1: state prep
    qc.gate_history.append(f"StatePrep |b>  Ry angles on state={n_clock+1}..{total-1}")
    for i in range(n_clock + 1, total):
        idx   = min(i - n_clock - 1, len(b_state) - 1)
        angle = float(2 * math.asin(min(abs(b_state[idx]), 1.0)))
        qc.rz(angle, i)

    # Stage 2: QPE — H on clock
    for q in range(n_clock):
        qc.h(q)

    # Stage 2: controlled e^{iAt}
    for k in range(n_clock):
        qc.gate_history.append(
            f"Controlled-e^{{iA*{2*math.pi*(2**k):.4f}}} "
            f"ctrl={k} state={n_clock+1}..{total-1}"
        )

    # Stage 2: inverse QFT
    for i in range(n_clock // 2):
        qc.swap(i, n_clock - 1 - i)
    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(f"CRZ({angle:.4f}) ctrl={kk} target={j}  [IQFT]")

    # Stage 3: eigenvalue inversion
    for lam in eigenvalues:
        if abs(lam) > 1e-10:
            theta = 2 * math.asin(min(abs(C / lam), 1.0))
            qc.gate_history.append(
                f"Controlled-Ry({theta:.4f}) [1/lambda={1/lam:.4f}] ancilla={ancilla}"
            )
    qc.rz(0.0, ancilla)

    # Stage 4: uncompute QPE
    qc.gate_history.append("QFT on clock  [uncompute QPE]")
    for j in range(n_clock - 1, -1, -1):
        for kk in range(n_clock - 1, j, -1):
            angle = math.pi / (2 ** (kk - j))
            qc.rz(angle, j)
        qc.h(j)
    for i in range(n_clock // 2):
        qc.swap(i, n_clock - 1 - i)

    qc.gate_history.append(f"Measure ancilla={ancilla} post-select |1>")
    qc.gate_history.append("Measure state register -> x")
    return qc


# ═══════════════════════════════════════════════════════════════
#  Core: hhl_solve
# ═══════════════════════════════════════════════════════════════

def hhl_solve(A, b, n_clock=None, verbose=True):
    """
    Solve A x = b using the HHL algorithm.

    Parameters
    ----------
    A       : Hermitian array (N x N), max 8x8
    b       : array (N,)
    n_clock : QPE clock qubits (auto-sized if None, capped at 6)
    verbose : bool

    Returns
    -------
    HHLResult — .x_quantum, .x_classical, .residual, .relative_error,
                .eigenvalues, .condition_number, .circuit, .n_clock,
                .b_coeffs, .C, .state_prep_angles
    """
    steps = []
    def log(msg):
        steps.append(msg)
        if verbose: print("  " + msg)

    A, b = _validate(A, b)
    n    = A.shape[0]
    log(f"System: {n}x{n}  Hermitian")

    b_norm_val = float(np.linalg.norm(b))
    b_state    = b / b_norm_val
    sp_angles  = _state_prep_angles(b_state.real)
    log(f"||b|| = {b_norm_val:.6f}")
    log(f"State-prep Ry angles: {np.round(sp_angles, 4).tolist()}")

    eigenvalues, eigenvectors = np.linalg.eigh(A)
    log(f"Eigenvalues of A: {np.round(eigenvalues, 6).tolist()}")

    abs_ev  = np.abs(eigenvalues)
    nonzero = abs_ev[abs_ev > 1e-10]
    kappa   = float(nonzero.max() / nonzero.min()) if len(nonzero) > 1 else 1.0
    log(f"Condition number kappa = {kappa:.4f}")

    if n_clock is None:
        n_clock = min(max(3, math.ceil(math.log2(kappa + 1)) + 1), 6)
    else:
        n_clock = min(int(n_clock), 6)
    log(f"QPE clock qubits: {n_clock}  (auto-sized from kappa, capped at 6)")

    b_coeffs = eigenvectors.conj().T @ b_state
    log(f"|b> in eigenbasis: {np.round(np.abs(b_coeffs), 4).tolist()}")

    C = float(nonzero.min())
    log(f"Normalisation constant C = {C:.6f}")

    inv_coeffs = np.where(
        np.abs(eigenvalues) > 1e-10,
        b_coeffs * C / eigenvalues,
        0.0 + 0.0j
    )
    log("Eigenvalue inversion complete  (ancilla post-selection simulated)")

    x_raw     = eigenvectors @ inv_coeffs
    x_quantum = (x_raw * b_norm_val / C).real
    log(f"x_quantum (raw): {np.round(x_quantum, 6).tolist()}")

    circuit = _build_hhl_circuit(A, b_state, eigenvalues, C, n_clock)
    log(f"HHL circuit: {circuit.num_qubits} qubits, {len(circuit.gate_history)} gates")

    x_classical = np.linalg.solve(A, b).real
    residual    = float(np.linalg.norm(A @ x_quantum - b.real))
    rel_err     = float(np.linalg.norm(x_quantum - x_classical) /
                        (np.linalg.norm(x_classical) + 1e-14))
    log(f"Residual  ||Ax-b||          = {residual:.4e}")
    log(f"Relative error ||xq-xc||/||xc|| = {rel_err:.4e}")

    return HHLResult(A, b, x_quantum, x_classical, residual, rel_err,
                     eigenvalues, eigenvectors, kappa, C,
                     b_coeffs, circuit, n_clock, sp_angles, steps)


# ═══════════════════════════════════════════════════════════════
#  VQLS
# ═══════════════════════════════════════════════════════════════

def variational_qls(A, b, iterations=120, verbose=True):
    """
    Variational Quantum Linear Solver (VQLS).
    Minimises ||A|x> - |b>|| via parameter-shift gradient descent.
    Returns QLSResult (.success threshold = 0.05 residual).
    """
    steps = []
    def log(msg):
        steps.append(msg)
        if verbose: print("  " + msg)

    A, b  = _validate(A, b)
    A_r   = A.real.astype(float)
    b_r   = b.real.astype(float)
    n     = A_r.shape[0]
    log(f"VQLS: {n}x{n} system, {iterations} iterations")

    b_state = b_r / (np.linalg.norm(b_r) + 1e-14)
    rng     = np.random.default_rng(42)
    theta   = rng.uniform(0, 2 * math.pi, n)
    lr      = 0.1
    best_cost, best_theta = float("inf"), theta.copy()

    for it in range(iterations):
        x_t = np.cos(theta)
        nm  = np.linalg.norm(x_t)
        x_t = x_t / nm if nm > 1e-12 else np.ones(n) / math.sqrt(n)

        Ax   = A_r @ x_t
        num  = float(b_state @ Ax) ** 2
        den  = float(np.linalg.norm(Ax)) ** 2 * float(np.linalg.norm(b_state)) ** 2
        cost = 1.0 - num / (den + 1e-12)

        if cost < best_cost:
            best_cost, best_theta = cost, theta.copy()

        grad = np.zeros(n)
        for j in range(n):
            for sign, s in [(+1, +math.pi/2), (-1, -math.pi/2)]:
                tp = theta.copy(); tp[j] += s
                xp = np.cos(tp); xp /= (np.linalg.norm(xp) + 1e-12)
                Axp = A_r @ xp
                cp  = 1.0 - (b_state @ Axp) ** 2 / (
                      np.linalg.norm(Axp) ** 2 * np.linalg.norm(b_state) ** 2 + 1e-12)
                grad[j] += sign * cp / 2
        theta -= lr * grad

        if (it + 1) % 30 == 0:
            log(f"  iter {it+1:3d}/{iterations}  cost={cost:.6f}")

    x_q  = np.cos(best_theta)
    x_c  = np.linalg.solve(A_r, b_r)
    scale = np.dot(x_q, x_c) / (np.dot(x_q, x_q) + 1e-14)
    x_q  *= scale

    residual = float(np.linalg.norm(A_r @ x_q - b_r))
    rel_err  = float(np.linalg.norm(x_q - x_c) / (np.linalg.norm(x_c) + 1e-14))
    log(f"Final cost={best_cost:.6f}  residual={residual:.4e}  rel_err={rel_err:.4e}")
    return QLSResult("VQLS", A, b, x_q, x_c, residual, rel_err, steps)


# ═══════════════════════════════════════════════════════════════
#  Quantum Conjugate Gradient
# ═══════════════════════════════════════════════════════════════

def quantum_conjugate_gradient(A, b, tol=1e-10, max_iter=200, verbose=True):
    """
    Quantum-inspired Conjugate Gradient. Exact for SPD A.
    Converges in at most N iterations.
    """
    steps = []
    def log(msg):
        steps.append(msg)
        if verbose: print("  " + msg)

    A, b = _validate(A, b)
    A_r  = A.real.astype(float)
    b_r  = b.real.astype(float)
    n    = A_r.shape[0]
    log(f"QCG: {n}x{n} system  tol={tol}")

    x      = np.zeros(n)
    r      = b_r - A_r @ x
    p      = r.copy()
    rs_old = float(r @ r)
    log(f"Initial residual: {math.sqrt(rs_old):.4e}")

    for i in range(min(max_iter, n * 10)):
        Ap     = A_r @ p
        denom  = float(p @ Ap)
        if abs(denom) < 1e-14: break
        alpha  = rs_old / denom
        x      = x + alpha * p
        r      = r - alpha * Ap
        rs_new = float(r @ r)
        res    = math.sqrt(rs_new)

        if (i + 1) % 20 == 0 or res < tol:
            log(f"  iter {i+1}  residual={res:.4e}")
        if res < tol:
            log(f"  Converged at iteration {i+1}")
            break

        p      = r + (rs_new / (rs_old + 1e-14)) * p
        rs_old = rs_new

    x_c      = np.linalg.solve(A_r, b_r)
    residual = float(np.linalg.norm(A_r @ x - b_r))
    rel_err  = float(np.linalg.norm(x - x_c) / (np.linalg.norm(x_c) + 1e-14))
    log(f"Final residual={residual:.4e}  rel_err={rel_err:.4e}")
    return QLSResult("QuantumCG", A, b, x, x_c, residual, rel_err, steps)


# ═══════════════════════════════════════════════════════════════
#  Randomised Kaczmarz
# ═══════════════════════════════════════════════════════════════

def randomised_kaczmarz(A, b, iterations=500, verbose=True):
    """
    Randomised Kaczmarz row-action solver.
    Converges for any consistent Hermitian system.
    """
    steps = []
    def log(msg):
        steps.append(msg)
        if verbose: print("  " + msg)

    A, b = _validate(A, b)
    A_r  = A.real.astype(float)
    b_r  = b.real.astype(float)
    n    = A_r.shape[0]
    log(f"Kaczmarz: {n}x{n} system, {iterations} iterations")

    row_norms_sq = np.sum(A_r ** 2, axis=1)
    probs        = row_norms_sq / row_norms_sq.sum()
    x            = np.zeros(n)
    rng          = np.random.default_rng(0)

    for it in range(iterations):
        i     = int(rng.choice(n, p=probs))
        ai    = A_r[i]
        denom = row_norms_sq[i]
        if denom < 1e-14: continue
        x     = x + ((b_r[i] - ai @ x) / denom) * ai

        if (it + 1) % 100 == 0:
            res = float(np.linalg.norm(A_r @ x - b_r))
            log(f"  iter {it+1}  residual={res:.4e}")

    x_c      = np.linalg.solve(A_r, b_r)
    residual = float(np.linalg.norm(A_r @ x - b_r))
    rel_err  = float(np.linalg.norm(x - x_c) / (np.linalg.norm(x_c) + 1e-14))
    log(f"Final residual={residual:.4e}  rel_err={rel_err:.4e}")
    return QLSResult("Kaczmarz", A, b, x, x_c, residual, rel_err, steps)


# ═══════════════════════════════════════════════════════════════
#  Quantum-Preconditioned GMRES (FIXED)
# ═══════════════════════════════════════════════════════════════

def quantum_preconditioned_gmres(A, b, max_iter=50, tol=1e-10, verbose=True):
    """
    Quantum-Preconditioned GMRES.
    Uses diagonal preconditioner P = diag(1/sqrt(|A_ii|)),
    then solves the preconditioned system with full GMRES.
    Works for any invertible Hermitian A.
    """
    steps = []
    def log(msg):
        steps.append(msg)
        if verbose: print("  " + msg)

    A, b = _validate(A, b)
    A_r  = A.real.astype(float)
    b_r  = b.real.astype(float)
    n    = A_r.shape[0]
    log(f"Q-GMRES: {n}x{n} system  max_iter={max_iter}")

    # Diagonal preconditioner
    diag       = np.abs(np.diag(A_r))
    diag       = np.where(diag > 1e-10, diag, 1.0)
    P_inv      = np.diag(1.0 / np.sqrt(diag))
    log(f"Preconditioner: diag({np.round(np.diag(P_inv), 4).tolist()})")

    # Preconditioned system: (P A P) y = P b,  x = P y
    PAP = P_inv @ A_r @ P_inv
    Pb  = P_inv @ b_r

    # ── Full GMRES on PAP y = Pb ─────────────────────────────
    m   = min(max_iter, n)
    y   = np.zeros(n)
    r0  = Pb - PAP @ y
    b0  = float(np.linalg.norm(r0))

    if b0 < tol:
        x = P_inv @ y
        x_c      = np.linalg.solve(A_r, b_r)
        residual = float(np.linalg.norm(A_r @ x - b_r))
        rel_err  = float(np.linalg.norm(x - x_c) / (np.linalg.norm(x_c) + 1e-14))
        log("Converged immediately")
        return QLSResult("Q-GMRES", A, b, x, x_c, residual, rel_err, steps)

    # Arnoldi basis Q (n x m+1), upper Hessenberg H (m+1 x m)
    Q  = np.zeros((n, m + 1))
    H  = np.zeros((m + 1, m))
    Q[:, 0] = r0 / b0

    # Givens rotation accumulators
    cs = np.zeros(m)
    sn = np.zeros(m)
    g  = np.zeros(m + 1)
    g[0] = b0

    solved = False
    j_end  = m

    for j in range(m):
        # Arnoldi step
        w = PAP @ Q[:, j]
        for i in range(j + 1):
            H[i, j] = float(Q[:, i] @ w)
            w        = w - H[i, j] * Q[:, i]
        H[j + 1, j] = float(np.linalg.norm(w))

        if H[j + 1, j] > 1e-12:
            Q[:, j + 1] = w / H[j + 1, j]

        # Apply previous Givens rotations to new column
        for i in range(j):
            tmp         =  cs[i] * H[i, j] + sn[i] * H[i + 1, j]
            H[i + 1, j] = -sn[i] * H[i, j] + cs[i] * H[i + 1, j]
            H[i, j]     = tmp

        # Compute new Givens rotation
        denom = math.sqrt(H[j, j] ** 2 + H[j + 1, j] ** 2)
        if denom < 1e-14:
            cs[j], sn[j] = 1.0, 0.0
        else:
            cs[j] = H[j, j]     / denom
            sn[j] = H[j + 1, j] / denom

        # Apply to H and g
        H[j, j]     =  cs[j] * H[j, j] + sn[j] * H[j + 1, j]
        H[j + 1, j] = 0.0
        g[j + 1]    = -sn[j] * g[j]
        g[j]        =  cs[j] * g[j]

        res_j = abs(g[j + 1])
        if (j + 1) % 10 == 0:
            log(f"  iter {j+1}  residual={res_j:.4e}")

        if res_j < tol:
            log(f"  Converged at iteration {j+1}")
            j_end  = j + 1
            solved = True
            break

    if not solved:
        j_end = m

    # Back-substitution on upper triangular H[:j_end, :j_end]
    Hj = H[:j_end, :j_end]
    gj = g[:j_end]
    try:
        y_sol = np.linalg.solve(Hj, gj)
    except np.linalg.LinAlgError:
        y_sol, _, _, _ = np.linalg.lstsq(Hj, gj, rcond=None)

    y = Q[:, :j_end] @ y_sol

    # Recover x = P y  (undo left preconditioning)
    x        = P_inv @ y
    x_c      = np.linalg.solve(A_r, b_r)
    residual = float(np.linalg.norm(A_r @ x - b_r))
    rel_err  = float(np.linalg.norm(x - x_c) / (np.linalg.norm(x_c) + 1e-14))
    log(f"Final residual={residual:.4e}  rel_err={rel_err:.4e}")
    return QLSResult("Q-GMRES", A, b, x, x_c, residual, rel_err, steps)
