#!/usr/bin/env python
#
# file: results/hw_kernel_job.py
#
# revision history:
#  20260810 (am): rebuilt to project coding standards
#  20260805 (am): initial version
#
# Measures a quantum kernel on real IBM hardware and caches it, so that
# Segment 4 of the notebook never has to wait on a queue.
#
# THIS SPENDS REAL QPU TIME. The Open plan allowance is about ten
# minutes per month, so the script prints an estimate and asks for
# confirmation before it submits anything.
#
#   python hw_kernel_job.py            # estimate, then confirm
#   python hw_kernel_job.py --yes      # skip the confirmation
#
# Before the first run, paste your api key into IBM_API_KEY below. Leave
# CRN as None and it is looked up for you -- see backends.discover_crn
# and notebook segment 4.5 for why the crn matters.
#------------------------------------------------------------------------------

# import system modules
#
import json
import logging
import sys
import time
import warnings
from pathlib import Path

# quiet the noisy account resolver before qiskit is imported
#
warnings.filterwarnings("ignore")
logging.getLogger("qiskit_ibm_runtime.accounts.account").setLevel(
    logging.CRITICAL)

# import third-party modules
#
import numpy as np
from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
from qiskit_ibm_runtime import SamplerV2

# make the lab package importable when run from this directory
#
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

# import qmllab modules
#
from qmllab import backends, data, kernels, models          # noqa: E402

#------------------------------------------------------------------------------
#
# global variables are listed here
#
#------------------------------------------------------------------------------

# paste your IBM Quantum api key between the quotes. get one free at
# https://quantum.cloud.ibm.com. leave CRN as None and it is found for
# you. if you would rather keep the key in a file, set TOKEN_PATH to it
# instead and leave IBM_API_KEY empty.
#
IBM_API_KEY = ""
TOKEN_PATH = None
CRN = None

# the experiment. a symmetric n x n kernel costs n(n-1)/2 circuits, so
# these numbers drive the whole bill -- raise them carefully.
#
DATASET = "two_moons"
SCALER = "minmax"
N_TRAIN = 20
N_TEST = 12
SHOTS = 1024
REPS = 2

# where the cached result is written
#
OUT = Path(__file__).resolve().parent / "hw_kernel_20x20.json"

#------------------------------------------------------------------------------
#
# functions are listed here
#
#------------------------------------------------------------------------------

