#!/usr/bin/env python
#
# file: qmllab/experiments.py
#
# revision history:
#  20260810 (am): initial version for workshop lab 02
#
# One entry point for every experiment in the lab.
#
# Everything the notebook shows can be reproduced from a config dict,
# which is what makes run_experiment.py possible. Participants can keep
# exploring long after the ninety minutes are up.
#
# Three routines do most of the work:
#
#   run()        one method on one dataset
#   compare()    classical against quantum on the same split, which is
#                the question the whole lab is built around
#   sweep()      one parameter varied over a list of values
#
# Every routine accepts an output directory. When one is given the run
# writes a json record, a decision boundary, and a readable summary, so
# a result can be reviewed after the fact instead of scrolling back.
#------------------------------------------------------------------------------

# import system modules
#
import json
import time
from pathlib import Path

# import third-party modules
#
import numpy as np

# import qmllab modules
#
from . import data as _data
from . import kernels as _kernels
from . import models as _models

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

# where results and figures land when no output directory is given
#
LAB_DIR = Path(__file__).resolve().parent.parent
RESULTS_DIR = LAB_DIR / "results"
FIGURES_DIR = LAB_DIR / "figures"

# every knob an experiment understands, with its default
#
DEFAULTS = {
    "dataset": "two_moons",
    "dataset_path": None,
    "eval_path": None,
    "eval_n_samples": None,
    "scaler": "minmax",
    "n_samples": 300,
    "test_size": 0.2,
    "random_state": 42,
    "method": "qsvm",
    "feature_map": "zz",
    "reps": 2,
    "entanglement": "full",
    "kernel_backend": "exact",
    "shots": 1024,
    "layers": 2,
    "maxiter": 300,
    "optimizer": "COBYLA",
}

# the methods run() knows how to build
#
METHODS = ["classical", "qsvm", "vqc", "vqc_naive"]

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

def run(config=None, verbose=True, output_dir=None, plot=True, **overrides):
    """
    function: run

    arguments:
     config:     a dict of settings; missing keys fall back to DEFAULTS
     verbose:    print a one-line summary
     output_dir: where to write json, figures and a summary
     plot:       save a decision boundary when output_dir is given
     overrides:  individual settings, which win over config

    return:
     a dict of results

    description:
     Runs one method on one dataset, end to end. When output_dir is
     given the run is also written to disk so it can be reviewed later.
    """

    # merge the defaults, the config and any direct overrides
    #
    cfg = {**DEFAULTS, **(config or {}), **overrides}

    # n_samples defaults to a lab-sized 300 so the notebook stays quick.
    # that is the wrong default for someone who handed us a file: it
    # would silently throw away most of their data. so when a path was
    # supplied and n_samples was not asked for explicitly, use all of it.
    #
    supplied = set((config or {}).keys()) | set(overrides.keys())
    if cfg["dataset_path"] is not None and "n_samples" not in supplied:
        cfg["n_samples"] = None

    # load the data. when eval_path is set there is no split: the whole
    # dataset trains and the evaluation set comes from that file.
    #
    X_train, X_test, y_train, y_test, _ = _data.prepare(
        name=cfg["dataset"], scaler=cfg["scaler"], path=cfg["dataset_path"],
        n_samples=cfg["n_samples"], test_size=cfg["test_size"],
        random_state=cfg["random_state"], eval_path=cfg["eval_path"],
        eval_n_samples=cfg["eval_n_samples"])

    # record what we are about to do, including where the scores came
    # from -- a held-out split and a supplied file are not the same claim
    #
    out = {"config": cfg,
           "n_train": len(X_train),
           "n_test": len(X_test),
           "eval_source": ("split" if cfg["eval_path"] is None
                           else str(cfg["eval_path"])),
           "range_width": _data.range_width(X_train)}
    start = time.time()

    # a plain classical svm
    #
    if cfg["method"] == "classical":
        res = _models.classical_svm(X_train, y_train, X_test, y_test)

    # path a: a quantum kernel feeding a classical svm
    #
    elif cfg["method"] == "qsvm":
        feature_map = _kernels.make_feature_map(
            cfg["feature_map"], X_train.shape[1], cfg["reps"],
            cfg["entanglement"])
        kernel_fn, kernel_kwargs = _resolve_kernel(cfg)
        res = _models.qsvm(X_train, y_train, X_test, y_test,
                           feature_map=feature_map, kernel_fn=kernel_fn,
                           **kernel_kwargs)
        out["kernel_stats"] = _kernels.kernel_stats(res["K_train"])

    # path b: a data re-uploading circuit
    #
    elif cfg["method"] == "vqc":
        res = _models.vqc_reuploading(
            X_train, y_train, X_test, y_test, layers=cfg["layers"],
            maxiter=cfg["maxiter"], optimizer=cfg["optimizer"],
            seed=cfg["random_state"])

    # path b done the obvious way, kept as a control
    #
    elif cfg["method"] == "vqc_naive":
        res = _models.vqc_naive(X_train, y_train, X_test, y_test,
                                reps=cfg["reps"], maxiter=cfg["maxiter"],
                                optimizer=cfg["optimizer"])

    # exit ungracefully -- unknown method
    #
    else:
        raise ValueError(
            f"unknown method {cfg['method']!r}; choose from {METHODS}")

    # keep the fitted model aside; it is not json-serializable
    #
    model = res.pop("model", None)
    res.pop("circuit", None)
    out.update(res)
    out["total_time"] = time.time() - start

    # write the run to disk when the caller asked for it
    #
    if output_dir is not None:
        out["output_dir"] = str(_write_run(out, model, cfg, X_train, y_train,
                                           output_dir, plot))

    # print a one-line summary
    #
    if verbose:
        _print_row(out)

    # exit gracefully
    #
    return out
