Ejecuta cargas de trabajo cuánticas con QRMI

python3 -m venv ~/.venvs/qrmi-ibm
source ~/.venvs/qrmi-ibm/bin/activate
python -m pip install "qrmi[ibm]" python-dotenv pyscf ffsim qiskit-addon-sqd matplotlib ipykernel
python -m ipykernel install --user --name qrmi-ibm --display-name "QRMI IBM"

import os
import time
import json
import numpy as np
from dotenv import load_dotenv

from qrmi import (
    QuantumResource,
    ResourceType,
    Payload,
    TaskStatus,
    get_job_qpu_resources_and_types,
)
from qrmi.primitives import QRMIService
from qrmi.primitives.ibm import SamplerV2, get_target

from qiskit import QuantumCircuit, qasm3
from qiskit.circuit.library import efficient_su2
from qiskit.primitives.containers.sampler_pub import SamplerPub
from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager

# Load credentials from .env without overriding already-set scheduler environment variables
load_dotenv(override=False)

# Preserve resources if injected by Slurm SPANK plugin; fallback to default for interactive run
BACKEND_NAME = os.environ.get("QRMI_JOB_QPU_RESOURCES", "ibm_kingston")
os.environ.setdefault("QRMI_JOB_QPU_RESOURCES", BACKEND_NAME)
os.environ.setdefault("QRMI_JOB_QPU_TYPES", "ibm-quantum-compute-service")

print(f"Backend: {BACKEND_NAME}")
print("Environment ready.")

# ── IBM Quantum ───────────────────────────────────────────────────────
qrmi = QuantumResource(BACKEND_NAME, ResourceType.IBMQuantumComputeService)
# ResourceType.IBMQuantumSystem is the alternative for directly provisioned systems

print(f"Resource id:   {qrmi.resource_id()}")
print(f"Resource type: {qrmi.resource_type()}")
print(f"Accessible:    {qrmi.is_accessible()}")

# Acquire exclusive access — open try/finally immediately so every
# subsequent failure (target retrieval, transpilation, submission) is covered.
# Release is skipped when running under Slurm: the SPANK plugin owns the
# session lifecycle and will release it when the job finishes.
lock = qrmi.acquire()
print(f"Lock token:    {lock}")
try:
    # Retrieve backend capabilities
    transpiler_target = get_target(
        qrmi
    )  # calls qrmi.target() and parses the JSON
    target_json = json.loads(qrmi.target().value)
    config = target_json.get("configuration", {})
    print(f"\nBackend: {config.get('backend_name', 'unknown')}")
    print(f"Qubits:  {config.get('n_qubits', 'unknown')}")
    print(f"Gates:   {config.get('basis_gates', [])}")

    # ── IBM Quantum ───────────────────────────────────────────────────

    # Build a Bell state circuit
    qc = QuantumCircuit(2)
    qc.h(0)
    qc.cx(0, 1)
    qc.measure_all()
    print(qc.draw("text"))

    # Transpile to ISA using the target retrieved in Step 1
    pm = generate_preset_pass_manager(
        optimization_level=1, target=transpiler_target
    )
    isa_circuit = pm.run(qc)
    print(f"\nTranspiled gate counts: {isa_circuit.count_ops()}")

    # Build the QRMI payload
    # Payload.QiskitPrimitive wraps the IBM SamplerV2 input schema:
    #   pubs: list of [qasm3_string, parameter_values]  (shots goes at top level)
    #   program_id: "sampler" or "estimator"
    shots = 1024
    pub = SamplerPub.coerce((isa_circuit,), shots)
    qasm3_str = qasm3.dumps(
        pub.circuit,
        disable_constants=True,
        allow_aliasing=True,
        experimental=qasm3.ExperimentalFeatures.SWITCH_CASE_V1,
    )
    # Parameter values as a flat list (empty for non-parametric circuits)
    param_array = pub.parameter_values.as_array(
        pub.circuit.parameters
    ).tolist()

    input_json = {
        "pubs": [
            [qasm3_str, param_array]
        ],  # list-of-lists; shots at top level
        "version": 2,
        "support_qiskit": False,  # True returns binary-encoded Qiskit result
        "shots": shots,
    }
    payload = Payload.QiskitPrimitive(
        input=json.dumps(input_json), program_id="sampler"
    )
    print("Payload ready")

    # ── IBM Quantum ───────────────────────────────────────────────────

    # Submit the job
    job_id = qrmi.task_start(payload)
    print(f"Job submitted: {job_id}")

    # Poll until complete
    while True:
        status = qrmi.task_status(job_id)
        print(f"  Status: {status}")
        if status not in [TaskStatus.Running, TaskStatus.Queued]:
            break
        time.sleep(5)

    print(f"\nFinal status: {status}")

    # Retrieve results
    # support_qiskit=False → plain JSON; parse directly without ResultDecoder
    if status == TaskStatus.Completed:
        raw = qrmi.task_result(job_id).value
        result = json.loads(raw)
        # IBM QCS plain-JSON result shape: {"results": [{"data": {"meas": {"samples": [...]}}}]}
        # samples is a list of hex-encoded integers; decode to zero-padded bitstrings
        samples = result["results"][0]["data"]["meas"]["samples"]
        num_bits = sum(reg.size for reg in isa_circuit.cregs)
        from collections import Counter

        counts = Counter(format(int(s, 16), f"0{num_bits}b") for s in samples)
        print(f"\nMeasurement counts: {dict(counts.most_common(8))}")
        qrmi.task_stop(job_id)
    else:
        print(f"Job did not complete. Logs:\n{qrmi.task_logs(job_id)}")