def main(assume_yes=False):
    """
    function: main

    arguments:
     assume_yes: skip the confirmation prompt

    return:
     the path of the cached result

    description:
     Builds the compute-uncompute circuits for a train and a test
     kernel, submits them as a single job, then caches the measured
     matrices beside the exact ones for comparison.
    """

    # stop early if no key was supplied, before anything is built
    #
    if not (IBM_API_KEY.strip() or TOKEN_PATH):
        print("no api key set. open this file and paste yours into "
              "IBM_API_KEY near the top.")
        print("get one free at https://quantum.cloud.ibm.com")
        return None

    # load and split exactly as the notebook does, so the cached result
    # lines up with what segment 4 expects
    #
    X_train, X_test, y_train, y_test, _ = data.prepare(
        DATASET, scaler=SCALER, n_samples=1000, test_size=0.2,
        random_state=42)
    A, ya = X_train[:N_TRAIN], y_train[:N_TRAIN]
    B, yb = X_test[:N_TEST], y_test[:N_TEST]

    # compute the exact kernels first; they are free and give us a
    # reference to score the hardware against
    #
    feature_map = kernels.make_feature_map("zz", 2, reps=REPS)
    K_train_exact = kernels.exact_kernel(A, feature_map=feature_map)
    K_test_exact = kernels.exact_kernel(B, A, feature_map=feature_map)

    # build one circuit per pair we actually need
    #
    circuits, tags = [], []
    for i in range(N_TRAIN):
        for j in range(i + 1, N_TRAIN):
            circuits.append(kernels.fidelity_circuit(A[i], A[j], feature_map))
            tags.append(("train", i, j))
    for i in range(N_TEST):
        for j in range(N_TRAIN):
            circuits.append(kernels.fidelity_circuit(B[i], A[j], feature_map))
            tags.append(("test", i, j))

    # price the job and let the user back out
    #
    cost = backends.qpu_cost(N_TRAIN, N_TEST)
    print(f"{len(circuits)} circuits, {SHOTS} shots each")
    print(f"estimated QPU: {cost['qpu_seconds']:.0f}s "
          f"({cost['percent_of_monthly']:.1f}% of a monthly Open plan)")
    if not assume_yes:
        if input("submit this job? [y/N] ").strip().lower() != "y":
            print("cancelled -- no QPU time spent")
            return None

    # connect. the crn is looked up automatically unless one was set.
    #
    service = backends.get_service(api_key=IBM_API_KEY or None,
                                   token_path=TOKEN_PATH, crn=CRN)

    # pick the processor with the shortest queue
    #
    backend = service.least_busy(operational=True, simulator=False)

    # transpile to the device's native gates
    #
    pm = generate_preset_pass_manager(backend=backend, optimization_level=3)
    isa = [pm.run(c) for c in circuits]
    print(f"backend={backend.name}, transpiled")

    # submit everything as a single job
    #
    start = time.time()
    job = SamplerV2(mode=backend).run(isa, shots=SHOTS)
    print("job id:", job.job_id(), flush=True)
    result = job.result()
    wall = time.time() - start

    # the kernel entry is the probability of the all-zeros outcome
    #
    K_train = np.eye(N_TRAIN)
    K_test = np.zeros((N_TEST, N_TRAIN))
    for (kind, i, j), pub in zip(tags, result):
        value = pub.data.meas.get_counts().get("00", 0) / SHOTS
        if kind == "train":
            K_train[i, j] = K_train[j, i] = value
        else:
            K_test[i, j] = value

    # score the hardware against the exact reference
    #
    try:
        usage = job.metrics().get("usage", {})
    except Exception:
        usage = {}
    off = ~np.eye(N_TRAIN, dtype=bool)
    corr = np.corrcoef(K_train[off], K_train_exact[off])[0, 1]
    acc_hw = models.qsvm_from_kernels(K_train, ya, K_test, yb,
                                      C=10)["test_acc"]
    acc_exact = models.qsvm_from_kernels(K_train_exact, ya, K_test_exact, yb,
                                         C=10)["test_acc"]

    # report
    #
    print(f"\nwall {wall:.0f}s   usage {usage}")
    print(f"corr(hw, exact) train kernel : {corr:.4f}")
    print(f"mean |hw - exact|            : "
          f"{np.abs(K_train - K_train_exact)[off].mean():.4f}")
    print(f"QSVM acc  hardware kernel    : {acc_hw:.4f}")
    print(f"QSVM acc  exact kernel       : {acc_exact:.4f}")

    # cache everything segment 4 needs to replay this run
    #
    record = {"backend": backend.name, "job_id": job.job_id(),
              "n_train": N_TRAIN, "n_test": N_TEST, "shots": SHOTS,
              "wall_s": wall, "usage": usage, "acc_hw": acc_hw,
              "acc_exact": acc_exact, "corr": corr,
              "X_train": A.tolist(), "y_train": ya.tolist(),
              "X_test": B.tolist(), "y_test": yb.tolist(),
              "K_train_hw": K_train.tolist(),
              "K_test_hw": K_test.tolist(),
              "K_train_exact": K_train_exact.tolist(),
              "K_test_exact": K_test_exact.tolist()}
    with open(OUT, "w") as fp:
        json.dump(record, fp, indent=1)
    print("cached ->", OUT)

    # exit gracefully
    #
    return OUT
#
# end of function

#------------------------------------------------------------------------------
#
# the main program starts here
#
#------------------------------------------------------------------------------

if __name__ == "__main__":
    main(assume_yes="--yes" in sys.argv)

#
# end of file