#
# end of function


def compare(config=None, output_dir=None, verbose=True, **overrides):
    """
    function: compare

    arguments:
     config:     a dict of settings; missing keys fall back to DEFAULTS
     output_dir: where to write json, figures and a summary
     verbose:    print the comparison table
     overrides:  individual settings, which win over config

    return:
     a dict holding each method's results and a verdict

    description:
     Runs a classical svm and a quantum kernel svm on exactly the same
     split, then reports which one won and by how much. This is the
     question the whole lab is built around, so it gets a routine of its
     own rather than being assembled by hand each time.

     Adding "vqc" to the methods list also trains a variational circuit,
     which costs considerably more time.
    """

    # keep only what the caller actually supplied, so that run() can
    # still tell a deliberate setting from an untouched default
    #
    user = {**(config or {}), **overrides}
    methods = user.pop("methods", ["classical", "qsvm"])

    # run each method on the same data and seed, giving each its own
    # subdirectory so their records do not overwrite one another
    #
    rows = {}
    for method in methods:
        sub = None if output_dir is None else Path(output_dir) / method
        rows[method] = run({**user, "method": method}, verbose=False,
                           output_dir=sub, plot=sub is not None)

    # read the fully-resolved settings back off the first run
    #
    cfg = rows[methods[0]]["config"]

    # decide who won on the test set
    #
    best = max(rows, key=lambda m: rows[m]["test_acc"])
    classical = rows.get("classical", {}).get("test_acc")
    quantum = rows.get("qsvm", {}).get("test_acc")
    summary = {"dataset": cfg["dataset"],
               "scaler": cfg["scaler"],
               "n_train": rows[methods[0]]["n_train"],
               "n_test": rows[methods[0]]["n_test"],
               "eval_source": rows[methods[0]]["eval_source"],
               "best": best,
               "accuracies": {m: rows[m]["test_acc"] for m in methods}}

    # describe the outcome in words, since a number alone is not a result
    #
    if classical is not None and quantum is not None:
        gap = quantum - classical
        summary["gap"] = gap
        if abs(gap) < 0.02:
            summary["verdict"] = "tie"
        elif gap > 0:
            summary["verdict"] = "quantum wins"
        else:
            summary["verdict"] = "classical wins"

    # print the comparison
    #
    if verbose:
        _print_comparison(summary, rows)

    # write the comparison alongside the individual runs
    #
    if output_dir is not None:
        path = Path(output_dir) / "comparison.json"
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as fp:
            json.dump({"summary": summary,
                       "runs": {m: _strip(rows[m]) for m in rows}},
                      fp, indent=1, default=float)
        summary["path"] = str(path)

    # exit gracefully
    #
    return {"summary": summary, "runs": rows}
