#!/usr/bin/env python
#
# file: qmllab/data.py
#
# revision history:
#  20260810 (am): initial version for workshop lab 02
#
# Datasets and scaling for Lab 02.
#
# Data is stored in the IMLD csv format used throughout the lab:
#
#   # filename: ./data/imld_two_moons.csv
#   # classes: [0,1]
#   # colors: [#1f77b4,#ff7f0e]
#   # limits: [-1.0,1.0,-1.0,1.0]
#   #
#   0, 0.229829, 0.436207
#
# the label is column 0, the features follow, values carry six decimal
# places, and every header line begins with a '#'.
#
# The scaler is not a detail in this lab. A ZZFeatureMap encodes a
# product of features, 2(pi - x_i)(pi - x_j), so the width of the scaled
# range controls how fast the kernel oscillates. See SCALERS below and
# segment 2 of the notebook.
#------------------------------------------------------------------------------

# import system modules
#
from pathlib import Path

# import third-party modules
#
import numpy as np
from sklearn.datasets import make_circles, make_moons
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import (MaxAbsScaler, MinMaxScaler, Normalizer,
                                   RobustScaler, StandardScaler)

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

# the directory holding the shipped csv files
#
DATA_DIR = Path(__file__).resolve().parent.parent / "data"

# datasets this module knows how to generate
#
GENERATORS = ["two_moons", "two_spirals", "checkerboard", "yin_yang",
              "noisy_xor", "toroidal", "circles"]

# the default plotting colors and limits written into an IMLD header
#
DEFAULT_COLORS = ["#1f77b4", "#ff7f0e"]
DEFAULT_LIMITS = [-1.0, 1.0, -1.0, 1.0]

# scalers ordered by the width of the range they produce. that ordering
# is the segment 2 result: narrower is better for a ZZ feature map.
#
SCALERS = {
    "minmax":   MinMaxScaler,
    "maxabs":   MaxAbsScaler,
    "l1":       lambda: Normalizer(norm="l1"),
    "l2":       lambda: Normalizer(norm="l2"),
    "robust":   RobustScaler,
    "standard": StandardScaler,
    "none":     None,
}

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

def get_scaler(name):
    """
    function: get_scaler

    arguments:
     name: the scaler name, a key of SCALERS

    return:
     a fresh scaler instance, or None when name is "none"

    description:
     Builds a new scaler each call so that a scaler fitted on one split
     is never reused on another.
    """

    # reject an unknown name
    #
    if name not in SCALERS:
        raise ValueError(
            f"unknown scaler {name!r}; choose from {list(SCALERS)}")

    # build the scaler, passing through the no-scaling case
    #
    factory = SCALERS[name]

    # exit gracefully
    #
    return None if factory is None else factory()
#
# end of function


def write_imld(path, X, y, colors=None, limits=None):
    """
    function: write_imld

    arguments:
     path:   destination csv file
     X:      an (n, 2) array of features
     y:      an (n,) array of integer labels
     colors: one plot color per class (defaults to DEFAULT_COLORS)
     limits: [xmin, xmax, ymin, ymax] (defaults to DEFAULT_LIMITS)

    return:
     the path that was written

    description:
     Writes data in the IMLD csv format: a commented header, then one
     row per sample as "label, x1, x2" with six decimal places.
    """

    # normalize the arguments
    #
    path = Path(path)
    X = np.asarray(X, dtype=float)
    y = np.asarray(y).astype(int)
    classes = sorted(set(y.tolist()))
    colors = colors or DEFAULT_COLORS[:len(classes)]
    limits = limits or DEFAULT_LIMITS

    # make sure the destination directory exists
    #
    path.parent.mkdir(parents=True, exist_ok=True)

    # assemble the header block
    #
    header = [
        f"# filename: ./data/{path.name}",
        "# classes: [" + ",".join(str(c) for c in classes) + "]",
        "# colors: [" + ",".join(colors) + "]",
        "# limits: [" + ",".join(f"{v:.1f}" for v in limits) + "]",
        "#",
    ]

    # write the header followed by one row per sample
    #
    with open(path, "w") as fp:
        fp.write("\n".join(header) + "\n")
        for label, row in zip(y, X):
            values = ", ".join(f"{v:.6f}" for v in row)
            fp.write(f"{label}, {values}\n")

    # exit gracefully
    #
    return path
#
# end of function


def read_imld(path):
    """
    function: read_imld

    arguments:
     path: an IMLD csv file

    return:
     a tuple (X, y) of features and integer labels

    description:
     Reads the IMLD csv format. Header lines start with '#' and are
     skipped. Column 0 is the label and the remaining columns are the
     features.
    """

    # reject a missing file early so the message is useful
    #
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"dataset not found: {path}")

    # read every non-comment row
    #
    raw = np.genfromtxt(path, delimiter=",", comments="#")

    # guard against a single-row file collapsing to one dimension
    #
    if raw.ndim == 1:
        raw = raw.reshape(1, -1)

    # split the label column from the feature columns
    #
    y = raw[:, 0].astype(int)
    X = raw[:, 1:]

    # exit gracefully
    #
    return X, y
