"""Datasets for federated grid fault detection.

Two datasets share **one loader interface** —
``load_partition(client_index, num_clients, dataset=...)`` returns numpy
``(X_train, y_train, X_test, y_test)`` — so swapping data never touches the FL or
topology code:

- ``pq_synth`` (default): IEEE-1159-style power-quality waveforms generated *on
  the node* — no download/account, full run in minutes. 8 classes, 2 channels.
- ``vsb``: **real** partial-discharge measurements on medium-voltage covered
  conductors (Kaggle "VSB Power Line Fault Detection"). Binary (PD vs no-PD), one
  channel. Downloaded + cached on the node on first use (Kaggle creds required).

Each dataset advertises its shape via ``spec(name)`` -> :class:`DatasetSpec`
(number of classes, input channels, signal length, which class is the headline
fault). The 1D-CNN's global-average-pool head is length-invariant, so only
``num_classes`` and ``n_channels`` actually reach the model.
"""

from __future__ import annotations

import os
import sys
from typing import List, Optional, Tuple

import numpy as np

# --- pq_synth signal/label geometry (also the default DatasetSpec) -----------
CLASS_NAMES = [
    "normal",        # 0
    "sag",           # 1
    "swell",         # 2
    "interruption",  # 3
    "harmonics",     # 4
    "transient",     # 5  <- incipient-fault signature (headline)
    "flicker",       # 6
    "notch",         # 7
]
NUM_CLASSES = len(CLASS_NAMES)
FAULT_CLASS = 5           # transient / incipient-fault class
NORMAL_CLASS = 0
SIGNAL_LEN = 256          # samples per window
N_CHANNELS = 2            # voltage + current
CYCLES = 4.0              # fundamental cycles spanned by one window

# --- vsb geometry ------------------------------------------------------------
# Each raw VSB signal is 800,000 samples of one phase over a 20 ms window. We
# peak-preserving-downsample it to VSB_SIGNAL_LEN so the same compact CNN applies;
# 800000 is divisible by 400, so each output sample summarizes a 2000-sample bin.
VSB_SIGNAL_LEN = 400
VSB_COMPETITION = "vsb-power-line-fault-detection"


class DatasetSpec:
    """Shape + label geometry a topology needs to build the model and metrics."""

    def __init__(self, name, class_names, n_channels, signal_len,
                 fault_class, normal_class):
        self.name = name
        self.class_names = class_names
        self.num_classes = len(class_names)
        self.n_channels = n_channels
        self.signal_len = signal_len
        self.fault_class = fault_class
        self.normal_class = normal_class


PQ_SYNTH_SPEC = DatasetSpec("pq_synth", CLASS_NAMES, N_CHANNELS, SIGNAL_LEN,
                            FAULT_CLASS, NORMAL_CLASS)
VSB_SPEC = DatasetSpec("vsb", ["no_pd", "pd_fault"], 1, VSB_SIGNAL_LEN,
                       fault_class=1, normal_class=0)


def spec(dataset: str = "pq_synth", allow_fallback: bool = True) -> DatasetSpec:
    """Return the :class:`DatasetSpec` for ``dataset`` (pq_synth on unknown+fallback)."""
    if dataset == "vsb":
        return VSB_SPEC
    if dataset == "pq_synth":
        return PQ_SYNTH_SPEC
    if not allow_fallback:
        raise NotImplementedError(
            f"dataset '{dataset}' is roadmap (see weave.md §4); implemented: "
            f"'pq_synth', 'vsb'. Set allow_fallback=True to use pq_synth."
        )
    print(f"### PROGRESS: dataset '{dataset}' unknown — spec falls back to pq_synth",
          file=sys.stderr)
    return PQ_SYNTH_SPEC


# =============================================================================
# pq_synth — on-node synthetic power-quality waveforms
# =============================================================================