#
# end of function


def sweep(param, values, config=None, output_dir=None, verbose=True,
          **overrides):
    """
    function: sweep

    arguments:
     param:      the setting to vary
     values:     the values it should take
     config:     a dict of settings
     output_dir: where to write each run
     verbose:    print one line per run
     overrides:  individual settings

    return:
     a list of result dicts, one per value

    description:
     Varies a single parameter and collects the results. This is the
     lab's main tool for asking whether something actually matters.
    """

    # run once per value, keeping everything else fixed
    #
    rows = []
    for value in values:
        sub = None
        if output_dir is not None:
            sub = Path(output_dir) / f"{param}_{value}"
        rows.append(run({**(config or {}), **overrides, param: value},
                        verbose=verbose, output_dir=sub))

    # exit gracefully
    #
    return rows
#
# end of function


def scaler_study(config=None, scalers=("minmax", "maxabs", "l1", "l2",
                                       "robust", "standard"),
                 output_dir=None, verbose=False, **overrides):
    """
    function: scaler_study

    arguments:
     config:     a dict of settings
     scalers:    the scalers to compare
     output_dir: where to write each run
     verbose:    print one line per run
     overrides:  individual settings

    return:
     a dict of scaler name -> width, accuracy and kernel concentration

    description:
     Reproduces the segment 2 headline: quantum kernel accuracy tracks
     the width of the encoded range, because the ZZ feature map's angle
     is a product of features.
    """

    # run the same experiment once per scaler
    #
    out = {}
    for scaler in scalers:
        sub = None
        if output_dir is not None:
            sub = Path(output_dir) / f"scaler_{scaler}"
        res = run({**(config or {}), **overrides, "scaler": scaler},
                  verbose=verbose, output_dir=sub)
        out[scaler] = {
            "width": res["range_width"],
            "test_acc": res["test_acc"],
            "n_support": res.get("n_support"),
            "mean_offdiag": res.get("kernel_stats", {}).get("mean_offdiag")}

    # exit gracefully
    #
    return out
#
# end of function


def load_hardware_results(name="hw_kernel_20x20.json"):
    """
    function: load_hardware_results

    arguments:
     name: the cached result file to read

    return:
     a dict with the arrays restored to numpy

    description:
     Loads a cached hardware run so that segment 4 never waits on a
     queue. Real qpu time is scarce, so these results are measured once
     and replayed thereafter.
    """

    # reject a missing cache with a message that says what to do
    #
    path = RESULTS_DIR / name
    if not path.exists():
        raise FileNotFoundError(
            f"{path} not found -- run results/hw_kernel_job.py first "
            "(costs real QPU time)")

    # read the record and restore the arrays
    #
    with open(path) as fp:
        record = json.load(fp)
    for key in ("K_train_hw", "K_test_hw", "K_train_exact", "K_test_exact",
                "X_train", "X_test", "y_train", "y_test"):
        if key in record:
            record[key] = np.array(record[key])

    # exit gracefully
    #
    return record
#
# end of function


def save(results, filename):
    """
    function: save

    arguments:
     results:  anything json-serializable
     filename: the name to write under the results directory

    return:
     the path that was written

    description:
     Writes a result record into the lab's results directory.
    """

    # make sure the directory exists, then write the record
    #
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    path = RESULTS_DIR / filename
    with open(path, "w") as fp:
        json.dump(results, fp, indent=1, default=float)

    # exit gracefully
    #
    return path
#
# end of function