#
# end of function


def generate(name="two_moons", n_samples=1000, noise=0.15, random_state=42):
    """
    function: generate

    arguments:
     name:         one of GENERATORS
     n_samples:    how many points to produce
     noise:        the amount of jitter applied to the clean shape
     random_state: the seed controlling every random draw

    return:
     a tuple (X, y) with X scaled into [-1, 1] on both axes

    description:
     Builds one of the lab's two-dimensional, two-class datasets. Every
     shape is rescaled into [-1, 1] so that all of them share the IMLD
     header limits and remain comparable to one another.
    """

    # seed a local generator so the caller's global state is untouched
    #
    rng = np.random.RandomState(random_state)

    # two interleaving crescents
    #
    if name == "two_moons":
        X, y = make_moons(n_samples=n_samples, noise=noise,
                          random_state=random_state)

    # two concentric rings
    #
    elif name == "circles":
        X, y = make_circles(n_samples=n_samples, noise=noise, factor=0.4,
                            random_state=random_state)

    # two interleaved spiral arms
    #
    elif name == "two_spirals":
        half = n_samples // 2
        theta = np.sqrt(rng.rand(half)) * 4.0 * np.pi
        r = theta / (4.0 * np.pi)
        arm = np.c_[r * np.cos(theta), r * np.sin(theta)]
        X = np.vstack([arm, -arm])
        y = np.r_[np.zeros(half, int), np.ones(half, int)]
        X = X + noise * 0.5 * rng.randn(*X.shape)

    # a checkerboard of alternating cells
    #
    elif name == "checkerboard":
        X = rng.uniform(-1.0, 1.0, size=(n_samples, 2))
        cells = np.floor((X + 1.0) * 2.0).astype(int)
        y = ((cells[:, 0] + cells[:, 1]) % 2).astype(int)
        X = X + noise * 0.15 * rng.randn(*X.shape)

    # the classic yin-yang split, two lobes plus two eyes
    #
    elif name == "yin_yang":
        X = rng.uniform(-1.0, 1.0, size=(n_samples, 2))
        keep = np.linalg.norm(X, axis=1) <= 1.0
        X = X[keep]
        upper = np.linalg.norm(X - [0.0, 0.5], axis=1) <= 0.5
        lower = np.linalg.norm(X - [0.0, -0.5], axis=1) <= 0.5
        y = (X[:, 0] > 0).astype(int)
        y[upper] = 1
        y[lower] = 0
        y[np.linalg.norm(X - [0.0, 0.5], axis=1) <= 0.15] = 0
        y[np.linalg.norm(X - [0.0, -0.5], axis=1) <= 0.15] = 1
        X = X + noise * 0.1 * rng.randn(*X.shape)

    # exclusive-or on the sign of each coordinate
    #
    elif name == "noisy_xor":
        X = rng.uniform(-1.0, 1.0, size=(n_samples, 2))
        y = ((X[:, 0] > 0) ^ (X[:, 1] > 0)).astype(int)
        flip = rng.rand(len(y)) < (noise * 0.5)
        y[flip] = 1 - y[flip]

    # an inner annulus against an outer one
    #
    elif name == "toroidal":
        angle = rng.uniform(0.0, 2.0 * np.pi, size=n_samples)
        inner = rng.rand(n_samples) < 0.5
        radius = np.where(inner, 0.35, 0.85)
        radius = radius + noise * 0.25 * rng.randn(n_samples)
        X = np.c_[radius * np.cos(angle), radius * np.sin(angle)]
        y = (~inner).astype(int)

    # reject anything else
    #
    else:
        raise ValueError(
            f"unknown dataset {name!r}; choose from {GENERATORS}")

    # rescale both axes into [-1, 1] to match the IMLD header limits
    #
    span = X.max(axis=0) - X.min(axis=0)
    span[span == 0.0] = 1.0
    X = 2.0 * (X - X.min(axis=0)) / span - 1.0

    # exit gracefully
    #
    return X, np.asarray(y).astype(int)
#
# end of function


def load_dataset(name="two_moons", path=None, n_samples=None, noise=0.15,
                 random_state=42):
    """
    function: load_dataset

    arguments:
     name:         a dataset name; ignored when path is supplied
     path:         an IMLD csv file to read instead of generating
     n_samples:    subsample down to this many points when given
     noise:        jitter passed to the generator
     random_state: the seed for generation and for subsampling

    return:
     a tuple (X, y)

    description:
     Resolves a dataset from one of three places, in order: an explicit
     csv path, a shipped imld_<name>.csv file, or the generator. The
     ad_hoc dataset is a special case handled by qiskit-machine-learning.
    """

    # an explicit path always wins
    #
    if path is not None:
        X, y = read_imld(path)

    # ad_hoc is generated by qiskit, not by this module
    #
    elif name == "ad_hoc":
        X, y = _load_ad_hoc(n_samples, random_state)

    # otherwise prefer a shipped file and fall back to generating one
    #
    else:
        shipped = DATA_DIR / f"imld_{name}.csv"
        if shipped.exists():
            X, y = read_imld(shipped)
        else:
            X, y = generate(name, n_samples or 1000, noise, random_state)

    # subsample when the caller asked for fewer points than we have
    #
    if n_samples is not None and len(X) > n_samples:
        rng = np.random.RandomState(random_state)
        idx = rng.choice(len(X), n_samples, replace=False)
        X, y = X[idx], y[idx]

    # exit gracefully
    #
    return X, y
