"""Pinned paper inputs and exact participant/split ordering."""
from pathlib import Path
import json
import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
REPO_ID = 'Zipeng365/WristHARBench'
REVISION = '5e6af3c5c6dca050e73c06f49272f9656c0f259f'


def preprocess_signal(signal):
    # Arrow -> pandas represents a nested list as an object array of row arrays.
    # Stack its rows explicitly; load_dataset and local Parquet now share a path.
    a = np.stack(signal).astype(np.float64)
    a = pd.DataFrame(a).ffill().bfill().to_numpy(dtype=np.float64)
    if a.ndim != 2 or not np.isfinite(a).all():
        raise ValueError('Non-finite signal after registered imputation')
    mu, sd = a.mean(0, keepdims=True), a.std(0, keepdims=True)
    sd[sd == 0] = 1
    return (a - mu) / sd


def config_name(dataset_id):
    if dataset_id == 'capture24_wearable_activity_v1':
        return dataset_id + '_paper_snapshot_20260828'
    return dataset_id


def export_task(dataset_id, output, *, local_data=None, folds=(0, 1, 2, 3, 4)):
    scope = json.loads((ROOT / 'configs/scope.json').read_text())
    if dataset_id not in scope['datasets']:
        raise ValueError(f'Unknown paper dataset: {dataset_id}')
    settings = json.loads((ROOT / 'configs/datasets.json').read_text())[dataset_id]
    if local_data is None:
        from datasets import load_dataset
        rows = load_dataset(REPO_ID, config_name(dataset_id), split='full', revision=REVISION)
        frame = rows.to_pandas()
    else:
        files = sorted((Path(local_data) / 'data' / config_name(dataset_id)).glob('*.parquet'))
        if not files:
            raise FileNotFoundError(f'No Parquet files for {config_name(dataset_id)}')
        frame = pd.concat([pd.read_parquet(f) for f in files], ignore_index=True)
    if len(frame) != settings['sequences'] or frame.series_id.duplicated().any():
        raise ValueError('Input scope does not match the frozen paper contract')
    frame['series_id'] = frame.series_id.astype(str)
    frame['subject_id'] = frame.subject_id.astype(str)
    vocabulary = sorted(frame.label.astype(str).unique())
    encoded = {label: i for i, label in enumerate(vocabulary)}
    ordering = pd.read_csv(ROOT / f'configs/ordering/{dataset_id}.csv.gz', dtype={'series_id': str, 'subject_id': str})
    if set(ordering.series_id) != set(frame.series_id):
        raise ValueError('Ordering sidecar does not cover the downloaded rows')
    output = Path(output)
    for fold in folds:
        if fold not in range(5):
            raise ValueError('Only the five paper splits are supported')
        groups = []
        for split in ['train', 'validation', 'test']:
            side = ordering[(ordering.fold == fold) & (ordering.split == split)].sort_values('split_row')
            part = side.merge(frame, on=['series_id', 'subject_id'], how='left', validate='one_to_one')
            if not part[f'fold_{fold}'].eq(split).all():
                raise ValueError('Published split assignments disagree with the frozen export')
            groups.append(set(part.subject_id))
            X = np.stack([preprocess_signal(x) for x in part.signal]).astype(np.float32)
            if X.shape[1:] != (settings['window_steps'], len(settings['channels'])):
                raise ValueError(f'Unexpected shape for {dataset_id}: {X.shape}')
            folder = output / dataset_id / f'fold{fold}'
            folder.mkdir(parents=True, exist_ok=True)
            name = 'valid' if split == 'validation' else split
            destination = folder / f'{name}.npz'
            if destination.exists():
                raise FileExistsError(f'Refusing to overwrite existing inputs: {destination}')
            np.savez_compressed(destination, X=X,
                                y=np.array([encoded[str(x)] for x in part.label], dtype=np.int64),
                                subject=part.subject_id.to_numpy(dtype=str),
                                series_id=part.series_id.to_numpy(dtype=str),
                                time_index=part.time_index.to_numpy(dtype=float),
                                label_names=np.asarray(vocabulary, dtype=str),
                                sampling_rate_hz=np.asarray([settings['sampling_rate_hz']]),
                                metadata_json=np.asarray(json.dumps(settings.get('metadata', {}))))
        if any(groups[i] & groups[j] for i, j in [(0, 1), (0, 2), (1, 2)]):
            raise ValueError('Participant leakage across split partitions')
    return {'dataset_id': dataset_id, 'windows': len(frame), 'folds': list(folds),
            'source_repo': REPO_ID, 'revision': REVISION, 'configuration': config_name(dataset_id)}


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Export exact paper NPZ inputs from pinned WristHARBench')
    parser.add_argument('--dataset', required=True)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--fold', type=int, action='append')
    parser.add_argument('--local-data', type=Path)
    args = parser.parse_args()
    print(json.dumps(export_task(args.dataset, args.output, local_data=args.local_data,
                                folds=tuple(args.fold) if args.fold is not None else tuple(range(5))), indent=2))