def _resolve_kernel(cfg):
    """
    function: _resolve_kernel

    arguments:
     cfg: the merged configuration

    return:
     a tuple (kernel function, keyword arguments)

    description:
     Maps a kernel_backend string onto the routine that implements it.
    """

    # exact linear algebra, no sampling at all
    #
    kind = cfg["kernel_backend"]
    if kind == "exact":
        return _kernels.exact_kernel, {}

    # sampled from a finite number of shots
    #
    if kind == "sampled":
        return _kernels.fidelity_kernel, {"shots": cfg["shots"]}

    # run on a simulator, optionally carrying a device noise model
    #
    if kind in ("aer", "noisy"):
        from .backends import get_backend
        backend = get_backend("ideal" if kind == "aer" else "noisy")
        return _kernels.fidelity_kernel, {"shots": cfg["shots"],
                                          "backend": backend}

    # refuse to spend qpu time from inside a sweep
    #
    if kind == "hardware":
        raise ValueError(
            "run hardware kernels via results/hw_kernel_job.py -- they cost "
            "real QPU time and should be cached, not called from a sweep")

    # exit ungracefully -- unknown backend
    #
    raise ValueError(f"unknown kernel_backend {kind!r}")
#
# end of function


def _write_run(out, model, cfg, X_train, y_train, output_dir, plot):
    """
    function: _write_run

    arguments:
     out:        the result dict
     model:      the fitted model, or None
     cfg:        the merged configuration
     X_train:    the training features
     y_train:    the training labels
     output_dir: where to write
     plot:       whether to save a decision boundary

    return:
     the output directory

    description:
     Writes one run to disk: a json record, a readable summary and, when
     asked, a decision boundary rendered on the simulator.
    """

    # make the output directory
    #
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # write the machine-readable record
    #
    with open(output_dir / "results.json", "w") as fp:
        json.dump(_strip(out), fp, indent=1, default=float)

    # write a short human-readable summary
    #
    with open(output_dir / "summary.txt", "w") as fp:
        fp.write(_format_summary(out))

    # render a decision boundary when there is a model to draw
    #
    if plot and model is not None:
        _save_boundary(model, cfg, X_train, y_train,
                       output_dir / "decision_boundary.png")

    # exit gracefully
    #
    return output_dir
#
# end of function


def _save_boundary(model, cfg, X_train, y_train, path):
    """
    function: _save_boundary

    arguments:
     model:   the fitted model
     cfg:     the merged configuration
     X_train: the training features
     y_train: the training labels
     path:    the png to write

    return:
     the path that was written, or None when there is nothing to draw

    description:
     Renders a decision boundary for whatever model run() just fitted.
     Always drawn on the simulator: on hardware this single picture
     would cost hours of qpu time.
    """

    # nothing to draw without a model
    #
    if model is None:
        return None

    # use a non-interactive backend so this works over ssh
    #
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    # import here to avoid a circular import at module load
    #
    from . import plotting as _plotting

    # describe the panel
    #
    spec = {"model": model, "X": X_train, "y": y_train, "X_train": X_train,
            "title": (f"{cfg['method']} · {cfg['dataset']} · "
                      f"scaler={cfg['scaler']}")}

    # a kernel model needs its kernel to evaluate the grid
    #
    if cfg["method"] == "qsvm":
        spec["kernel_fn"] = _kernels.exact_kernel
        spec["feature_map"] = _kernels.make_feature_map(
            cfg["feature_map"], X_train.shape[1], cfg["reps"],
            cfg["entanglement"])

    # ring the support vectors when the model has them
    #
    if hasattr(model, "support_"):
        spec["support_idx"] = model.support_

    # draw and save
    #
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    _plotting.plot_decision_boundary(**spec)
    plt.savefig(path, dpi=120, bbox_inches="tight")
    plt.close()

    # exit gracefully
    #
    return path
#
# end of function


def _strip(out):
    """
    function: _strip

    arguments:
     out: a result dict

    return:
     a copy safe to serialize as json

    description:
     Converts numpy arrays to lists so a record can be written out.
    """

    # convert every array-valued entry
    #
    clean = {}
    for key, value in out.items():
        if isinstance(value, np.ndarray):
            clean[key] = value.tolist()
        else:
            clean[key] = value

    # exit gracefully
    #
    return clean
