# Examples

Complete code skeletons for common training scenarios. Each scenario is self-contained.

## Scenario A: Standard binary classification (procedure-style)

The recommended pattern for most EEG/signal classification tasks.

```python
import torch
import torch.nn as nn
from zyplib.train import (
    default_fabric, FabricTrainingContext, procedure,
    TorchMetricRecorder, clf_metrics_builder, seed_everything,
)
from zyplib.train.common.callbacks import EarlyStopping, LRSchedulerCallback

seed_everything(42)

# 1. Fabric with custom callbacks
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
fabric = default_fabric(
    dir_name='lightning_logs',
    keep_best_k=3,
    monitor='loss',
    fn_custom_callback=lambda v: [
        EarlyStopping(monitor='loss', patience=20),
        LRSchedulerCallback(scheduler),
    ],
)

# 2. Metric recorders
def new_recorder():
    return TorchMetricRecorder(
        clf_metrics_builder(binary=True, use_auc=True, use_specificity=True),
        before_metric='sigmoid',
    )

# 3. Context
ctx = FabricTrainingContext(
    model=model,
    optimizer=optimizer,
    fabric=fabric,
    fn_loss=nn.BCEWithLogitsLoss(),
    train_metric_recorder=new_recorder(),
    val_metric_recorder=new_recorder(),
)
ctx.launch()

# 4. Train
procedure.train_loop(
    ctx, train_loader,
    max_epochs=100,
    val_loader=val_loader,
    hparams={'lr': 1e-3, 'batch_size': 32},
)

# 5. Evaluate
ctx.load_best_ckpt('replace')
scores = procedure.test(ctx, test_loader)
```

## Scenario B: Multiclass classification (procedure-style)

Key differences from binary: `CrossEntropyLoss`, `before_metric='softmax'`, `num_classes`.

```python
import torch.nn as nn
from zyplib.train import (
    default_fabric, FabricTrainingContext, procedure,
    TorchMetricRecorder, clf_metrics_builder,
)

num_classes = 4

fabric = default_fabric(dir_name='lightning_logs', keep_best_k=3)

def new_recorder():
    return TorchMetricRecorder(
        clf_metrics_builder(task='multiclass', num_classes=num_classes),
        before_metric='softmax',
    )

ctx = FabricTrainingContext(
    model=model,
    optimizer=optimizer,
    fabric=fabric,
    fn_loss=nn.CrossEntropyLoss(),
    train_metric_recorder=new_recorder(),
    val_metric_recorder=new_recorder(),
)
ctx.launch()

procedure.train_loop(ctx, train_loader, max_epochs=100, val_loader=val_loader)
ctx.load_best_ckpt('replace')
scores = procedure.test(ctx, test_loader)
```

## Scenario C: Trainer-style (BasicFabricTrainer)

OOP alternative. Same underlying logic, different API surface.

```python
import torch.nn as nn
from zyplib.train import default_fabric, clf_metrics_builder
from zyplib.train.common import BasicFabricTrainer

fabric = default_fabric(dir_name='lightning_logs', keep_best_k=3)

trainer = BasicFabricTrainer(
    fabric, model, optimizer,
    fn_loss=nn.BCEWithLogitsLoss(),
    metric_builder=lambda: clf_metrics_builder(binary=True, use_auc=True),
    before_metric='sigmoid',
    hparams={'lr': 1e-3},
    launch_fabric_internally=True,  # auto-calls ctx.launch()
)

trainer.fit(train_loader, val_loader, max_epochs=100)
trainer.load_best_ckpt()
trainer.test(test_loader)
```

### Trainer with custom step (subclass)

```python
class MyTrainer(BasicFabricTrainer):
    def training_step(self, *, ctx, batch, batch_idx, epoch_idx=None):
        x, y = batch
        ctx.model.train()
        ctx.optimizer.zero_grad()
        y_pred = ctx.model(x)
        loss = self.fn_loss(y_pred, y)
        self.fabric.backward(loss)
        ctx.optimizer.step()
        return y_pred, {'loss': loss.item()}
```

### Trainer with dynamic step (set_training_step)

```python
trainer = BasicFabricTrainer(fabric, model, optimizer, ...)
trainer.set_training_step(my_custom_train_step)
trainer.set_validation_step(my_custom_val_step)
trainer.fit(train_loader, val_loader, max_epochs=100)
```

## Scenario D: Custom step function with multi-loss

When you need to track multiple loss components.

