# Paper-scoped implementation; see SOURCE_PROVENANCE.json.
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Iterable, Sequence
import numpy as np
from numpy.typing import NDArray
from scipy.stats import kurtosis, skew
from sklearn.base import BaseEstimator, TransformerMixin

def _as_3d(X: Any) -> NDArray[np.float64]:
    arr = np.asarray(X, dtype=np.float64)
    if arr.ndim == 2:
        arr = arr[:, :, None]
    if arr.ndim != 3:
        raise ValueError(f'Expected X shape (n,T,C) or (n,T); got {arr.shape}')
    return arr

def _safe_zscore_per_window(X: NDArray[np.float64], eps: float=1e-06) -> NDArray[np.float64]:
    mu = np.nanmean(X, axis=1, keepdims=True)
    sd = np.nanstd(X, axis=1, keepdims=True)
    sd = np.where(sd < eps, 1.0, sd)
    return np.nan_to_num((X - mu) / sd, nan=0.0, posinf=0.0, neginf=0.0)

def _slope(y: NDArray[np.float64]) -> NDArray[np.float64]:
    T = y.shape[1]
    t = np.arange(T, dtype=np.float64)
    tc = t - t.mean()
    denom = float(np.dot(tc, tc)) or 1.0
    yc = y - np.mean(y, axis=1, keepdims=True)
    return np.einsum('t,ntc->nc', tc, yc) / denom

def _zero_crossing_rate(X: NDArray[np.float64]) -> NDArray[np.float64]:
    s = np.signbit(X)
    return np.mean(s[:, 1:, :] != s[:, :-1, :], axis=1)

def _segment_view(X: NDArray[np.float64], w: int) -> NDArray[np.float64]:
    n, T, C = X.shape
    if w <= 0:
        raise ValueError('w must be positive')
    seg_len = int(np.ceil(T / w))
    T2 = seg_len * w
    if T2 != T:
        pad = np.repeat(X[:, -1:, :], T2 - T, axis=1)
        X = np.concatenate([X, pad], axis=1)
    return X.reshape(n, w, seg_len, C)

def _spectral_features(X: NDArray[np.float64], bands: Sequence[tuple[float, float]]) -> NDArray[np.float64]:
    n, T, C = X.shape
    x = X - np.mean(X, axis=1, keepdims=True)
    fft = np.fft.rfft(x, axis=1)
    power = np.abs(fft) ** 2 / max(T, 1)
    freqs = np.fft.rfftfreq(T, d=1.0)
    total = np.sum(power[:, 1:, :], axis=1) + 1e-12 if power.shape[1] > 1 else np.ones((n, C))
    feats: list[NDArray[np.float64]] = []
    for lo, hi in bands:
        mask = (freqs >= lo) & (freqs < hi)
        if not np.any(mask):
            feats.append(np.zeros((n, C), dtype=np.float64))
        else:
            feats.append(np.sum(power[:, mask, :], axis=1) / total)
    p = power[:, 1:, :] / (total[:, None, :] if power.shape[1] > 1 else 1.0)
    entropy = -np.sum(np.where(p > 0, p * np.log(p + 1e-12), 0.0), axis=1) / np.log(max(p.shape[1], 2)) if power.shape[1] > 1 else np.zeros((n, C))
    centroid = np.sum(freqs[1:, None] * power[:, 1:, :], axis=1) / total if power.shape[1] > 1 else np.zeros((n, C))
    peak = freqs[1:][np.argmax(power[:, 1:, :], axis=1)] if power.shape[1] > 1 else np.zeros((n, C))
    feats.extend([entropy, centroid, peak])
    return np.concatenate([f.reshape(n, -1) for f in feats], axis=1)

def _autocorr_features(X: NDArray[np.float64], lags: Sequence[int]) -> NDArray[np.float64]:
    n, T, C = X.shape
    x = X - np.mean(X, axis=1, keepdims=True)
    denom = np.sum(x * x, axis=1) + 1e-12
    feats = []
    for lag in lags:
        if lag <= 0 or lag >= T:
            feats.append(np.zeros((n, C), dtype=np.float64))
        else:
            feats.append(np.sum(x[:, :-lag, :] * x[:, lag:, :], axis=1) / denom)
    return np.concatenate([f.reshape(n, -1) for f in feats], axis=1) if feats else np.empty((n, 0))