def _one_waveform(cls: int, rng: np.random.Generator, snr_db: float) -> np.ndarray:
    """Synthesize a single (N_CHANNELS, SIGNAL_LEN) power-quality window."""
    t = np.linspace(0.0, CYCLES * 2.0 * np.pi, SIGNAL_LEN, endpoint=False)
    phase = rng.uniform(0, 2 * np.pi)
    # Voltage and (slightly phase-shifted) current share the disturbance.
    v = np.sin(t + phase)
    i = np.sin(t + phase - rng.uniform(0.1, 0.5))

    # A disturbance occupies a contiguous segment of the window.
    a = rng.integers(SIGNAL_LEN // 8, SIGNAL_LEN // 2)
    b = min(SIGNAL_LEN, a + rng.integers(SIGNAL_LEN // 8, SIGNAL_LEN // 2))
    seg = slice(a, b)

    if cls == 0:                                   # normal
        pass
    elif cls == 1:                                 # sag (voltage dip)
        f = rng.uniform(0.3, 0.8)
        v[seg] *= f
        i[seg] *= (2.0 - f)                        # current rises as voltage sags
    elif cls == 2:                                 # swell
        f = rng.uniform(1.1, 1.6)
        v[seg] *= f
        i[seg] *= f
    elif cls == 3:                                 # interruption
        v[seg] *= rng.uniform(0.0, 0.1)
        i[seg] *= rng.uniform(0.0, 0.1)
    elif cls == 4:                                 # harmonics
        for h, amp in ((3, 0.25), (5, 0.15), (7, 0.1)):
            ph = rng.uniform(0, 2 * np.pi)
            v += amp * rng.uniform(0.6, 1.4) * np.sin(h * t + ph)
            i += amp * rng.uniform(0.6, 1.4) * np.sin(h * t + ph)
    elif cls == 5:                                 # transient / incipient fault
        # Short, high-frequency decaying spike — the insulation-degradation
        # precursor the proposal targets.
        c = rng.integers(a, b)
        width = rng.integers(4, 16)
        env = np.exp(-np.arange(SIGNAL_LEN) / width)
        env = np.roll(env, c)
        env[:c] = 0.0
        hf = np.sin(rng.uniform(8, 20) * t)
        spike = rng.uniform(0.5, 1.5) * env * hf
        v += spike
        i += rng.uniform(0.5, 1.5) * spike
    elif cls == 6:                                 # flicker (amplitude modulation)
        fm = 1.0 + rng.uniform(0.1, 0.3) * np.sin(rng.uniform(0.5, 2.0) * t)
        v *= fm
        i *= fm
    elif cls == 7:                                 # notch
        step = rng.integers(8, 20)
        depth = rng.uniform(0.4, 0.9)
        v[::step] *= depth
        i[::step] *= depth

    x = np.stack([v, i], axis=0)

    # Additive white noise at the client's SNR.
    sig_power = np.mean(x ** 2)
    noise_power = sig_power / (10.0 ** (snr_db / 10.0))
    x = x + rng.normal(0.0, np.sqrt(noise_power), size=x.shape)
    return x.astype(np.float32)


def _generate(labels: np.ndarray, rng: np.random.Generator, snr_db: float) -> np.ndarray:
    return np.stack([_one_waveform(int(c), rng, snr_db) for c in labels], axis=0)


def _load_pq_synth(
    client_index: int, num_clients: int, n_train: int, n_test: int,
    alpha: float, seed: int,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Non-IID synthetic partition: Dirichlet(alpha) class mix + per-client SNR."""
    rng = np.random.default_rng(seed * 100003 + client_index)
    class_probs = rng.dirichlet(alpha * np.ones(NUM_CLASSES))
    snr_db = float(rng.uniform(20.0, 40.0))  # client-specific channel quality

    y_train = rng.choice(NUM_CLASSES, size=n_train, p=class_probs)
    y_test = rng.choice(NUM_CLASSES, size=n_test, p=class_probs)
    X_train = _generate(y_train, rng, snr_db)
    X_test = _generate(y_test, rng, snr_db)
    return X_train, y_train.astype(np.int64), X_test, y_test.astype(np.int64)


# =============================================================================
# vsb — real partial-discharge measurements (Kaggle), downloaded on the node
# =============================================================================

def _vsb_data_dir() -> str:
    """Optional pre-staged VSB dir — override with $FL_DATA_DIR (default ./data).

    If ``metadata_train.csv`` + ``train.parquet`` are already here, they are used
    as-is (offline / manual staging). Otherwise the files come from Kaggle via
    kagglehub (see :func:`_ensure_vsb`).
    """
    return os.path.join(os.environ.get("FL_DATA_DIR", "data"), "vsb")


def _ensure_vsb() -> Tuple[str, str]:
    """Return ``(metadata_csv, train_parquet)`` paths, fetching from Kaggle if absent.

    Uses **kagglehub**, which authenticates with a modern ``KGAT_*`` access token
    read from ``$KAGGLE_API_TOKEN`` or ``~/.kaggle/access_token`` (falling back to
    legacy ``KAGGLE_USERNAME``/``KAGGLE_KEY`` or ``~/.kaggle/kaggle.json``), and
    caches downloads under ``~/.cache/kagglehub`` (override with ``$KAGGLEHUB_CACHE``).
    Requires one-time acceptance of the competition rules. A pre-staged copy under
    ``$FL_DATA_DIR/vsb`` short-circuits the download (used for offline testing).
    """
    staged = _vsb_data_dir()
    meta = os.path.join(staged, "metadata_train.csv")
    parq = os.path.join(staged, "train.parquet")
    if os.path.exists(meta) and os.path.exists(parq):
        return meta, parq

    try:
        import kagglehub  # lazy: pq_synth-only nodes need neither kagglehub nor a token
    except Exception as exc:  # noqa: BLE001
        raise RuntimeError(
            "VSB needs the 'kagglehub' package (pip install kagglehub). "
            f"Import failed: {exc}"
        ) from exc

    paths = {}
    for fname in ("metadata_train.csv", "train.parquet"):
        print(f"### PROGRESS: [vsb] fetching {fname} via kagglehub "
              f"(cached after first use)...", file=sys.stderr)
        try:
            paths[fname] = kagglehub.competition_download(VSB_COMPETITION, path=fname)
        except Exception as exc:  # noqa: BLE001
            raise RuntimeError(
                f"kagglehub download of '{fname}' from '{VSB_COMPETITION}' failed. "
                "HTTP 401 => bad/missing KGAT_ token (set $KAGGLE_API_TOKEN or "
                "~/.kaggle/access_token); HTTP 403 => competition rules not accepted "
                f"at kaggle.com/competitions/{VSB_COMPETITION}/rules. Error: {exc}"
            ) from exc
    return paths["metadata_train.csv"], paths["train.parquet"]


def _downsample_peak(sig: np.ndarray, out_len: int) -> np.ndarray:
    """Peak-preserving downsample of a 1-D signal to ``out_len`` samples.

    PD pulses are microsecond-scale spikes; plain decimation would step right over
    them. We bin the signal and keep each bin's largest-magnitude sample, so the
    discharge transients that define the fault survive the length reduction.
    """
    n = len(sig)
    bin_size = n // out_len
    sig = np.asarray(sig[: bin_size * out_len], dtype=np.float32).reshape(out_len, bin_size)
    idx = np.argmax(np.abs(sig), axis=1)
    return sig[np.arange(out_len), idx]


def _read_vsb_signals(parquet_path: str, signal_ids: List[int], out_len: int,
                      batch: int = 256) -> np.ndarray:
    """Read the given signal columns from the parquet, downsampled to ``out_len``.

    Reads in column batches so peak memory stays bounded (one raw 800k-sample
    column is ~0.8 MB; we downsample and discard it immediately).
    """
    import pyarrow.parquet as pq

    pf = pq.ParquetFile(parquet_path)
    available = set(pf.schema.names)
    out = np.empty((len(signal_ids), out_len), dtype=np.float32)
    pos = 0
    for i in range(0, len(signal_ids), batch):
        names = [str(s) for s in signal_ids[i:i + batch] if str(s) in available]
        if not names:
            continue
        table = pf.read(columns=names)
        for name in names:                         # index by name -> order-safe
            sig = table.column(name).to_numpy(zero_copy_only=False)
            out[pos] = _downsample_peak(sig, out_len)
            pos += 1
    return out[:pos]


def _load_vsb(
    client_index: int, num_clients: int, n_train: int, n_test: int, seed: int,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Real VSB partition: this client owns a disjoint set of measurement lines.

    Federated split by ``id_measurement % num_clients`` — each client sees a
    distinct subset of physical lines (and thus its own PD-prevalence / noise
    profile), which is genuine, not-injected non-IID. ``n_train``/``n_test`` cap
    the number of single-phase signals used per client (VSB is large + imbalanced).
    """
    import pandas as pd

    meta_csv, parq = _ensure_vsb()
    meta = pd.read_csv(meta_csv)

    mine = meta[(meta["id_measurement"] % num_clients) == client_index]
    if mine.empty:
        raise RuntimeError(
            f"VSB: no measurements fell to client {client_index}/{num_clients} — "
            f"check num_clients."
        )

    rng = np.random.default_rng(seed * 100003 + client_index)
    sig_ids = mine["signal_id"].to_numpy()
    targets = mine["target"].to_numpy().astype(np.int64)
    perm = rng.permutation(len(sig_ids))
    sig_ids, targets = sig_ids[perm], targets[perm]

    cap = n_train + n_test
    if len(sig_ids) > cap:
        sig_ids, targets = sig_ids[:cap], targets[:cap]
    k = len(sig_ids)
    n_te = max(1, min(n_test, k // 5))             # ~20% test, bounded by n_test

    te_ids, y_test = sig_ids[:n_te], targets[:n_te]
    tr_ids, y_train = sig_ids[n_te:], targets[n_te:]

    X_train = _read_vsb_signals(parq, list(tr_ids), VSB_SIGNAL_LEN)
    X_test = _read_vsb_signals(parq, list(te_ids), VSB_SIGNAL_LEN)

    # Per-signal standardization, then add the single channel axis -> (N, 1, L).
    def _norm(x):
        mu = x.mean(axis=1, keepdims=True)
        sd = x.std(axis=1, keepdims=True) + 1e-6
        return ((x - mu) / sd)[:, None, :].astype(np.float32)

    return _norm(X_train), y_train, _norm(X_test), y_test


# =============================================================================
# public loader
# =============================================================================

def load_partition(
    client_index: int,
    num_clients: int,
    dataset: str = "pq_synth",
    n_train: int = 2000,
    n_test: int = 500,
    alpha: float = 0.5,
    seed: int = 0,
    allow_fallback: bool = True,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Return ``(X_train, y_train, X_test, y_test)`` for one client.

    ``pq_synth`` builds a Dirichlet(``alpha``) non-IID synthetic partition on the
    fly. ``vsb`` reads a disjoint set of real measurement lines (``alpha`` unused;
    the non-IID split is physical). Unknown names fall back to ``pq_synth`` when
    ``allow_fallback`` (else ``NotImplementedError``).
    """
    if dataset == "vsb":
        return _load_vsb(client_index, num_clients, n_train, n_test, seed)
    if dataset != "pq_synth":
        if not allow_fallback:
            raise NotImplementedError(
                f"dataset '{dataset}' is roadmap (see weave.md §4); implemented: "
                f"'pq_synth', 'vsb'. Set allow_fallback=True to use pq_synth."
            )
        print(
            f"### PROGRESS: dataset '{dataset}' not built — falling back to pq_synth",
            file=sys.stderr,
        )
    return _load_pq_synth(client_index, num_clients, n_train, n_test, alpha, seed)


def class_distribution(y: np.ndarray, dataset_spec: Optional[DatasetSpec] = None) -> dict:
    """Count labels by class name (uses ``dataset_spec`` names, else pq_synth)."""
    sp = dataset_spec or PQ_SYNTH_SPEC
    counts = np.bincount(np.asarray(y, dtype=np.int64), minlength=sp.num_classes)
    return {sp.class_names[i]: int(c) for i, c in enumerate(counts[: sp.num_classes])}


if __name__ == "__main__":
    # Sanity check without torch: `python3 -m fl.data [dataset]`
    ds = sys.argv[1] if len(sys.argv) > 1 else "pq_synth"
    sp = spec(ds)
    xtr, ytr, xte, yte = load_partition(client_index=0, num_clients=4, dataset=ds)
    print(f"[{ds}] X_train", xtr.shape, xtr.dtype, "y_train", ytr.shape)
    print(f"[{ds}] X_test ", xte.shape, "class mix:", class_distribution(ytr, sp))