finally:
    # Release only in interactive sessions; under Slurm the SPANK plugin
    # manages the session lifecycle and calling release() here would
    # prematurely close a session it does not own.
    if not os.environ.get("SLURM_JOB_ID"):
        qrmi.release(lock)
        print("\nSession released.")

# QRMIService reads QRMI_JOB_QPU_RESOURCES / QRMI_JOB_QPU_TYPES set in Setup or Slurm
service = QRMIService()
qrmi_svc = service.resources()[0]
print(f"Using: {qrmi_svc.resource_id()} ({qrmi_svc.resource_type()})")

# Build an EfficientSU2 circuit
circuit = efficient_su2(5, entanglement="linear")
circuit.measure_all()
param_values = np.random.rand(circuit.num_parameters)

pm = generate_preset_pass_manager(
    optimization_level=1, target=get_target(qrmi_svc)
)
isa_circuit = pm.run(circuit)

# SamplerV2 executes jobs against the QRMI resource and decodes results into primitive containers
sampler = SamplerV2(qrmi_svc, options={"default_shots": 1024})
job = sampler.run([(isa_circuit, param_values)])
print(f"Job ID: {job.job_id()} | Status: {job.status()}")

# Poll with retry — re-raise immediately on permanent failures;
# only retry on transient network/timeout errors (connection resets, 503s).
_TRANSIENT = (
    "503",
    "Service Unavailable",
    "ConnectionError",
    "TimeoutError",
    "timed out",
    "Connection reset",
)
result = None
for attempt in range(60):
    try:
        result = job.result()  # blocks until complete
        break
    except Exception as e:
        if not any(tok in str(e) for tok in _TRANSIENT):
            raise
        print(f"  Transient error on attempt {attempt + 1}: {e}")
        time.sleep(10)

if result is not None:
    counts = result[0].data.meas.get_counts()
    print(f"Counts (first 5): {dict(list(counts.items())[:5])}")
else:
    print("Job did not complete after retries.")

if job.errored():
    print(f"Logs:\n{job.logs()}")

#SBATCH --gres=qpu:1
#SBATCH --qpu=ibm_kingston
python my_workflow.py   # QRMI_JOB_QPU_RESOURCES and QRMI_JOB_QPU_TYPES are already set

# get_job_qpu_resources_and_types() reads QRMI_JOB_QPU_RESOURCES / QRMI_JOB_QPU_TYPES
# set by the Slurm SPANK plugin (or manually above in Setup)
qpus, qpu_types = get_job_qpu_resources_and_types()
print("Resources allocated by scheduler:")
for qpu, qpu_type in zip(qpus, qpu_types):
    print(f"  {qpu}  ({qpu_type})")

# QRMIService wraps this into a list of ready QuantumResource objects
for r in QRMIService().resources():
    print(
        f"\nQRMIService found: {r.resource_id()}  accessible={r.is_accessible()}"
    )

from qrmi.primitives.ibm import get_backend
import math
import os
import time
from functools import partial
from dotenv import load_dotenv
import numpy as np
import matplotlib.pyplot as plt

import pyscf
import pyscf.gto
import pyscf.scf
import pyscf.cc
import pyscf.mcscf
import pyscf.ao2mo

import ffsim
import ffsim.qiskit
from qiskit import QuantumCircuit, QuantumRegister
from qiskit_addon_sqd.fermion import (
    SCIResult,
    diagonalize_fermionic_hamiltonian,
    solve_sci_batch,
)
from qrmi.primitives import QRMIService
from qrmi.primitives.ibm import SamplerV2, get_target

load_dotenv(override=False)
os.environ.setdefault("QRMI_JOB_QPU_RESOURCES", "ibm_kingston")
os.environ.setdefault("QRMI_JOB_QPU_TYPES", "ibm-quantum-compute-service")

