"""
Opening the door on quantum ML -- best-shot methods + a quantum-native control.

The earlier sweep used off-the-shelf fidelity kernels. This gives quantum ML its
BEST honest shot on the detection task, and adds a control that proves the test can
SEE a quantum advantage when one really exists:

  Methods (each vs a fairly-tuned classical baseline):
    - fidelity quantum kernel (IQP / Havlicek embedding)
    - PROJECTED quantum kernel (Huang et al. 2021 -- often stronger)
    - TRAINABLE quantum kernel (kernel-target alignment: fit the embedding scale to data)
    - data re-uploading variational classifier (Perez-Salinas 2020, a universal QML model)

  Tasks:
    - FIRE: the SEVIRI-calibrated fire/non-fire task (real-shaped)
    - QUANTUM-NATIVE control: labels generated FROM the quantum embedding, so a quantum
      kernel matching that structure SHOULD beat classical RBF. If quantum wins here but
      not on fire, the conclusion is precise: quantum ML helps iff the data has
      quantum-native structure -- and fire-detection data does not.

CPU/simulator, small (nq=4, n<=100). No QPU, no DGX.
"""
from __future__ import annotations
import os, time, json
import numpy as np
import pennylane as qml
from sklearn.svm import SVC
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.decomposition import PCA
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import roc_auc_score
from qsweep_core import make_seviri_dataset

NQ = 4


def _fire_features(n, seed):
    X, y = make_seviri_dataset(n, seed=seed)
    Xk = MinMaxScaler((0, np.pi)).fit_transform(PCA(NQ, random_state=0).fit_transform(
        StandardScaler().fit_transform(X)))
    return Xk, y


def _quantum_native(n, seed):
    """Labels generated from the SAME quantum embedding the kernel uses -> quantum-friendly."""
    rng = np.random.default_rng(seed)
    X = rng.uniform(0, np.pi, (n, NQ))
    dev = qml.device("default.qubit", wires=NQ)

    @qml.qnode(dev)
    def obs(x):
        qml.IQPEmbedding(x, wires=range(NQ), n_repeats=2)
        return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1) @ qml.PauliZ(2))

    vals = np.array([obs(x) for x in X])
    y = (vals > np.median(vals)).astype(int)
    return X, y


# ---- kernels -------------------------------------------------------------
def _dev():
    return qml.device("default.qubit", wires=NQ)


def fidelity_kernel(A, B, scale=1.0):
    dev = _dev()

    @qml.qnode(dev)
    def ov(x1, x2):
        qml.IQPEmbedding(scale * x1, wires=range(NQ), n_repeats=2)
        qml.adjoint(qml.IQPEmbedding)(scale * x2, wires=range(NQ), n_repeats=2)
        return qml.probs(wires=range(NQ))

    K = np.zeros((len(A), len(B)))
    for i, a in enumerate(A):
        for j, b in enumerate(B):
            K[i, j] = ov(a, b)[0]
    return K


def projected_features(X, scale=1.0):
    """Huang projected kernel features: per-qubit <X>,<Y>,<Z> after the embedding."""
    dev = _dev()

    @qml.qnode(dev)
    def feats(x):
        qml.IQPEmbedding(scale * x, wires=range(NQ), n_repeats=2)
        return ([qml.expval(qml.PauliX(i)) for i in range(NQ)] +
                [qml.expval(qml.PauliY(i)) for i in range(NQ)] +
                [qml.expval(qml.PauliZ(i)) for i in range(NQ)])

    return np.array([feats(x) for x in X])


def _svc_auc_precomputed(Kfn, X, y, seed, folds=4):
    skf = StratifiedKFold(folds, shuffle=True, random_state=seed)
    a = []
    for tr, te in skf.split(X, y):
        Ktr = Kfn(X[tr], X[tr]); Kte = Kfn(X[te], X[tr])
        m = SVC(kernel="precomputed", C=2.0).fit(Ktr, y[tr])
        a.append(roc_auc_score(y[te], m.decision_function(Kte)))
    return float(np.mean(a))


def _svc_auc_features(feat_fn, X, y, seed, folds=4):
    skf = StratifiedKFold(folds, shuffle=True, random_state=seed)
    a = []
    F = feat_fn(X)
    for tr, te in skf.split(X, y):
        m = SVC(kernel="rbf", C=2.0, gamma="scale").fit(F[tr], y[tr])
        a.append(roc_auc_score(y[te], m.decision_function(F[te])))
    return float(np.mean(a))


