QVLASOV

"""
QVLASOV lane -- classical reference for Miyamoto et al. (PRResearch 6, 013200, 2024), Sec. IV.

Linearized Vlasov equation for a light species (neutrino) in an external CDM force, in 1D
phase space (x, v). Central-difference discretization -> df/dt = A f with A REAL ANTISYMMETRIC
(paper Eq. 16, 23). Since iA is Hermitian this is the Schroedinger-form the quantum algorithm
simulates; here we integrate it CLASSICALLY (exp(A T) f0) as the ground-truth the quantum
Trotter circuits are benchmarked against.

Toy parameters (paper Sec. IV): L=2, V=1, n_gr=64, F_CDM(x)=A sin(Kx) with A=-1, K=pi,
Maxwell IC f0(x,u) = exp(-u^2/2 sigma^2)/sqrt(2 pi sigma^2), sigma=0.1, evolve to T=0.1, 0.2.
Expected (Figs. 3-4): the CDM force (single mode K) drives ONLY the k=K density mode.
"""
import numpy as np
import scipy.sparse as sp
from scipy.sparse.linalg import expm_multiply


def build_operator(n=64, L=2.0, V=1.0, A=-1.0, K=np.pi):
    """Return (A_matrix, x, v). f indexed as flat[i*n + j], i=x, j=v.
    Streaming -v d/dx (periodic in x); force -F(x) d/dv (Dirichlet f=0 outside v-grid)."""
    dx = L / n
    x = np.arange(n) * dx
    dv = 2.0 * V / (n + 1)
    v = -V + (np.arange(n) + 1) * dv          # Dirichlet grid, excludes +/-V (Eq. 11)
    F = A * np.sin(K * x)                      # CDM force (Eq. 68)

    rows, cols, vals = [], [], []

    def add(r, c, val):
        rows.append(r); cols.append(c); vals.append(val)

    for i in range(n):
        ip, im = (i + 1) % n, (i - 1) % n      # periodic in x
        for j in range(n):
            r = i * n + j
            # streaming: -v_j (f[i+1,j] - f[i-1,j]) / (2 dx)
            add(r, ip * n + j, -v[j] / (2 * dx))
            add(r, im * n + j, +v[j] / (2 * dx))
            # force: -F_i (f[i,j+1] - f[i,j-1]) / (2 dv), Dirichlet
            if j + 1 < n:
                add(r, i * n + (j + 1), -F[i] / (2 * dv))
            if j - 1 >= 0:
                add(r, i * n + (j - 1), +F[i] / (2 * dv))

    Amat = sp.csr_matrix((vals, (rows, cols)), shape=(n * n, n * n))
    return Amat, x, v


def maxwell_ic(n, v, sigma=0.1):
    """f0 constant in x, Maxwell in v (Eq. 70). Returns flat vector (n*n,)."""
    fv = np.exp(-v ** 2 / (2 * sigma ** 2)) / np.sqrt(2 * np.pi * sigma ** 2)
    f0 = np.tile(fv, (n, 1))                   # (n_x, n_v), constant in x
    return f0.reshape(-1)


def density(f_flat, n, dv):
    """rho(x) = sum_v f dv  ->  delta(x) = (rho - mean)/mean."""
    f = f_flat.reshape(n, n)
    rho = f.sum(axis=1) * dv
    return rho, (rho - rho.mean()) / rho.mean()


def evolve(Amat, f0, T):
    """f(T) = exp(A T) f0  (A antisymmetric -> orthogonal -> norm preserving)."""
    return expm_multiply(Amat * T, f0)


def run():
    n, L, V, K = 64, 2.0, 1.0, np.pi
    Amat, x, v = build_operator(n=n, L=L, V=V, A=-1.0, K=K)
    dv = 2.0 * V / (n + 1)
    f0 = maxwell_ic(n, v, sigma=0.1)

    print(f"QVLASOV classical: n_gr={n} (dim {n*n}), L={L} V={V} K=pi, A antisymmetric="
          f"{np.allclose((Amat + Amat.T).toarray(), 0, atol=1e-12)}")
    results = {}
    for T in (0.0, 0.1, 0.2):
        fT = evolve(Amat, f0, T) if T > 0 else f0.copy()
        rho, delta = density(fT, n, dv)
        norm = np.linalg.norm(fT) / np.linalg.norm(f0)
        dk = np.abs(np.fft.rfft(delta))
        kmode = np.argmax(dk[1:]) + 1  # dominant nonzero mode
        results[T] = dict(delta=delta, dk=dk, norm=norm, kmode=int(kmode))
        print(f"T={T:.1f}: |f|/|f0|={norm:.6f} (orthogonal->1)  dominant density mode k_index={kmode} "
              f"(K=pi -> index {int(round(K*L/(2*np.pi)))})  max|delta|={np.abs(delta).max():.3e}")
    np.savez("data/vlasov_classical.npz",
             **{f"delta_T{int(T*10)}": results[T]["delta"] for T in results},
             **{f"dk_T{int(T*10)}": results[T]["dk"] for T in results}, x=x)
    # check: at T>0 the dominant mode is k-index 1 (= K=pi over L=2 -> one wavelength)
    kexpect = int(round(K * L / (2 * np.pi)))  # = 1
    ok = all(results[T]["kmode"] == kexpect for T in (0.1, 0.2)) and \
         all(abs(results[T]["norm"] - 1) < 1e-6 for T in results)
    print("QVLASOV classical:", "PASS" if ok else "CHECK",
          f"(single mode k={kexpect} driven, norm conserved)")
    return results, x


