# Paper-scoped implementation; see SOURCE_PROVENANCE.json.
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal
import numpy as np
from numpy.typing import NDArray
from sklearn.model_selection import GroupShuffleSplit, StratifiedShuffleSplit, train_test_split
from sklearn.preprocessing import LabelEncoder

@dataclass
class HARData:
    X: NDArray[np.float64]
    y: NDArray[np.int64]
    subject: NDArray[Any] | None = None
    dataset: NDArray[Any] | None = None
    device: NDArray[Any] | None = None
    time_index: NDArray[Any] | None = None
    label_names: list[str] | None = None
    source: str = 'unknown'
    sampling_rate_hz: float | None = None
    metadata: dict[str, Any] | None = None

    def __post_init__(self) -> None:
        X = np.asarray(self.X, dtype=np.float64)
        if X.ndim == 2:
            X = X[:, :, None]
        if X.ndim != 3:
            raise ValueError(f'X must have shape (n,T,C) or (n,T); got {X.shape}')
        y = np.asarray(self.y)
        if y.ndim != 1 or y.shape[0] != X.shape[0]:
            raise ValueError('y must be 1D and have same n as X')
        self.X = X
        self.y = y.astype(np.int64, copy=False)
        n = X.shape[0]
        for name in ('subject', 'dataset', 'device', 'time_index'):
            val = getattr(self, name)
            if val is not None and len(val) != n:
                raise ValueError(f'{name} must have length n={n}')

    @property
    def n_classes(self) -> int:
        if self.label_names is not None:
            return int(len(self.label_names))
        return int(np.max(self.y)) + 1 if self.y.size else 0

    @property
    def shape(self) -> tuple[int, int, int]:
        return tuple((int(v) for v in self.X.shape))

    def subset(self, idx: NDArray[np.int64] | list[int]) -> 'HARData':
        idx_arr = np.asarray(idx, dtype=np.int64)
        return HARData(X=self.X[idx_arr], y=self.y[idx_arr], subject=None if self.subject is None else np.asarray(self.subject)[idx_arr], dataset=None if self.dataset is None else np.asarray(self.dataset)[idx_arr], device=None if self.device is None else np.asarray(self.device)[idx_arr], time_index=None if self.time_index is None else np.asarray(self.time_index)[idx_arr], label_names=self.label_names, source=self.source, sampling_rate_hz=self.sampling_rate_hz, metadata=None if self.metadata is None else dict(self.metadata))

def _label_encode(y: NDArray[Any]) -> tuple[NDArray[np.int64], list[str]]:
    if np.issubdtype(np.asarray(y).dtype, np.integer):
        y_int = np.asarray(y, dtype=np.int64)
        classes = [str(c) for c in sorted(np.unique(y_int).tolist())]
        le = LabelEncoder().fit(y_int)
        return (le.transform(y_int).astype(np.int64), [str(c) for c in le.classes_.tolist()])
    le = LabelEncoder().fit(np.asarray(y).astype(str))
    return (le.transform(np.asarray(y).astype(str)).astype(np.int64), [str(c) for c in le.classes_.tolist()])

def load_dataset(data: str | Path, *, synthetic_seed: int=42) -> HARData:
    path = Path(data)
    if not path.exists():
        raise FileNotFoundError(path)
    if path.suffix.lower() != '.npz':
        raise ValueError('Only .npz is supported by the native loader. Export WristHARBench to npz first.')
    obj = np.load(path, allow_pickle=True)
    if 'X' not in obj or 'y' not in obj:
        raise ValueError('NPZ file must contain keys X and y')
    raw_y = np.asarray(obj['y'])
    if 'label_names' in obj and np.issubdtype(raw_y.dtype, np.integer):
        y = raw_y.astype(np.int64, copy=False)
        names = [str(x) for x in np.asarray(obj['label_names'], dtype=object).tolist()]
        if y.size and (int(np.min(y)) < 0 or int(np.max(y)) >= len(names)):
            raise ValueError('Integer y is incompatible with the supplied label_names vocabulary')
    else:
        y, names = _label_encode(raw_y)

    def opt(key: str) -> NDArray[Any] | None:
        return obj[key] if key in obj else None
    return HARData(X=np.asarray(obj['X'], dtype=np.float64), y=y, subject=opt('subject'), dataset=opt('dataset'), device=opt('device'), time_index=opt('time_index'), label_names=names, source=str(path), sampling_rate_hz=float(np.asarray(obj['sampling_rate_hz']).reshape(-1)[0]) if 'sampling_rate_hz' in obj else None, metadata=dict(np.asarray(obj['metadata'], dtype=object).reshape(-1)[0]) if 'metadata' in obj else None)

def concatenate_har(*parts: HARData, source: str | None=None) -> HARData:
    if not parts:
        raise ValueError('At least one HARData object is required')
    shape_tail = parts[0].X.shape[1:]
    n_classes = parts[0].n_classes
    label_names = parts[0].label_names
    for part in parts[1:]:
        if part.X.shape[1:] != shape_tail:
            raise ValueError(f'Incompatible time/channel shape {part.X.shape[1:]} != {shape_tail}')
        if part.n_classes != n_classes:
            raise ValueError(f'Incompatible class vocabulary {part.n_classes} != {n_classes}')
        if label_names is not None and part.label_names is not None and (part.label_names != label_names):
            raise ValueError('label_names differ across HARData objects')

    def cat_optional(name: str) -> NDArray[Any] | None:
        vals = [getattr(part, name) for part in parts]
        if all((v is None for v in vals)):
            return None
        if any((v is None for v in vals)):
            raise ValueError(f'Cannot concatenate partially missing {name}')
        return np.concatenate([np.asarray(v) for v in vals], axis=0)
    metadata = dict(parts[0].metadata or {})
    metadata['concatenated_sources'] = [part.source for part in parts]
    return HARData(X=np.concatenate([part.X for part in parts], axis=0), y=np.concatenate([part.y for part in parts], axis=0), subject=cat_optional('subject'), dataset=cat_optional('dataset'), device=cat_optional('device'), time_index=cat_optional('time_index'), label_names=label_names, source=source or '+'.join((part.source for part in parts)), sampling_rate_hz=parts[0].sampling_rate_hz, metadata=metadata)