# ── Step 1: Map classical inputs to a quantum problem ─────────────────

# Build N2 molecule at 1.0 Å bond distance
mol = pyscf.gto.Mole()
mol.build(
    atom=[["N", (0, 0, 0)], ["N", (1.0, 0, 0)]],
    basis="cc-pvdz",
    symmetry="Dooh",
)

# Define active space: freeze 2 core orbitals
n_frozen = 2
active_space = range(n_frozen, mol.nao_nr())

# Get molecular integrals
scf = pyscf.scf.RHF(mol).run()
norb = len(active_space)
n_electrons = int(sum(scf.mo_occ[active_space]))
n_alpha = (n_electrons + mol.spin) // 2
n_beta = (n_electrons - mol.spin) // 2
nelec = (n_alpha, n_beta)

cas = pyscf.mcscf.CASCI(scf, norb, nelec)
mo = cas.sort_mo(active_space, base=0)
hcore, nuclear_repulsion_energy = cas.get_h1cas(mo)
eri = pyscf.ao2mo.restore(1, cas.get_h2cas(mo), norb)

# Reference energy from external SCI calculation
reference_energy = -109.22802921665716

print(
    f"N₂/cc-pVDZ active space: {norb} orbitals ({2 * norb} qubits), {nelec} electrons"
)
print(f"SCF energy:       {scf.e_tot:.8f} Ha")
print(f"Reference energy: {reference_energy:.8f} Ha")

# Get CCSD amplitudes for initializing the LUCJ ansatz
ccsd = pyscf.cc.CCSD(
    scf, frozen=[i for i in range(mol.nao_nr()) if i not in active_space]
).run()
t1 = ccsd.t1
t2 = ccsd.t2
print(f"CCSD energy:      {ccsd.e_tot:.8f} Ha")

# Discover backend via QRMIService (QRMI_JOB_QPU_RESOURCES set in Setup)
service = QRMIService()
qrmi_sqd = service.resources()[0]
print(f"Using QRMI resource: {qrmi_sqd.resource_id()}")

# get_backend() wraps the QRMI resource as a Qiskit backend for layout synthesis

backend = get_backend(qrmi_sqd)

# Set ansatz properties
n_reps = 1
pairs_aa = [(p, p + 1) for p in range(norb - 1)]
pairs_ab = None

# Create pass manager adapted to hardware heavy-hex topology
pass_manager, pairs_ab = ffsim.qiskit.generate_lucj_pass_manager(
    backend=backend,
    norb=norb,
    connectivity="heavy-hex",
    interaction_pairs=(pairs_aa, pairs_ab),
    optimization_level=3,
)

# Create the compressed LUCJ ansatz operator
ucj_op = ffsim.UCJOpSpinBalanced.from_t_amplitudes(
    t2=t2,
    t1=t1,
    n_reps=n_reps,
    interaction_pairs=(pairs_aa, pairs_ab),
    optimize=True,
    options=dict(maxiter=1000),
)

# Assemble the circuit
qubits = QuantumRegister(2 * norb, name="q")
circuit = QuantumCircuit(qubits)
circuit.append(ffsim.qiskit.PrepareHartreeFockJW(norb, nelec), qubits)
circuit.append(ffsim.qiskit.UCJOpSpinBalancedJW(ucj_op), qubits)
circuit.measure_all()
print(f"LUCJ circuit: {circuit.num_qubits} qubits, depth {circuit.depth()}")

# ── Step 2: Optimize for quantum hardware execution ───────────────────

isa_circuit = pass_manager.run(circuit)
print(f"Transpiled gate counts: {isa_circuit.count_ops()}")

# ── Step 3: Execute using Qiskit primitives (QRMI SamplerV2) ─────────

sampler = SamplerV2(qrmi_sqd, options={"default_shots": 100_000})
# sampler.options.environment.job_tags = ["TUT_SQD"]
job = sampler.run([(isa_circuit,)])
print(f"Job submitted via QRMI: {job.job_id()} | Status: {job.status()}")
print("Waiting for results from hardware...")

_TRANSIENT = (
    "503",
    "Service Unavailable",
    "ConnectionError",
    "TimeoutError",
    "timed out",
    "Connection reset",
)
primitive_result = None
for attempt in range(120):
    try:
        primitive_result = job.result()
        break
    except Exception as e:
        if not any(tok in str(e) for tok in _TRANSIENT):
            raise
        print(f"  Transient error on attempt {attempt + 1}: {e}")
        time.sleep(10)

if primitive_result is None:
    raise RuntimeError("Job did not complete after retries")