if __name__ == "__main__":
    run()

"""
QVLASOV quantum stage 1 -- Trotterized Hamiltonian simulation of the linearized Vlasov
operator, emulated EXACTLY on a statevector (numpy). This is the algorithm the paper runs
on an FTQC; here we verify at TOY scale that:
  (a) H = iA is Hermitian (A antisymmetric) -> exp(-iHt) is a genuine quantum evolution;
  (b) Trotter(exp(-iH_x t/n) exp(-iH_v t/n))^n -> exact exp(-iHt) as n grows (2nd-order);
  (c) the encoded state reproduces the CLASSICAL neutrino density mode (Fig. 4).

n_gr in {4, 8}: state dimension n_gr^2 = 16 (4 qubits) or 64 (6 qubits) -- QPU-scale.
"""
import numpy as np
from scipy.linalg import expm


def build_parts(n, L=2.0, V=1.0, A=-1.0, K=np.pi):
    """Dense A_stream (-v d/dx) and A_force (-F d/dv); flat index i*n+j (i=x, j=v)."""
    dx, dv = L / n, 2.0 * V / (n + 1)
    x = np.arange(n) * dx
    v = -V + (np.arange(n) + 1) * dv
    F = A * np.sin(K * x)
    As = np.zeros((n * n, n * n)); Af = np.zeros((n * n, n * n))
    for i in range(n):
        ip, im = (i + 1) % n, (i - 1) % n
        for j in range(n):
            r = i * n + j
            As[r, ip * n + j] += -v[j] / (2 * dx)
            As[r, im * n + j] += +v[j] / (2 * dx)
            if j + 1 < n:
                Af[r, i * n + (j + 1)] += -F[i] / (2 * dv)
            if j - 1 >= 0:
                Af[r, i * n + (j - 1)] += +F[i] / (2 * dv)
    return As, Af, x, v, dv


def maxwell(n, v, sigma=0.1):
    fv = np.exp(-v ** 2 / (2 * sigma ** 2)) / np.sqrt(2 * np.pi * sigma ** 2)
    return np.tile(fv, (n, 1)).reshape(-1)


def density_mode(state, n, dv, kidx=1):
    rho = state.reshape(n, n).real.sum(axis=1) * dv
    delta = (rho - rho.mean()) / rho.mean()
    return np.abs(np.fft.rfft(delta))[kidx]


def main():
    for n in (4, 8):
        As, Af, x, v, dv = build_parts(n)
        A = As + Af
        H = 1j * A
        herm = np.allclose(H, H.conj().T, atol=1e-12)
        f0 = maxwell(n, v); f0 = f0 / np.linalg.norm(f0)
        T = 0.2
        exact = expm(-1j * H * T) @ f0      # = expm(A T) f0
        # (b) Trotter convergence
        print(f"--- n_gr={n} (dim {n*n} = {int(np.log2(n*n))} qubits) | H Hermitian={herm}")
        Hx, Hv = 1j * As, 1j * Af
        for nt in (1, 2, 4, 8, 16, 32):
            Ux = expm(-1j * Hx * (T / nt))
            Uv = expm(-1j * Hv * (T / nt))
            state = f0.copy()
            for _ in range(nt):
                state = Uv @ (Ux @ state)
            err = np.linalg.norm(state - exact)
            print(f"   Trotter steps={nt:2d}: |psi_trot - psi_exact| = {err:.2e}")
        # (c) density mode matches classical
        km_exact = density_mode(exact, n, dv)
        km_f0 = density_mode(f0, n, dv)
        print(f"   density k=1 mode: t=0 -> {km_f0:.2e},  T={T} -> {km_exact:.2e} "
              f"(enhanced, cf. classical Fig. 4)")