#
# end of function


def _load_ad_hoc(n_samples, random_state):
    """
    function: _load_ad_hoc

    arguments:
     n_samples:    the total number of points to build
     random_state: the seed for the global qiskit generator

    return:
     a tuple (X, y)

    description:
     Builds the Havlicek ad_hoc dataset, which is constructed so that a
     depth-2 ZZ feature map is the natural kernel for it. The generator
     draws from a global rng, so we seed it here or the lab would not
     reproduce.
    """

    # import here so that a plain import of this module stays light
    #
    from qiskit_machine_learning.datasets import ad_hoc_data
    from qiskit_machine_learning.utils import algorithm_globals

    # seed the global generator used inside ad_hoc_data
    #
    algorithm_globals.random_seed = random_state

    # build a balanced train and test half, then merge them
    #
    half = (n_samples or 200) // 2
    x_tr, y_tr, x_te, y_te = ad_hoc_data(
        training_size=half // 2, test_size=half // 2, n=2, gap=0.3,
        one_hot=False)

    # exit gracefully
    #
    return np.vstack([x_tr, x_te]), np.concatenate([y_tr, y_te]).astype(int)
#
# end of function


def prepare(name="two_moons", scaler="minmax", path=None, n_samples=None,
            test_size=0.2, random_state=42, noise=0.15, eval_path=None,
            eval_n_samples=None):
    """
    function: prepare

    arguments:
     name:           a dataset name; ignored when path is supplied
     scaler:         a key of SCALERS
     path:           an IMLD csv file to read instead of generating
     n_samples:      subsample the training data to this many points
     test_size:      the held-out fraction, ignored when eval_path is set
     random_state:   the seed for loading and splitting
     noise:          jitter passed to the generator
     eval_path:      an IMLD csv holding a separate evaluation set
     eval_n_samples: subsample the evaluation data to this many points

    return:
     a tuple (X_train, X_eval, y_train, y_eval, scaler)

    description:
     Loads, splits, then scales. The scaler is always fitted on the
     training data alone, which is what keeps the reported score
     meaningful.

     Two modes:

      - no eval_path: the dataset is split into train and test using
        test_size, stratified by label. this is the default and what
        the notebook uses.

      - eval_path given: no split happens at all. the whole dataset
        becomes training data and the evaluation set is read from that
        file. use this when you already hold out your own test data.

     In both cases the return shape is identical, so every caller works
     unchanged.
    """

    # load the training data
    #
    X, y = load_dataset(name, path=path, n_samples=n_samples, noise=noise,
                        random_state=random_state)

    # an explicit evaluation file replaces the split entirely
    #
    if eval_path is not None:
        X_train, y_train = X, y
        X_eval, y_eval = read_imld(eval_path)

        # subsample the evaluation set when the caller asked for that
        #
        if eval_n_samples is not None and len(X_eval) > eval_n_samples:
            rng = np.random.RandomState(random_state)
            idx = rng.choice(len(X_eval), eval_n_samples, replace=False)
            X_eval, y_eval = X_eval[idx], y_eval[idx]

        # a mismatch here is a silent disaster later, so check it now
        #
        if X_eval.shape[1] != X_train.shape[1]:
            raise ValueError(
                f"evaluation data has {X_eval.shape[1]} features but the "
                f"training data has {X_train.shape[1]}")

    # otherwise hold out a stratified split as usual
    #
    else:
        X_train, X_eval, y_train, y_eval = train_test_split(
            X, y, test_size=test_size, random_state=random_state, stratify=y)

    # fit the scaler on training data only, then apply it to both
    #
    scale = get_scaler(scaler)
    if scale is not None:
        X_train = scale.fit_transform(X_train)
        X_eval = scale.transform(X_eval)

    # exit gracefully
    #
    return X_train, X_eval, y_train, y_eval, scale
#
# end of function


def range_width(X):
    """
    function: range_width

    arguments:
     X: an array of scaled features

    return:
     the width of the encoded range as a float

    description:
     Reports the single number that predicts quantum kernel accuracy in
     segment 2. A wide range makes the ZZ feature map's product term
     oscillate faster than the data varies.
    """

    # exit gracefully
    #
    return float(X.max() - X.min())
#
# end of function

#
# end of file