pub_result = primitive_result[0]
bit_array = pub_result.data.meas
print(f"Total shots collected: {bit_array.num_shots}")

# ── Step 4: Post-process and return result in classical format ────────


def is_valid_bitstring(
    bitstring: str, norb: int, nelec: tuple[int, int]
) -> bool:
    n_a, n_b = nelec
    return (
        len(bitstring) == 2 * norb
        and bitstring[norb:].count("1") == n_a
        and bitstring[:norb].count("1") == n_b
    )


num_valid = sum(
    is_valid_bitstring(b, norb, nelec) for b in bit_array.get_bitstrings()
)
valid_fraction = num_valid / bit_array.num_shots
expected_random = (
    math.comb(norb, n_alpha) * math.comb(norb, n_beta) / (2 ** (2 * norb))
)
print(f"Fraction of valid configurations sampled: {valid_fraction:.5f}")
print(f"Expected fraction from uniform random:     {expected_random:.4e}")

# Configure SQD eigensolver
energy_tol = 1e-3
occupancies_tol = 1e-3
max_iterations = 5
num_batches = 3
samples_per_batch = 300
symmetrize_spin = True
carryover_threshold = 1e-4
max_cycle = 200

# Hartree-Fock initial occupancy guess
initial_occupancies = (
    np.array([1] * n_alpha + [0] * (norb - n_alpha)),
    np.array([1] * n_beta + [0] * (norb - n_beta)),
)

sci_solver = partial(solve_sci_batch, spin_sq=0.0, max_cycle=max_cycle)
result_history = []


def callback(results: list[SCIResult]):
    result_history.append(results)
    iteration = len(result_history)
    print(f"Iteration {iteration}")
    for i, res in enumerate(results):
        subspace_dim = np.prod(res.sci_state.amplitudes.shape)
        print(
            f"  Subsample {i}: Energy = {res.energy + nuclear_repulsion_energy:.8f} Ha | Subspace dim = {subspace_dim}"
        )


print("\nRunning SQD post-processing...")
rng = np.random.default_rng(42)
sqd_result = diagonalize_fermionic_hamiltonian(
    hcore,
    eri,
    bit_array,
    samples_per_batch=samples_per_batch,
    norb=norb,
    nelec=nelec,
    num_batches=num_batches,
    energy_tol=energy_tol,
    occupancies_tol=occupancies_tol,
    max_iterations=max_iterations,
    sci_solver=sci_solver,
    symmetrize_spin=symmetrize_spin,
    initial_occupancies=initial_occupancies,
    carryover_threshold=carryover_threshold,
    callback=callback,
    seed=rng,
)

final_energy = sqd_result.energy + nuclear_repulsion_energy
energy_error = final_energy - reference_energy

print("\n=== Energy Summary (N₂/cc-pVDZ active space) ===")
print(f"SCF energy:       {scf.e_tot:.8f} Ha")
print(f"Reference energy: {reference_energy:.8f} Ha")
print(f"Final SQD energy: {final_energy:.8f} Ha")
print(
    f"Energy error:     {energy_error:.8f} Ha ({abs(energy_error) * 627.5:.4f} kcal/mol)"
)

# ── Visualization ─────────────────────────────────────────────────────

x1 = range(len(result_history))
min_e = [
    min(res, key=lambda r: r.energy).energy + nuclear_repulsion_energy
    for res in result_history
]
e_diff = [abs(e - reference_energy) for e in min_e]
chem_accuracy = 0.001  # ~1 mHa / ~0.6 kcal/mol

y2 = np.sum(sqd_result.orbital_occupancies, axis=0)
x2 = range(len(y2))

fig, axs = plt.subplots(1, 2, figsize=(12, 5))

# Energies convergence plot
axs[0].plot(x1, e_diff, label="Energy error", marker="o")
axs[0].set_xticks(list(x1))
axs[0].set_xticklabels(list(x1))
axs[0].set_yscale("log")
axs[0].axhline(
    y=chem_accuracy,
    color="#BF5700",
    linestyle="--",
    label="Chemical accuracy (1 mHa)",
)
axs[0].set_title("SQD Energy Error vs Iteration")
axs[0].set_xlabel("Iteration")
axs[0].set_ylabel("Energy Error (Ha)")
axs[0].legend()

# Spatial orbital occupancy plot
axs[1].bar(x2, y2, width=0.8)
axs[1].set_xticks(list(x2)[::2])
axs[1].set_xticklabels(list(x2)[::2])
axs[1].set_title("Avg Occupancy per Spatial Orbital")
axs[1].set_xlabel("Spatial Orbital Index")
axs[1].set_ylabel("Avg Occupancy")

plt.tight_layout()
plt.show()