if __name__ == "__main__":
    main()

"""
QVLASOV quantum stage 2 -- gate-level circuit for real hardware (IonQ via Braket).

We Pauli-decompose H_x=iA_stream and H_v=iA_force (n_gr=4 -> 4 qubits), build a Trotter
circuit of Pauli-exponential gadgets (exp(-i c_k P_k tau) = basis-change + CNOT-ladder +
Rz), and VALIDATE it (numpy) against the exact evolution before any hardware.
This native-gate circuit is what Braket transpiles for IonQ.
"""
import numpy as np
from itertools import product
from trotter_emulation import build_parts, maxwell

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


def kron_string(s):
    M = np.array([[1]], dtype=complex)
    for ch in s:
        M = np.kron(M, PAULI[ch])
    return M


def pauli_decompose(H, nq, tol=1e-9):
    """H (2^nq x 2^nq Hermitian) -> {pauli_string: real_coeff}."""
    d = 2 ** nq
    terms = {}
    for combo in product("IXYZ", repeat=nq):
        s = "".join(combo)
        c = np.trace(kron_string(s).conj().T @ H) / d
        if abs(c) > tol:
            terms[s] = float(np.real(c))
    return terms


def exp_pauli_matrix(s, theta, nq):
    """exp(-i theta P) as a full matrix (for validation)."""
    P = kron_string(s)
    return np.cos(theta) * np.eye(2 ** nq) - 1j * np.sin(theta) * P


def trotter_unitary(terms_x, terms_v, T, n_trotter, nq):
    """(prod exp(-i c P tau) over v-terms, then x-terms)^n_trotter, tau=T/n_trotter."""
    tau = T / n_trotter
    U = np.eye(2 ** nq, dtype=complex)
    for _ in range(n_trotter):
        for s, c in terms_x.items():
            U = exp_pauli_matrix(s, c * tau, nq) @ U
        for s, c in terms_v.items():
            U = exp_pauli_matrix(s, c * tau, nq) @ U
    return U


def validate(n_gr=4):
    from scipy.linalg import expm
    nq = int(np.log2(n_gr * n_gr))
    As, Af, x, v, dv = build_parts(n_gr)
    Hx, Hv = 1j * As, 1j * Af
    H = Hx + Hv
    tx = pauli_decompose(Hx, nq)
    tv = pauli_decompose(Hv, nq)
    print(f"n_gr={n_gr} -> {nq} qubits | Pauli terms: H_x={len(tx)}, H_v={len(tv)}")
    # (a) decomposition exact
    Hx_rec = sum(c * kron_string(s) for s, c in tx.items())
    print(f"  decomposition error |H_x - sum c P| = {np.abs(Hx - Hx_rec).max():.2e}")
    # (b) Trotter circuit unitary -> exact evolution
    T = 0.2
    exact = expm(-1j * H * T)
    for nt in (1, 2, 4, 8):
        U = trotter_unitary(tx, tv, T, nt, nq)
        err = np.abs(U - exact).max()
        print(f"  Trotter steps={nt}: |U_circuit - U_exact|_max = {err:.2e}")
    # (c) act on Maxwell IC, check density mode grows
    f0 = maxwell(n_gr, v); f0 = f0 / np.linalg.norm(f0)
    psi = trotter_unitary(tx, tv, T, 8, nq) @ f0
    rho = psi.reshape(n_gr, n_gr).real.sum(axis=1) * dv
    delta = (rho - rho.mean()) / rho.mean()
    print(f"  density k=1 mode after circuit: {np.abs(np.fft.rfft(delta))[1]:.3e} (grows from 0)")
    return tx, tv


if __name__ == "__main__":
    validate(4)

"""
QVLASOV -- Braket gate circuit (n_gr=4, 4 qubits) for IonQ.
Circuit = state-prep(Maxwell f0, a product state: uniform in x, Gaussian in v)
          + Trotter( exp(-i H_x tau) exp(-i H_v tau) ) via native Pauli-exponential gadgets.
Validated on the LOCAL simulator against the classical evolution BEFORE any cloud/QPU cost.
"""
import numpy as np
from braket.circuits import Circuit
from braket.devices import LocalSimulator
from trotter_emulation import build_parts, maxwell
from quantum_circuit import pauli_decompose, trotter_unitary

N_GR = 4
NQ = 4  # qubits 0,1 = x (MSB), qubits 2,3 = v (LSB)