def _transition_features(x: NDArray[np.float64], n_bins: int=6) -> NDArray[np.float64]:
    n, T, C = x.shape
    outs: list[NDArray[np.float64]] = []
    for c in range(C):
        xc = x[:, :, c]
        q = np.quantile(xc, np.linspace(0, 1, n_bins + 1)[1:-1], axis=1).T
        states = np.sum(xc[:, :, None] > q[:, None, :], axis=2).astype(np.int64)
        mat = np.zeros((n, n_bins, n_bins), dtype=np.float64)
        for i in range(n):
            np.add.at(mat[i], (states[i, :-1], states[i, 1:]), 1.0)
        row = mat.sum(axis=2, keepdims=True)
        mat = np.where(row > 0, mat / np.maximum(row, 1.0), 0.0)
        hist = np.mean(np.eye(n_bins)[states], axis=1)
        flat = mat.reshape(n, -1)
        ent = -np.sum(np.where(flat > 0, flat * np.log(flat + 1e-12), 0.0), axis=1, keepdims=True)
        outs.extend([hist, flat, ent])
    return np.concatenate(outs, axis=1) if outs else np.empty((n, 0))

@dataclass
class WearableFeatureExtractor(BaseEstimator, TransformerMixin):
    zscore: bool = True
    include_raw_stats: bool = True
    include_magnitude: bool = True
    include_jerk: bool = True
    include_multiscale: bool = True
    segment_sizes: tuple[int, ...] = (4, 8, 16)
    include_spectral: bool = True
    spectral_bands: tuple[tuple[float, float], ...] = ((0.0, 0.05), (0.05, 0.15), (0.15, 0.3), (0.3, 0.51))
    include_autocorr: bool = True
    autocorr_lags: tuple[int, ...] = (1, 2, 4, 8, 16, 32)
    include_cross_channel: bool = True
    include_symbolic: bool = True
    symbolic_bins: int = 6
    feature_names_: list[str] = field(default_factory=list, init=False)

    def fit(self, X: Any, y: Any | None=None) -> 'WearableFeatureExtractor':
        self.transform(X[:min(len(X), 3)] if hasattr(X, '__len__') else X)
        return self

    def transform(self, X: Any) -> NDArray[np.float64]:
        X3 = _as_3d(X)
        if self.zscore:
            Xn = _safe_zscore_per_window(X3)
        else:
            Xn = np.nan_to_num(X3, nan=0.0, posinf=0.0, neginf=0.0)
        variants = [Xn]
        names_prefix = ['ch']
        if self.include_magnitude and Xn.shape[2] >= 2:
            mag = np.linalg.norm(Xn[:, :, :min(3, Xn.shape[2])], axis=2, keepdims=True)
            variants.append(mag)
            names_prefix.append('mag')
        if self.include_jerk:
            jerk = np.diff(Xn, axis=1, prepend=Xn[:, :1, :])
            variants.append(jerk)
            names_prefix.append('jerk')
            if Xn.shape[2] >= 2:
                variants.append(np.linalg.norm(jerk[:, :, :min(3, Xn.shape[2])], axis=2, keepdims=True))
                names_prefix.append('jerk_mag')
        feat_blocks: list[NDArray[np.float64]] = []
        feature_names: list[str] = []
        for V, prefix in zip(variants, names_prefix):
            n, T, C = V.shape
            if self.include_raw_stats:
                stats = [np.mean(V, axis=1), np.std(V, axis=1), np.min(V, axis=1), np.max(V, axis=1), np.median(V, axis=1), np.quantile(V, 0.25, axis=1), np.quantile(V, 0.75, axis=1), np.mean(V * V, axis=1), np.mean(np.diff(V, axis=1, prepend=V[:, :1, :]), axis=1), np.std(np.diff(V, axis=1, prepend=V[:, :1, :]), axis=1), _zero_crossing_rate(V), _slope(V), skew(V, axis=1, nan_policy='omit'), kurtosis(V, axis=1, nan_policy='omit')]
                block = np.concatenate([np.nan_to_num(s).reshape(n, -1) for s in stats], axis=1)
                feat_blocks.append(block)
                feature_names.extend([f'{prefix}_stat_{i}' for i in range(block.shape[1])])
            if self.include_multiscale:
                for w in self.segment_sizes:
                    if w > T:
                        continue
                    seg = _segment_view(V, int(w))
                    means = np.mean(seg, axis=2).reshape(n, -1)
                    stds = np.std(seg, axis=2).reshape(n, -1)
                    rng = (np.max(seg, axis=2) - np.min(seg, axis=2)).reshape(n, -1)
                    slopes = []
                    for j in range(int(w)):
                        slopes.append(_slope(seg[:, j, :, :]))
                    slopes_arr = np.stack(slopes, axis=1).reshape(n, -1)
                    block = np.concatenate([means, stds, rng, slopes_arr], axis=1)
                    feat_blocks.append(block)
                    feature_names.extend([f'{prefix}_seg{w}_{i}' for i in range(block.shape[1])])
            if self.include_spectral:
                block = _spectral_features(V, self.spectral_bands)
                feat_blocks.append(block)
                feature_names.extend([f'{prefix}_fft_{i}' for i in range(block.shape[1])])
            if self.include_autocorr:
                lags = tuple((l for l in self.autocorr_lags if l < T))
                block = _autocorr_features(V, lags)
                feat_blocks.append(block)
                feature_names.extend([f'{prefix}_acf_{i}' for i in range(block.shape[1])])
            if self.include_symbolic and V.shape[1] >= 8:
                Vsym = V if V.shape[2] <= 3 else V[:, :, :3]
                block = _transition_features(Vsym, n_bins=int(self.symbolic_bins))
                feat_blocks.append(block)
                feature_names.extend([f'{prefix}_sym_{i}' for i in range(block.shape[1])])
        if self.include_cross_channel and Xn.shape[2] > 1:
            n, _, C = Xn.shape
            corr_feats = []
            for i in range(C):
                for j in range(i + 1, C):
                    xi = Xn[:, :, i] - Xn[:, :, i].mean(axis=1, keepdims=True)
                    xj = Xn[:, :, j] - Xn[:, :, j].mean(axis=1, keepdims=True)
                    corr = np.sum(xi * xj, axis=1) / (np.sqrt(np.sum(xi * xi, axis=1) * np.sum(xj * xj, axis=1)) + 1e-12)
                    corr_feats.append(corr[:, None])
            if corr_feats:
                block = np.concatenate(corr_feats, axis=1)
                feat_blocks.append(block)
                feature_names.extend([f'corr_{i}' for i in range(block.shape[1])])
        if not feat_blocks:
            raise ValueError('No feature blocks enabled')
        F = np.concatenate(feat_blocks, axis=1)
        F = np.nan_to_num(F, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float64, copy=False)
        self.feature_names_ = feature_names
        return F

