# Examples: zyplib.dataset Typical Workflows

## Scenario A: Fixed train/val/test split (recommended)

Start from raw EEG files, split at file level, build DataLoaders.

```python
import numpy as np
from scipy.io import loadmat
from glob import glob

from zyplib.dataset import train_test_pipeline
from zyplib.dataset.torch_data import make_basic_dataloader, analyze_dataloader
from zyplib.preprocessing import resample_eeg, scale_eeg, segment_signal

# ---- 1. Define file lists ----
asd_files = glob('data/ASD/*.mat')
td_files = glob('data/TD/*.mat')

# ---- 2. Define callback functions ----
FS, NEW_FS = 1000, 250

def read_mat(file_path: str) -> np.ndarray:
    return loadmat(file_path)['data']  # [C, T]

def preprocess(signal: np.ndarray) -> np.ndarray:
    signal = resample_eeg(signal, FS, NEW_FS)
    signal = scale_eeg(signal, 'zscore', 'epoch')
    return signal

def segment(signal: np.ndarray) -> np.ndarray:
    return segment_signal(signal, NEW_FS, seg_duration=4)  # [N, C, T]

# ---- 3. File-level 3-way split + load ----
X_train, X_val, X_test, y_train, y_val, y_test = train_test_pipeline(
    split_size=(0.7, 0.15, 0.15),
    all_files=(asd_files, td_files),  # grouped -> stratified split
    random_seed=42,
    fn_pipeline='fileseg',
    pipeline_kwargs=dict(
        fn_read=read_mat,
        fn_segment=segment,
        label=lambda file_path: 1 if 'ASD' in file_path else 0,
        fn_before_segment=preprocess,
        random_shuffle=42,
    ),
)

# ---- 4. Build DataLoaders ----
train_loader, val_loader, test_loader = make_basic_dataloader(
    (X_train, y_train),
    (X_val, y_val),
    (X_test, y_test),
    batch_size=64,
    shuffle='shuffle-first',
)

# ---- 5. Verify ----
analyze_dataloader(train_loader, prefix_name='Train')
analyze_dataloader(val_loader, prefix_name='Val')
analyze_dataloader(test_loader, prefix_name='Test')
```

## Scenario B: File-level K-fold cross-validation

```python
from zyplib.dataset import kfold_pipeline
from zyplib.dataset.torch_data import make_basic_dataloader

# Reuse read_mat, preprocess, segment from Scenario A

all_files = glob('data/ASD/*.mat') + glob('data/TD/*.mat')

# You can also pass a custom pipeline function
def my_load_fn(files):
    from zyplib.dataset import basic_pipeline
    return basic_pipeline(
        files,
        fn_read=read_mat,
        fn_segment=segment,
        label=lambda file_path: 1 if 'ASD' in file_path else 0,
        fn_before_segment=preprocess,
        random_shuffle=42,
    )

results = {}
for fold, (X_train, X_val, y_train, y_val) in enumerate(
    kfold_pipeline(
        k=5,
        all_files=(asd_files, td_files),  # grouped to ensure class balance
        random_seed=42,
        fn_pipeline=my_load_fn,
    ),
    start=1,
):
    train_loader, val_loader = make_basic_dataloader(
        (X_train, y_train),
        (X_val, y_val),
        batch_size=64,
        shuffle='shuffle-first',
    )

    # ... train current fold ...
    # results[f'fold-{fold}'] = train_and_eval(train_loader, val_loader)
```

## Scenario C: Build DataLoaders from in-memory (X, y)

When data is already in memory (e.g. loaded from cache, obtained from another library), skip the Pipeline layer.

```python
import numpy as np
from zyplib.dataset.torch_data import (
    make_basic_dataloader,
    make_split_dataloader,
    make_kfold_dataloader,
    analyze_dataloader,
)

# Assume data already exists
X = np.random.randn(1000, 8, 250)  # [N, C, T]
y = np.random.randint(0, 2, size=(1000,))

# --- Option 1: Manual split then create ---
from zyplib.dataset import random_split
X_train, X_test = random_split(X, [0.8, 0.2], seed=42)
y_train, y_test = random_split(y, [0.8, 0.2], seed=42)

train_loader, test_loader = make_basic_dataloader(
    (X_train, y_train),
    (X_test, y_test),
    batch_size=64,
)

# --- Option 2: Auto split ---
train_loader, val_loader = make_split_dataloader(
    (X, y),
    split_size=[0.8, 0.2],
    batch_size=64,
    shuffle='shuffle-first',
)

# --- Option 3: Segment-level K-fold ---
# Note: this is segment-level splitting, NOT file-level!
for train_loader, val_loader in make_kfold_dataloader(
    xy=(X, y), k=5, batch_size=64, shuffle='shuffle-train',
):
    analyze_dataloader(train_loader, prefix_name='Train')
    # ... train ...
```

## Scenario D: Using fileloader_pipeline (flexible loading)

When each file has complex loading logic that doesn't fit the fileseg_pipeline read+segment pattern.

```python
from zyplib.dataset import simple_pipeline, train_test_pipeline
from zyplib.dataset.torch_data import make_basic_dataloader
import numpy as np

def my_custom_loader(file_path: str) -> tuple[np.ndarray, np.ndarray]:
    """Each file returns (X, y), X: [N, C, T], y: [N,]"""
    import h5py
    with h5py.File(file_path, 'r') as f:
        X = f['segments'][:]      # pre-segmented data
        y = f['labels'][:]
    return X, y

# Load all files directly
X, y = simple_pipeline(
    files=glob('data/*.h5'),
    fn_fileloader=my_custom_loader,
    fn_encode_label=lambda labels: np.asarray(labels),  # skip default onehot, keep original labels
    random_shuffle=42,
)

# Or combine with file-level splitting
X_train, X_test, y_train, y_test = train_test_pipeline(
    split_size=0.8,
    all_files=glob('data/*.h5'),
    random_seed=42,
    fn_pipeline='simple',
    pipeline_kwargs=dict(
        fn_fileloader=my_custom_loader,
        fn_encode_label=None,
    ),
)
```

## Scenario E: Using BasicXyDataset with custom transform

```python
from zyplib.dataset.torch_data import BasicXyDataset
from torch.utils.data import DataLoader
import torch

def my_augment(x, y):
    """Data augmentation: add Gaussian noise"""
    noise = torch.randn_like(x) * 0.01
    return x + noise, y

dataset = BasicXyDataset(
    X_train, y_train,
    transform=my_augment,
    p_trans=0.5,  # 50% probability to apply augmentation
)

train_loader = DataLoader(dataset, batch_size=64, shuffle=True)
```