def gaussian_prep_angles(n_gr=4, sigma=0.1, V=1.0):
    """4-amplitude Gaussian (v-grid) -> angles for a 2-qubit real state prep."""
    dv = 2 * V / (n_gr + 1)
    v = -V + (np.arange(n_gr) + 1) * dv
    a = np.exp(-v ** 2 / (2 * sigma ** 2)); a = a / np.linalg.norm(a)  # a[j], j=2*q2+q3
    r0 = np.hypot(a[0], a[1]); r1 = np.hypot(a[2], a[3])
    beta = 2 * np.arctan2(r1, r0)
    g0 = 2 * np.arctan2(a[1], a[0])
    g1 = 2 * np.arctan2(a[3], a[2])
    return beta, g0, g1, a


def add_state_prep(circ):
    """|0000> -> f0 = (uniform on x-qubits 0,1) (x) (Gaussian on v-qubits 2,3)."""
    circ.h(0); circ.h(1)                       # uniform in x
    beta, g0, g1, _ = gaussian_prep_angles()
    circ.ry(2, beta)                            # split between v-halves
    # uniformly-controlled Ry on qubit 3 controlled by qubit 2
    circ.ry(3, (g0 + g1) / 2)
    circ.cnot(2, 3)
    circ.ry(3, (g0 - g1) / 2)
    circ.cnot(2, 3)
    return circ


def add_pauli_exp(circ, s, theta):
    """exp(-i theta P): basis-change + CNOT ladder + Rz(2 theta) + uncompute."""
    active = [q for q, ch in enumerate(s) if ch != "I"]
    if not active or abs(theta) < 1e-12:
        return
    for q in active:            # basis change: map Pauli -> Z
        if s[q] == "X":
            circ.h(q)
        elif s[q] == "Y":
            circ.rx(q, np.pi / 2)
    for a, b in zip(active[:-1], active[1:]):
        circ.cnot(a, b)
    circ.rz(active[-1], 2 * theta)
    for a, b in zip(active[:-1][::-1], active[1:][::-1]):
        circ.cnot(a, b)
    for q in active:            # uncompute basis change
        if s[q] == "X":
            circ.h(q)
        elif s[q] == "Y":
            circ.rx(q, -np.pi / 2)


def build_circuit(T=0.2, n_trotter=4, prep=True):
    As, Af, x, v, dv = build_parts(N_GR)
    tx = pauli_decompose(1j * As, NQ)
    tv = pauli_decompose(1j * Af, NQ)
    tau = T / n_trotter
    circ = Circuit()
    if prep:
        add_state_prep(circ)
    for _ in range(n_trotter):
        for s, c in tx.items():
            add_pauli_exp(circ, s, c * tau)
        for s, c in tv.items():
            add_pauli_exp(circ, s, c * tau)
    return circ, (As, Af, x, v, dv)


def density_mode_from_state(state, n_gr, dv):
    rho = np.abs(state.reshape(n_gr, n_gr)).real
    # state amplitudes are real up to global phase; use real part with sign from phase
    amp = state.reshape(n_gr, n_gr)
    rho = amp.real.sum(axis=1) * dv
    delta = (rho - rho.mean()) / rho.mean()
    return np.abs(np.fft.rfft(delta))[1], delta


def validate():
    from scipy.linalg import expm
    T = 0.2
    As, Af, x, v, dv = build_parts(N_GR)
    f0 = maxwell(N_GR, v); f0 = f0 / np.linalg.norm(f0)
    exact = expm((As + Af) * T) @ f0
    exact_n = exact / np.linalg.norm(exact)
    dev = LocalSimulator()
    print(f"{'n_trot':>6} {'depth':>6} {'gates':>6} {'2q-gates':>8} {'fidelity':>10}")
    for nt in (1, 2, 3, 4):
        circ, _ = build_circuit(T=T, n_trotter=nt, prep=True)
        n2q = sum(1 for ins in circ.instructions if ins.operator.qubit_count == 2)
        sv = np.asarray(dev.run(circ.state_vector(), shots=0).result().values[0])
        fid = abs(np.vdot(sv, exact_n)) ** 2
        print(f"{nt:6d} {circ.depth:6d} {len(circ.instructions):6d} {n2q:8d} {fid:10.5f}")
    return None


if __name__ == "__main__":
    validate()

"""
Run the QVLASOV circuit on Braket devices and compare the measured phase-space distribution
to the classical ground truth. SV1 (noiseless cloud) validates the pipeline; IonQ Forte-1
(real QPU) is the hardware demo -- benchmarked classical <-> noiseless <-> real.

The hardware observable is the computational-basis probability P(bitstring) = |amplitude|^2
over the 16 phase-space grid points; classical target = |exp(A T) f0|^2. Distribution
fidelity F = (sum sqrt(p_i q_i))^2 quantifies how well the device reproduces the physics.
"""
import sys
import numpy as np
from scipy.linalg import expm
from braket.aws import AwsDevice
from braket.devices import LocalSimulator
from trotter_emulation import build_parts, maxwell
from braket_run import build_circuit, N_GR