#
# end of function


def _format_summary(out):
    """
    function: _format_summary

    arguments:
     out: a result dict

    return:
     a formatted multi-line string

    description:
     Renders one run as a readable report, so a result directory can be
     understood without loading any json.
    """

    # collect the settings that actually change an outcome
    #
    cfg = out["config"]
    lines = ["--- Lab 02 experiment ---", ""]
    for key in ("dataset", "dataset_path", "eval_path", "scaler", "method",
                "feature_map", "reps", "entanglement", "kernel_backend",
                "shots", "layers", "maxiter", "optimizer", "n_samples",
                "random_state"):
        lines.append(f"{key:16s}: {cfg.get(key)}")

    # report where the scores came from, the split and the encoded range
    #
    source = out.get("eval_source", "split")
    source = ("held-out split" if source == "split"
              else f"supplied file ({source})")
    lines += ["",
              f"{'evaluated on':16s}: {source}",
              f"{'train / eval':16s}: {out['n_train']} / {out['n_test']}",
              f"{'range width':16s}: {out['range_width']:.3f}",
              ""]

    # report the scores
    #
    for key in ("train_acc", "test_acc", "n_support", "C", "kernel_time",
                "train_time", "total_time"):
        if out.get(key) is not None:
            value = out[key]
            if isinstance(value, float):
                lines.append(f"{key:16s}: {value:.4f}")
            else:
                lines.append(f"{key:16s}: {value}")

    # report kernel concentration, the segment 2 diagnostic
    #
    if out.get("kernel_stats"):
        stats = out["kernel_stats"]
        lines += ["",
                  f"{'mean off-diag':16s}: {stats['mean_offdiag']:.4f}",
                  f"{'mean diagonal':16s}: {stats['mean_diag']:.4f}"]

    # exit gracefully
    #
    return "\n".join(lines) + "\n"
#
# end of function


def _print_row(out):
    """
    function: _print_row

    arguments:
     out: a result dict

    return:
     none

    description:
     Prints one run as a single aligned line.
    """

    # print the settings that vary plus the scores. mark the score when
    # it came from a supplied file rather than a held-out split.
    #
    cfg = out["config"]
    train = out.get("train_acc", float("nan"))
    test = out.get("test_acc", float("nan"))
    label = "eval" if out.get("eval_source", "split") != "split" else "test"
    print(f"{cfg['method']:10s} {cfg['dataset']:12s} "
          f"scaler={cfg['scaler']:9s} width={out['range_width']:.2f}  "
          f"train={train:.3f} {label}={test:.3f}  ({out['total_time']:.1f}s)")
#
# end of function


def _print_comparison(summary, rows):
    """
    function: _print_comparison

    arguments:
     summary: the verdict dict built by compare()
     rows:    each method's full result dict

    return:
     none

    description:
     Prints the classical against quantum comparison as a small table
     followed by a verdict in words.
    """

    # header describing what was compared, and on what
    #
    source = summary.get("eval_source", "split")
    label = "test" if source == "split" else "eval"
    print(f"\ndataset: {summary['dataset']}   scaler: {summary['scaler']}   "
          f"train/{label}: {summary['n_train']}/{summary['n_test']}")
    if source != "split":
        print(f"evaluated on supplied file: {source}")
    print(f"\n{'method':12s} {'train':>8s} {label:>8s} {'support':>9s} "
          f"{'time':>8s}")
    print("-" * 50)

    # one row per method
    #
    for method, res in rows.items():
        support = res.get("n_support")
        support = "-" if support is None else str(support)
        print(f"{method:12s} {res.get('train_acc', 0):8.3f} "
              f"{res.get('test_acc', 0):8.3f} {support:>9s} "
              f"{res.get('total_time', 0):7.1f}s")
    print("-" * 50)

    # state the verdict in words
    #
    if "verdict" in summary:
        print(f"verdict: {summary['verdict']} "
              f"(quantum - classical = {summary['gap']:+.3f})")
    else:
        print(f"best: {summary['best']}")
#
# end of function

#
# end of file