@dataclass
class RandomConvSketch(BaseEstimator, TransformerMixin):
    n_kernels: int = 256
    kernel_lengths: tuple[int, ...] = (7, 9, 11, 15)
    dilations: tuple[int, ...] = (1, 2, 4)
    random_state: int = 42
    zscore: bool = True
    kernels_: list[dict[str, Any]] = field(default_factory=list, init=False)

    def fit(self, X: Any, y: Any | None=None) -> 'RandomConvSketch':
        X3 = _as_3d(X)
        _, T, C = X3.shape
        rng = np.random.default_rng(self.random_state)
        kernels: list[dict[str, Any]] = []
        for _ in range(int(self.n_kernels)):
            length = int(rng.choice(self.kernel_lengths))
            dilation = int(rng.choice(self.dilations))
            max_span = (length - 1) * dilation + 1
            if max_span > T:
                dilation = max(1, (T - 1) // max(length - 1, 1))
            n_ch = int(rng.integers(1, min(C, 3) + 1))
            channels = np.sort(rng.choice(C, size=n_ch, replace=False)).astype(np.int64)
            weights = rng.normal(0, 1, size=(length, n_ch)).astype(np.float64)
            weights -= weights.mean(axis=0, keepdims=True)
            norm = np.linalg.norm(weights) + 1e-12
            weights /= norm
            bias = float(rng.normal(0, 0.25))
            kernels.append({'length': length, 'dilation': dilation, 'channels': channels, 'weights': weights, 'bias': bias})
        self.kernels_ = kernels
        return self

    def transform(self, X: Any) -> NDArray[np.float64]:
        X3 = _as_3d(X)
        if self.zscore:
            X3 = _safe_zscore_per_window(X3)
        n, T, _ = X3.shape
        feats = np.empty((n, len(self.kernels_) * 4), dtype=np.float64)
        for k, spec in enumerate(self.kernels_):
            length = int(spec['length'])
            dilation = int(spec['dilation'])
            channels = np.asarray(spec['channels'], dtype=np.int64)
            weights = np.asarray(spec['weights'], dtype=np.float64)
            span = (length - 1) * dilation + 1
            out_len = max(T - span + 1, 1)
            resp = np.zeros((n, out_len), dtype=np.float64)
            for ii in range(length):
                start = ii * dilation
                stop = start + out_len
                resp += X3[:, start:stop, :][:, :, channels] @ weights[ii]
            resp += float(spec['bias'])
            feats[:, 4 * k + 0] = np.mean(resp > 0.0, axis=1)
            feats[:, 4 * k + 1] = np.max(resp, axis=1)
            feats[:, 4 * k + 2] = np.mean(resp, axis=1)
            feats[:, 4 * k + 3] = np.std(resp, axis=1)
        return np.nan_to_num(feats, nan=0.0, posinf=0.0, neginf=0.0)

def _quantile_block(X: NDArray[np.float64], qs: Sequence[float]) -> NDArray[np.float64]:
    q = np.quantile(X, np.asarray(qs, dtype=np.float64), axis=1).transpose(1, 2, 0)
    return q.reshape(X.shape[0], -1)

def _gram_eig_features(X: NDArray[np.float64]) -> NDArray[np.float64]:
    n, _, c = X.shape
    feats = []
    for i in range(n):
        Xi = X[i] - X[i].mean(axis=0, keepdims=True)
        G = Xi.T @ Xi / max(Xi.shape[0] - 1, 1)
        vals = np.linalg.eigvalsh(G)
        vals = np.sort(np.maximum(vals, 0.0))[::-1]
        total = float(np.sum(vals)) + 1e-12
        ratios = vals / total
        pad = np.zeros((max(6 - 2 * c, 0),), dtype=np.float64)
        feats.append(np.concatenate([vals, ratios, pad], axis=0))
    return np.stack(feats, axis=0)

@dataclass
class IntervalDistributionSketch(BaseEstimator, TransformerMixin):
    n_random_intervals: int = 32
    random_state: int = 42
    quantiles: tuple[float, ...] = (0.1, 0.25, 0.5, 0.75, 0.9)
    include_dyadic: bool = True
    zscore: bool = True
    intervals_: list[tuple[int, int]] = field(default_factory=list, init=False)

    def fit(self, X: Any, y: Any | None=None) -> 'IntervalDistributionSketch':
        X3 = _as_3d(X)
        _, T, _ = X3.shape
        rng = np.random.default_rng(self.random_state)
        intervals: list[tuple[int, int]] = []
        if self.include_dyadic:
            for parts in (2, 4, 8):
                width = max(2, int(np.ceil(T / parts)))
                for start in range(0, T, width):
                    stop = min(T, start + width)
                    if stop - start >= 2:
                        intervals.append((start, stop))
        for _ in range(int(self.n_random_intervals)):
            lo = int(rng.integers(0, max(T - 2, 1)))
            max_len = max(3, T - lo)
            length = int(rng.integers(2, max_len + 1))
            hi = min(T, lo + length)
            if hi - lo >= 2:
                intervals.append((lo, hi))
        seen: set[tuple[int, int]] = set()
        self.intervals_ = []
        for it in intervals:
            if it not in seen:
                seen.add(it)
                self.intervals_.append(it)
        return self

    def transform(self, X: Any) -> NDArray[np.float64]:
        X3 = _as_3d(X)
        if self.zscore:
            X3 = _safe_zscore_per_window(X3)
        blocks: list[NDArray[np.float64]] = []
        for lo, hi in self.intervals_:
            V = X3[:, lo:hi, :]
            n = V.shape[0]
            stats = [np.mean(V, axis=1), np.std(V, axis=1), np.max(V, axis=1) - np.min(V, axis=1), np.mean(V * V, axis=1), _slope(V), _quantile_block(V, self.quantiles)]
            blocks.append(np.concatenate([s.reshape(n, -1) for s in stats], axis=1))
        if not blocks:
            return np.empty((X3.shape[0], 0), dtype=np.float64)
        return np.nan_to_num(np.concatenate(blocks, axis=1), nan=0.0, posinf=0.0, neginf=0.0)

@dataclass
class InvariantPhysicsFeatureExtractor(BaseEstimator, TransformerMixin):
    segment_sizes: tuple[int, ...] = (4, 8, 16)
    symbolic_bins: int = 6
    random_intervals: int = 24
    random_state: int = 42
    zscore: bool = True
    include_axis_features: bool = True
    interval_: IntervalDistributionSketch | None = field(default=None, init=False)

    def fit(self, X: Any, y: Any | None=None) -> 'InvariantPhysicsFeatureExtractor':
        X3 = _as_3d(X)
        self.interval_ = IntervalDistributionSketch(n_random_intervals=self.random_intervals, random_state=self.random_state, zscore=self.zscore).fit(X3, y)
        return self

    def transform(self, X: Any) -> NDArray[np.float64]:
        X3 = _as_3d(X)
        Xn = _safe_zscore_per_window(X3) if self.zscore else np.nan_to_num(X3)
        n, T, C = Xn.shape
        blocks: list[NDArray[np.float64]] = []
        if self.include_axis_features:
            axis = WearableFeatureExtractor(zscore=False, include_raw_stats=True, include_magnitude=False, include_jerk=False, include_multiscale=True, segment_sizes=self.segment_sizes, include_spectral=False, include_autocorr=False, include_cross_channel=True, include_symbolic=False).fit_transform(Xn)
            blocks.append(axis)
        base = Xn[:, :, :min(3, C)]
        mag = np.linalg.norm(base, axis=2, keepdims=True)
        jerk = np.diff(base, axis=1, prepend=base[:, :1, :])
        jerk_mag = np.linalg.norm(jerk, axis=2, keepdims=True)
        invariant_series = np.concatenate([mag, jerk_mag], axis=2)
        inv = WearableFeatureExtractor(zscore=False, include_raw_stats=True, include_magnitude=False, include_jerk=False, include_multiscale=True, segment_sizes=self.segment_sizes, include_spectral=True, include_autocorr=True, include_cross_channel=False, include_symbolic=True, symbolic_bins=self.symbolic_bins).fit_transform(invariant_series)
        blocks.append(inv)
        blocks.append(_gram_eig_features(base))
        mag_energy = np.mean(mag * mag, axis=1) + 1e-12
        jerk_energy = np.mean(jerk_mag * jerk_mag, axis=1)
        blocks.append((jerk_energy / mag_energy).reshape(n, -1))
        if self.interval_ is not None:
            blocks.append(self.interval_.transform(Xn))
        return np.nan_to_num(np.concatenate(blocks, axis=1), nan=0.0, posinf=0.0, neginf=0.0)

@dataclass
class RandomShapeletSketch(BaseEstimator, TransformerMixin):
    n_shapelets: int = 48
    lengths: tuple[int, ...] = (8, 12, 16, 24)
    dilations: tuple[int, ...] = (1, 2, 4)
    random_state: int = 42
    zscore: bool = True
    shapelets_: list[dict[str, Any]] = field(default_factory=list, init=False)

    def fit(self, X: Any, y: Any | None=None) -> 'RandomShapeletSketch':
        X3 = _as_3d(X)
        if self.zscore:
            X3 = _safe_zscore_per_window(X3)
        n, T, C = X3.shape
        rng = np.random.default_rng(self.random_state)
        self.shapelets_ = []
        for _ in range(int(self.n_shapelets)):
            length = int(rng.choice(self.lengths))
            dilation = int(rng.choice(self.dilations))
            span = (length - 1) * dilation + 1
            if span > T:
                dilation = max(1, (T - 1) // max(length - 1, 1))
                span = (length - 1) * dilation + 1
            i = int(rng.integers(0, n))
            c = int(rng.integers(0, C))
            start = int(rng.integers(0, max(T - span + 1, 1)))
            values = X3[i, start:start + span:dilation, c].astype(np.float64, copy=True)
            values = values - float(values.mean())
            sd = float(values.std()) or 1.0
            values = values / sd
            self.shapelets_.append({'values': values, 'channel': c, 'dilation': dilation})
        return self

    def transform(self, X: Any) -> NDArray[np.float64]:
        X3 = _as_3d(X)
        if self.zscore:
            X3 = _safe_zscore_per_window(X3)
        n, T, _ = X3.shape
        F = np.empty((n, len(self.shapelets_) * 2), dtype=np.float64)
        for k, spec in enumerate(self.shapelets_):
            sh = np.asarray(spec['values'], dtype=np.float64)
            c = int(spec['channel'])
            dilation = int(spec['dilation'])
            length = int(sh.shape[0])
            span = (length - 1) * dilation + 1
            out_len = max(T - span + 1, 1)
            best = np.full((n,), np.inf, dtype=np.float64)
            best_pos = np.zeros((n,), dtype=np.float64)
            for start in range(out_len):
                seg = X3[:, start:start + span:dilation, c]
                seg = seg - seg.mean(axis=1, keepdims=True)
                seg_sd = seg.std(axis=1, keepdims=True)
                seg = seg / np.where(seg_sd < 1e-06, 1.0, seg_sd)
                d = np.sqrt(np.mean((seg - sh[None, :]) ** 2, axis=1))
                improved = d < best
                best = np.where(improved, d, best)
                best_pos = np.where(improved, start / max(out_len - 1, 1), best_pos)
            F[:, 2 * k] = best
            F[:, 2 * k + 1] = best_pos
        return np.nan_to_num(F, nan=0.0, posinf=0.0, neginf=0.0)