def classical_auc(X, y, seed, folds=4):
    skf = StratifiedKFold(folds, shuffle=True, random_state=seed)
    best = 0.0
    for g in ["scale", 0.05, 0.1, 0.5, 1, 2, 5]:
        a = []
        for tr, te in skf.split(X, y):
            m = SVC(kernel="rbf", C=2.0, gamma=g).fit(X[tr], y[tr])
            a.append(roc_auc_score(y[te], m.decision_function(X[te])))
        best = max(best, float(np.mean(a)))
    # plus GBM
    ag = []
    for tr, te in skf.split(X, y):
        m = GradientBoostingClassifier(random_state=seed).fit(X[tr], y[tr])
        ag.append(roc_auc_score(y[te], m.predict_proba(X[te])[:, 1]))
    return max(best, float(np.mean(ag)))


def trainable_kernel_auc(X, y, seed):
    """Kernel-target alignment: pick embedding scale maximizing alignment on a subset."""
    rng = np.random.default_rng(seed)
    idx = rng.choice(len(y), min(40, len(y)), replace=False)
    Xs, ys = X[idx], y[idx]
    yt = np.where(ys == 1, 1.0, -1.0); T = np.outer(yt, yt)
    best_scale, best_align = 1.0, -1
    for scale in [0.5, 1.0, 1.5, 2.0, 3.0]:
        K = projected_features(Xs, scale)
        G = K @ K.T
        align = (G * T).sum() / (np.linalg.norm(G) * np.linalg.norm(T) + 1e-9)
        if align > best_align:
            best_align, best_scale = align, scale
    return _svc_auc_features(lambda Z: projected_features(Z, best_scale), X, y, seed), best_scale


def reuploading_auc(X, y, seed, layers=2, steps=12):
    from pennylane import numpy as pnp
    dev = _dev()
    rng = np.random.default_rng(seed)

    @qml.qnode(dev, interface="autograd")
    def circ(x, w):
        for l in range(layers):
            for i in range(NQ):
                qml.RY(x[i % X.shape[1]] * w[l, i, 0] + w[l, i, 1], wires=i)
            for i in range(NQ - 1):
                qml.CNOT(wires=[i, i + 1])
        return qml.expval(qml.PauliZ(0))

    skf = StratifiedKFold(3, shuffle=True, random_state=seed)
    aucs = []
    for tr, te in skf.split(X, y):
        w = pnp.array(0.1 * rng.standard_normal((layers, NQ, 2)), requires_grad=True)
        yt = np.where(y[tr] == 1, 1.0, -1.0)
        opt = qml.AdamOptimizer(0.15)

        def cost(w):
            p = pnp.stack([circ(X[i], w) for i in tr])
            return pnp.mean((p - yt) ** 2)

        for _ in range(steps):
            w = opt.step(cost, w)
        scores = np.array([circ(X[i], w) for i in te])
        aucs.append(roc_auc_score(y[te], scores))
    return float(np.mean(aucs))


def run_task(name, Xy_fn, n=100, seed=0):
    X, y = Xy_fn(n, seed)
    t = time.perf_counter()
    res = {
        "classical_best": round(classical_auc(X, y, seed), 3),
        "fidelity_kernel": round(_svc_auc_precomputed(fidelity_kernel, X, y, seed), 3),
        "projected_kernel": round(_svc_auc_features(projected_features, X, y, seed), 3),
    }
    tk, sc = trainable_kernel_auc(X, y, seed)
    res["trainable_kernel"] = round(tk, 3)
    if os.environ.get("QML_REUPLOAD") == "1":
        res["reuploading_clf"] = round(reuploading_auc(X, y, seed), 3)
    qbest = max(res["fidelity_kernel"], res["projected_kernel"], res["trainable_kernel"],
                res.get("reuploading_clf", 0.0))
    res["best_quantum"] = round(qbest, 3)
    res["quantum_beats_classical"] = bool(qbest > res["classical_best"] + 0.02)
    res["latency_s"] = round(time.perf_counter() - t, 1)
    print(f"[{name}] {json.dumps(res)}")
    return res


if __name__ == "__main__":
    out = {
        "fire_task": run_task("FIRE", _fire_features, n=80, seed=0),
        "quantum_native_control": run_task("QUANTUM-NATIVE", _quantum_native, n=80, seed=0),
    }
    out["verdict"] = (
        "Door test: quantum ML given its best shot (projected + trainable kernels, "
        "data-reuploading). If quantum wins the quantum-native control but NOT the fire "
        "task, the conclusion is precise -- quantum ML helps only when the data has "
        "quantum-native structure, which fire-detection data does not.")
    json.dump(out, open("results/qml_door.json", "w"), indent=2)
    print("\n" + out["verdict"])
    print("record -> results/qml_door.json")