```python
from zyplib.train import (
    default_fabric, FabricTrainingContext, procedure, TorchMetricRecorder,
)

def my_train_step(ctx, batch, *, batch_idx, epoch_idx):
    x, y = batch
    ctx.model.train()
    ctx.optimizer.zero_grad()
    y_pred = ctx.model(x)
    ce_loss = F.cross_entropy(y_pred, y)
    l2_reg = 0.01 * sum(p.pow(2).sum() for p in ctx.model.parameters())
    total = ce_loss + l2_reg
    ctx.fabric.backward(total)
    ctx.optimizer.step()
    return y_pred, {'loss': total.item(), 'ce_loss': ce_loss.item(), 'l2_reg': l2_reg.item()}

def my_val_step(ctx, batch, *, batch_idx, epoch_idx):
    x, y = batch
    ctx.model.eval()
    with torch.no_grad():
        y_pred = ctx.model(x)
        loss = F.cross_entropy(y_pred, y)
    return y_pred, {'loss': loss.item()}

# Use scalar_metrics='auto' to track all returned keys
fabric = default_fabric(dir_name='lightning_logs')
recorder_train = TorchMetricRecorder(
    clf_metrics_builder(task='multiclass', num_classes=4),
    before_metric='softmax',
    scalar_metrics='auto',  # auto-tracks 'loss', 'ce_loss', 'l2_reg'
)
recorder_val = TorchMetricRecorder(
    clf_metrics_builder(task='multiclass', num_classes=4),
    before_metric='softmax',
    scalar_metrics='loss',
)

ctx = FabricTrainingContext(
    model=model, optimizer=optimizer, fabric=fabric,
    train_metric_recorder=recorder_train,
    val_metric_recorder=recorder_val,
)
ctx.launch()

procedure.train_loop(
    ctx, train_loader, max_epochs=100, val_loader=val_loader,
    fn_train_step=my_train_step,
    fn_val_step=my_val_step,
)
```

## Scenario E: K-fold training loop

Combine with `zyplib.dataset.kfold_pipeline` for file-level K-fold.

```python
from zyplib.dataset import kfold_pipeline, make_basic_dataloader
from zyplib.train import (
    default_fabric, FabricTrainingContext, procedure,
    TorchMetricRecorder, clf_metrics_builder, seed_everything,
)

seed_everything(42)
fold_results = []

for fold_idx, (X_train, X_val, y_train, y_val) in enumerate(
    kfold_pipeline(k=5, all_files=(asd_files, td_files), fn_pipeline=my_load_fn, random_seed=42)
):
    train_loader, val_loader = make_basic_dataloader(
        (X_train, y_train), (X_val, y_val), batch_size=32
    )

    model = create_model()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    fabric = default_fabric(
        dir_name='kfold_logs',
        version=f'fold_{fold_idx}',
        keep_best_k=1,
    )

    def new_recorder():
        return TorchMetricRecorder(
            clf_metrics_builder(binary=True), before_metric='sigmoid',
        )

    ctx = FabricTrainingContext(
        model=model, optimizer=optimizer, fabric=fabric,
        fn_loss=nn.BCEWithLogitsLoss(),
        train_metric_recorder=new_recorder(),
        val_metric_recorder=new_recorder(),
    )
    ctx.launch()

    procedure.train_loop(ctx, train_loader, max_epochs=50, val_loader=val_loader)
    ctx.load_best_ckpt('replace')

    # Evaluate this fold (reuse val_loader as test)
    scores = procedure.test(ctx, val_loader)
    fold_results.append(scores)

# Aggregate results across folds
import numpy as np
for key in fold_results[0]:
    values = [r[key] for r in fold_results]
    print(f'{key}: {np.mean(values):.4f} +/- {np.std(values):.4f}')
```

## Scenario F: After-training evaluation

Loading a saved checkpoint and running evaluation.

```python
from zyplib.train import FabricTrainingContext, procedure, TorchMetricRecorder, clf_metrics_builder
from zyplib.nn.utils import load_checkpoint
from lightning.fabric import Fabric
import torch
import torch.nn as nn

# Minimal fabric for test (no loggers needed)
fabric = Fabric()

# Build model/context first
model = create_model()

# Create context — fn_loss is required because default test uses validation_step which needs it
optimizer = torch.optim.Adam(model.parameters())  # dummy, not used for test
ctx = FabricTrainingContext(
    model=model, optimizer=optimizer, fabric=fabric,
    fn_loss=nn.BCEWithLogitsLoss(),  # REQUIRED: test falls back to validation_step which resolves loss
    val_metric_recorder=TorchMetricRecorder(
        clf_metrics_builder(binary=True, use_auc=True),
        before_metric='sigmoid',
    ),
)
ctx.launch()  # recommended: launches Fabric and moves metric recorder to device

# Load checkpoint after launch so the target device is known
load_checkpoint(model, 'path/to/best.ckpt', device=fabric.device)

scores = procedure.test(ctx, test_loader)
print(scores)
```
