"""
kiri/data.py

DataLoader — wraps numpy arrays or file-backed datasets into
iterable mini-batches. Works on any backend.
"""

import numpy as np


class Dataset:
    """
    Base dataset class. Subclass this for custom data.

    Example
    -------
    class MyDataset(kiri.Dataset):
        def __init__(self, X, y):
            self.X, self.y = X, y

        def __len__(self):
            return len(self.X)

        def __getitem__(self, idx):
            return self.X[idx], self.y[idx]
    """
    def __len__(self):
        raise NotImplementedError

    def __getitem__(self, idx):
        raise NotImplementedError


class TensorDataset(Dataset):
    """Wraps numpy arrays into a Dataset."""
    def __init__(self, *arrays):
        lengths = [len(a) for a in arrays]
        assert len(set(lengths)) == 1, "All arrays must have the same length"
        self.arrays = arrays

    def __len__(self):
        return len(self.arrays[0])

    def __getitem__(self, idx):
        return tuple(a[idx] for a in self.arrays)


class DataLoader:
    """
    Iterates a Dataset in mini-batches.

    Parameters
    ----------
    dataset    : Dataset or (X, y) tuple
    batch_size : samples per batch
    shuffle    : shuffle before each epoch
    drop_last  : drop the final incomplete batch

    Usage
    -----
    loader = kiri.DataLoader((X_train, y_train), batch_size=64, shuffle=True)

    for epoch in range(10):
        for X_batch, y_batch in loader:
            ...
    """

    def __init__(self, dataset, batch_size=32, shuffle=True, drop_last=False):
        if isinstance(dataset, (tuple, list)):
            dataset = TensorDataset(*dataset)
        self.dataset    = dataset
        self.batch_size = batch_size
        self.shuffle    = shuffle
        self.drop_last  = drop_last
        self._indices   = np.arange(len(dataset))

    def __len__(self):
        n = len(self.dataset)
        if self.drop_last:
            return n // self.batch_size
        return (n + self.batch_size - 1) // self.batch_size

    def __iter__(self):
        indices = self._indices.copy()
        if self.shuffle:
            np.random.shuffle(indices)

        for start in range(0, len(indices), self.batch_size):
            batch_idx = indices[start:start + self.batch_size]
            if self.drop_last and len(batch_idx) < self.batch_size:
                break

            # collate
            items = [self.dataset[i] for i in batch_idx]
            if isinstance(items[0], tuple):
                yield tuple(
                    np.stack([item[k] for item in items])
                    for k in range(len(items[0]))
                )
            else:
                yield np.stack(items)

    def __repr__(self):
        return (f"DataLoader(dataset_size={len(self.dataset)}, "
                f"batch_size={self.batch_size}, "
                f"shuffle={self.shuffle}, "
                f"num_batches={len(self)})")