S3 = ("amazon-braket-us-east-1-297904677595", "qvlasov")
SV1 = "arn:aws:braket:::device/quantum-simulator/amazon/sv1"
IONQ = "arn:aws:braket:us-east-1::device/qpu/ionq/Forte-1"


def classical_distribution(T=0.2):
    As, Af, x, v, dv = build_parts(N_GR)
    f0 = maxwell(N_GR, v); f0 = f0 / np.linalg.norm(f0)
    psi = expm((As + Af) * T) @ f0
    return np.abs(psi) ** 2 / np.sum(np.abs(psi) ** 2)


def counts_to_dist(counts, nq=4):
    d = np.zeros(2 ** nq)
    tot = sum(counts.values())
    for bstr, c in counts.items():
        d[int(bstr, 2)] = c / tot
    return d


def fidelity(p, q):
    return float(np.sum(np.sqrt(np.clip(p, 0, None) * np.clip(q, 0, None))) ** 2)


def run(device_name, n_trotter=1, shots=1000, T=0.2):
    circ, _ = build_circuit(T=T, n_trotter=n_trotter, prep=True)
    cls = classical_distribution(T)
    # noiseless local reference
    loc = LocalSimulator().run(circ, shots=shots).result().measurement_counts
    loc_d = counts_to_dist(loc)
    print(f"[local noiseless] n_trotter={n_trotter} shots={shots} "
          f"fidelity vs classical = {fidelity(loc_d, cls):.4f}", flush=True)
    if device_name == "local":
        return
    arn = SV1 if device_name == "sv1" else IONQ
    dev = AwsDevice(arn)
    print(f"submitting to {device_name} ({arn.split('/')[-1]}) ...", flush=True)
    task = dev.run(circ, shots=shots, s3_destination_folder=S3)
    print(f"  task id: {task.id}", flush=True)
    if device_name == "sv1":
        res = task.result()  # SV1 returns quickly
        dev_d = counts_to_dist(res.measurement_counts)
        print(f"[SV1] fidelity vs classical = {fidelity(dev_d, cls):.4f}", flush=True)
    else:
        print(f"  IonQ task QUEUED -- poll with: {task.id}", flush=True)
        print(f"  status: {task.state()}", flush=True)
    return task


if __name__ == "__main__":
    dev = sys.argv[1] if len(sys.argv) > 1 else "local"
    nt = int(sys.argv[2]) if len(sys.argv) > 2 else 1
    run(dev, n_trotter=nt)

"""
Submit the QVLASOV Trotter circuits (n_trotter=1,2,3) to IonQ Forte-1 for the
classical <-> noiseless <-> real-QPU breaking-point curve. Saves task ARNs for polling.
"""
import json
import numpy as np
from braket.aws import AwsDevice
from braket_run import build_circuit
from run_cloud import S3, IONQ, classical_distribution, counts_to_dist, fidelity
from braket.devices import LocalSimulator

SHOTS = 300
NTS = [1, 2, 3]

dev = AwsDevice(IONQ)
cls = classical_distribution()
tasks = {}
print(f"IonQ Forte-1 -- submitting n_trotter={NTS} at {SHOTS} shots "
      f"(est ${len(NTS)*(SHOTS*0.08+0.30):.2f})", flush=True)
for nt in NTS:
    circ, _ = build_circuit(T=0.2, n_trotter=nt, prep=True)
    n2q = sum(1 for ins in circ.instructions if ins.operator.qubit_count == 2)
    loc = counts_to_dist(LocalSimulator().run(circ, shots=2000).result().measurement_counts)
    task = dev.run(circ, shots=SHOTS, s3_destination_folder=S3)
    tasks[nt] = {"arn": task.id, "depth": circ.depth, "n2q": n2q,
                 "noiseless_fidelity": fidelity(loc, cls)}
    print(f"  nt={nt}: depth={circ.depth} 2q={n2q} noiseless_fid={fidelity(loc, cls):.4f} "
          f"-> {task.id}  [{task.state()}]", flush=True)

with open("data/ionq_tasks.json", "w") as f:
    json.dump({"shots": SHOTS, "cls_dist": cls.tolist(), "tasks": tasks}, f, indent=2)
print("saved data/ionq_tasks.json -- poll with poll_ionq.py", flush=True)