"""Data loaders. Parse the raw IDX files directly (no torchvision dependency,
so it runs on the Pascal env on the 1080 farm). Datasets already live on the
shared NFS home; the box is offline for dataset mirrors."""

import os
import pickle
import struct
import hashlib
import numpy as np
import torch

DATA_DIR = os.environ.get("SDIL_DATA_DIR", "/home/yurenh2/sdrn/data")

# (subdir, mean, std) matching the standard torchvision normalisation used by
# the predecessor project, so numbers are comparable.
_STATS = {
    "mnist": ("MNIST", 0.1307, 0.3081),
    "fmnist": ("FashionMNIST", 0.2860, 0.3530),
}

_CIFAR_MEAN = np.array([0.4914, 0.4822, 0.4465], dtype=np.float32)
_CIFAR_STD = np.array([0.2470, 0.2435, 0.2616], dtype=np.float32)


def _read_images(path):
    with open(path, "rb") as f:
        magic, n, r, c = struct.unpack(">IIII", f.read(16))
        assert magic == 2051, f"bad image magic {magic}"
        buf = f.read(n * r * c)
    return np.frombuffer(buf, dtype=np.uint8).reshape(n, r * c).astype(np.float32)


def _read_labels(path):
    with open(path, "rb") as f:
        magic, n = struct.unpack(">II", f.read(8))
        assert magic == 2049, f"bad label magic {magic}"
        buf = f.read(n)
    return np.frombuffer(buf, dtype=np.uint8).astype(np.int64)


def _load_split(subdir, train, mean, std, data_dir):
    raw = f"{data_dir}/{subdir}/raw"
    pre = "train" if train else "t10k"
    x = _read_images(f"{raw}/{pre}-images-idx3-ubyte")
    y = _read_labels(f"{raw}/{pre}-labels-idx1-ubyte")
    x = (x / 255.0 - mean) / std                     # normalise + flatten (already 784)
    return torch.from_numpy(x), torch.from_numpy(y)


class _FastLoader:
    """Minimal shuffled minibatch iterator over in-memory (GPU) tensors."""
    def __init__(self, x, y, batch_size, shuffle, seed=0):
        self.x, self.y = x, y
        self.bs = batch_size
        self.shuffle = shuffle
        self.n = x.shape[0]
        self.g = torch.Generator(device="cpu").manual_seed(seed)

    def __iter__(self):
        idx = torch.randperm(self.n, generator=self.g) if self.shuffle else torch.arange(self.n)
        for i in range(0, self.n, self.bs):
            j = idx[i:i + self.bs]
            yield self.x[j], self.y[j]

    def __len__(self):
        return (self.n + self.bs - 1) // self.bs


class _CIFARImageLoader(_FastLoader):
    """Deterministic in-memory CIFAR iterator with standard train augmentation.

    Images are stored normalized in NCHW order.  Because torchvision applies
    zero padding *before* normalization, the normalized padding value is
    channel dependent; filling the padded tensor explicitly preserves that
    convention instead of padding normalized images with their channel mean.
    All shuffle/crop/flip draws use the loader's private CPU generator, so two
    loaders with the same seed produce bitwise identical epoch streams.
    """

    def __init__(self, x, y, batch_size, shuffle, augment=False, seed=0):
        super().__init__(x, y, batch_size, shuffle, seed=seed)
        if x.ndim != 4 or tuple(x.shape[1:]) != (3, 32, 32):
            raise ValueError(
                f"CIFAR image loader expects NCHW (N,3,32,32), got {tuple(x.shape)}")
        if augment and not shuffle:
            raise ValueError("CIFAR augmentation is restricted to the shuffled training loader")
        self.augment = bool(augment)
        self.pad = 4

    def _augment(self, x):
        batch = x.shape[0]
        if batch == 0:
            return x
        # Equivalent to RandomCrop(32, padding=4, fill=0) before Normalize.
        fill = torch.as_tensor(
            -_CIFAR_MEAN / _CIFAR_STD, dtype=x.dtype, device=x.device)
        padded = x.new_empty((batch, 3, 40, 40))
        padded[:] = fill.view(1, 3, 1, 1)
        padded[:, :, self.pad:self.pad + 32, self.pad:self.pad + 32] = x

        top = torch.randint(0, 2 * self.pad + 1, (batch,), generator=self.g)
        left = torch.randint(0, 2 * self.pad + 1, (batch,), generator=self.g)
        flip = torch.rand(batch, generator=self.g) < 0.5
        top = top.to(x.device)
        left = left.to(x.device)
        flip = flip.to(x.device)

        # Advanced indexing extracts a distinct 32x32 crop for every example.
        rows = top[:, None] + torch.arange(32, device=x.device)[None, :]
        cols = left[:, None] + torch.arange(32, device=x.device)[None, :]
        bhwc = padded.permute(0, 2, 3, 1)
        bidx = torch.arange(batch, device=x.device)[:, None, None]
        cropped = bhwc[bidx, rows[:, :, None], cols[:, None, :], :]
        cropped = cropped.permute(0, 3, 1, 2).contiguous()
        return torch.where(flip[:, None, None, None], cropped.flip(-1), cropped)

    def __iter__(self):
        idx = (torch.randperm(self.n, generator=self.g)
               if self.shuffle else torch.arange(self.n))
        for i in range(0, self.n, self.bs):
            j = idx[i:i + self.bs]
            x, y = self.x[j], self.y[j]
            yield self._augment(x) if self.augment else x, y


def _load_cifar(data_dir, n_classes=10):
    root = f"{data_dir}/cifar-10-batches-py"
    def unpickle(fn):
        with open(fn, "rb") as f:
            return pickle.load(f, encoding="bytes")
    xs, ys = [], []
    for i in range(1, 6):
        d = unpickle(f"{root}/data_batch_{i}")
        xs.append(d[b"data"]); ys.append(np.array(d[b"labels"]))
    xtr = np.concatenate(xs).astype(np.float32); ytr = np.concatenate(ys).astype(np.int64)
    dte = unpickle(f"{root}/test_batch")
    xte = np.array(dte[b"data"], dtype=np.float32); yte = np.array(dte[b"labels"], dtype=np.int64)

    def norm(x):  # x: (N, 3072) as R(1024)G(1024)B(1024)
        x = x.reshape(-1, 3, 1024) / 255.0
        x = (x - _CIFAR_MEAN[None, :, None]) / _CIFAR_STD[None, :, None]
        return x.reshape(-1, 3072).astype(np.float32)
    return (torch.from_numpy(norm(xtr)), torch.from_numpy(ytr),
            torch.from_numpy(norm(xte)), torch.from_numpy(yte))


def get_cifar_image_splits(batch_size=128, data_dir=DATA_DIR, device="cpu",
                           train_limit=None, val_examples=5000, split_seed=2027,
                           loader_seed=0, augment_train=True):
    """Load CIFAR-10 as NCHW images for convolutional protocols.

    Returns ``(train, validation, test, input_shape, n_out, metadata)``.  The
    validation split is stratified and constructed only from the original
    training set.  Training augmentation is the canonical random 32x32 crop
    after four-pixel zero padding followed by a 50% horizontal flip; validation
    and test images are never augmented.  This is intentionally a separate API
    so all frozen flattened-CIFAR experiments retain their historical tensors.
    """
    xtr, ytr, xte, yte = _load_cifar(data_dir)
    xtr = xtr.view(-1, 3, 32, 32)
    xte = xte.view(-1, 3, 32, 32)

    if val_examples:
        train_idx, val_idx = _stratified_validation_indices(
            ytr, val_examples, split_seed)
        xval, yval = xtr[val_idx], ytr[val_idx]
        xtr, ytr = xtr[train_idx], ytr[train_idx]
        digest = hashlib.sha256(val_idx.numpy().tobytes()).hexdigest()
        val_counts = {str(int(cls)): int((yval == cls).sum())
                      for cls in torch.unique(yval, sorted=True)}
    else:
        xval = yval = None
        digest = None
        val_counts = {}
    if train_limit is not None:
        if train_limit <= 0 or train_limit > xtr.shape[0]:
            raise ValueError(
                f"train_limit must be in [1, {xtr.shape[0]}], got {train_limit}")
        xtr, ytr = xtr[:train_limit], ytr[:train_limit]

    xtr, ytr = xtr.to(device), ytr.to(device)
    xte, yte = xte.to(device), yte.to(device)
    validation = None
    if xval is not None:
        xval, yval = xval.to(device), yval.to(device)
        validation = _CIFARImageLoader(
            xval, yval, 1000, False, augment=False, seed=loader_seed)
    metadata = {
        "dataset": "cifar10",
        "input_layout": "NCHW",
        "input_shape": [3, 32, 32],
        "normalization_mean": _CIFAR_MEAN.tolist(),
        "normalization_std": _CIFAR_STD.tolist(),
        "training_augmentation": (
            "random_crop_32_padding_4_zero_then_horizontal_flip_p0.5"
            if augment_train else "none"),
        "loader_seed": int(loader_seed),
        "split_seed": int(split_seed) if val_examples else None,
        "validation_examples": int(val_examples),
        "validation_index_sha256": digest,
        "validation_class_counts": val_counts,
        "train_examples": int(xtr.shape[0]),
        "test_examples": int(xte.shape[0]),
        "split_from_training_only": True,
    }
    train = _CIFARImageLoader(
        xtr, ytr, batch_size, True, augment=augment_train, seed=loader_seed)
    test = _CIFARImageLoader(
        xte, yte, 1000, False, augment=False, seed=loader_seed)
    return train, validation, test, (3, 32, 32), 10, metadata


def _stratified_validation_indices(y, n_val, seed):
    """Return deterministic disjoint train/validation indices.

    Validation examples are allocated as evenly as possible across observed
    classes. The returned validation indices are sorted before hashing so the
    split identity is independent of loader iteration order.
    """
    if n_val <= 0 or n_val >= y.numel():
        raise ValueError(f"n_val must be in [1, {y.numel() - 1}], got {n_val}")
    classes = torch.unique(y, sorted=True)
    base, remainder = divmod(n_val, classes.numel())
    generator = torch.Generator(device="cpu").manual_seed(seed)
    val_parts = []
    for position, cls in enumerate(classes):
        candidates = torch.nonzero(y == cls, as_tuple=False).flatten()
        take = base + int(position < remainder)
        if take > candidates.numel():
            raise ValueError(
                f"class {int(cls)} has only {candidates.numel()} examples, requested {take}")
        order = torch.randperm(candidates.numel(), generator=generator)
        val_parts.append(candidates[order[:take]])
    val_idx = torch.cat(val_parts).sort().values
    is_val = torch.zeros(y.numel(), dtype=torch.bool)
    is_val[val_idx] = True
    train_idx = torch.nonzero(~is_val, as_tuple=False).flatten()
    return train_idx, val_idx


def split_training_loader(loader, val_examples, split_seed, train_batch_size):
    """Split an in-memory synthetic training loader without touching its test set."""
    train_idx, val_idx = _stratified_validation_indices(
        loader.y.detach().cpu(), val_examples, split_seed)
    train_device_idx = train_idx.to(loader.x.device)
    val_device_idx = val_idx.to(loader.x.device)
    train = _FastLoader(loader.x[train_device_idx], loader.y[train_device_idx],
                        train_batch_size, True)
    validation = _FastLoader(loader.x[val_device_idx], loader.y[val_device_idx],
                             1000, False)
    digest = hashlib.sha256(val_idx.numpy().tobytes()).hexdigest()
    counts = {str(int(cls)): int((loader.y.detach().cpu()[val_idx] == cls).sum())
              for cls in torch.unique(loader.y.detach().cpu()[val_idx], sorted=True)}
    return train, validation, {
        "split_seed": split_seed,
        "validation_examples": int(val_examples),
        "validation_index_sha256": digest,
        "validation_class_counts": counts,
        "train_examples": int(train_idx.numel()),
        "split_from_training_only": True,
    }


def get_dataset_splits(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu",
                       shuffle_train=True, train_limit=None, val_examples=0, split_seed=2027):
    """Load real data with an optional training-only validation split.

    Returns ``(train, validation_or_none, test, n_in, n_out, metadata)``. Test
    examples never enter split construction. ``get_dataset`` below preserves
    the historical four-item API for all pre-freeze experiments.
    """
    if name == "cifar10":
        xtr, ytr, xte, yte = _load_cifar(data_dir)
        n_in = 3072
    else:
        subdir, mean, std = _STATS[name]
        xtr, ytr = _load_split(subdir, True, mean, std, data_dir)
        xte, yte = _load_split(subdir, False, mean, std, data_dir)
        n_in = 784

    val_loader = None
    if val_examples:
        train_idx, val_idx = _stratified_validation_indices(ytr, val_examples, split_seed)
        xval, yval = xtr[val_idx], ytr[val_idx]
        xtr, ytr = xtr[train_idx], ytr[train_idx]
        digest = hashlib.sha256(val_idx.numpy().tobytes()).hexdigest()
        val_counts = {str(int(cls)): int((yval == cls).sum())
                      for cls in torch.unique(yval, sorted=True)}
    else:
        val_idx = None
        xval = yval = None
        digest = None
        val_counts = {}
    if train_limit is not None:
        if train_limit <= 0 or train_limit > xtr.shape[0]:
            raise ValueError(f"train_limit must be in [1, {xtr.shape[0]}], got {train_limit}")
        xtr, ytr = xtr[:train_limit], ytr[:train_limit]
    xtr, ytr = xtr.to(device), ytr.to(device)
    xte, yte = xte.to(device), yte.to(device)
    if xval is not None:
        xval, yval = xval.to(device), yval.to(device)
        val_loader = _FastLoader(xval, yval, 1000, False)
    metadata = {
        "dataset": name,
        "split_seed": split_seed if val_examples else None,
        "validation_examples": int(val_examples),
        "validation_index_sha256": digest,
        "validation_class_counts": val_counts,
        "train_examples": int(xtr.shape[0]),
        "test_examples": int(xte.shape[0]),
        "split_from_training_only": True,
    }
    return (_FastLoader(xtr, ytr, batch_size, shuffle_train), val_loader,
            _FastLoader(xte, yte, 1000, False), n_in, 10, metadata)


def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", preload=True,
                shuffle_train=True, train_limit=None):
    train, _, test, n_in, n_out, _ = get_dataset_splits(
        name=name, batch_size=batch_size, data_dir=data_dir, device=device,
        shuffle_train=shuffle_train, train_limit=train_limit)
    return train, test, n_in, n_out


def make_teacher_student(n_in=128, n_classes=10, t_depth=8, t_width=64,
                         n_train=50000, n_test=10000, seed=0, residual=True,
                         t_scale=1.5, batch_size=128, device="cpu"):
    """Deep teacher-student classification. Labels = argmax of a fixed random
    deep (residual) teacher net. This is a candidate controlled composition
    task, not a guarantee that the selected student width, optimizer, and hard
    labels make added depth useful. Every configuration must pass the BP
    useful-depth screen before it is used for a credit-assignment claim."""
    from .core import SDILNet
    tsizes = [n_in] + [t_width] * t_depth + [n_classes]
    teacher = SDILNet(tsizes, act="tanh", device=device, seed=seed + 999,
                      w_scale=t_scale, residual=residual)
    g = torch.Generator(device="cpu").manual_seed(seed)
    xtr = torch.randn(n_train, n_in, generator=g).to(device)
    xte = torch.randn(n_test, n_in, generator=g).to(device)
    with torch.no_grad():
        ytr = teacher.logits(xtr).argmax(1)
        yte = teacher.logits(xte).argmax(1)
    return (_FastLoader(xtr, ytr, batch_size, True),
            _FastLoader(xte, yte, 1000, False), n_in, n_classes)


def make_hierarchical(levels=6, n_classes=10, n_train=50000, n_test=10000,
                      seed=0, batch_size=128, device="cpu"):
    """Hierarchical compositional target. Input
    has 2^levels features; C independent binary trees each fold adjacent pairs
    with a fixed random nonlinear combiner  v' = tanh(a·l + b·r + c·l·r)  (the
    l·r product is the composition-heavy term). Label = argmax over the C tree
    roots. The generator is deep, but a dense shallow student can approximate
    it with sufficient width; no depth lower bound is claimed for the student
    family used here. BP must empirically pass a predeclared depth-gain gate."""
    n_in = 2 ** levels
    g = torch.Generator(device="cpu").manual_seed(seed + 31)
    # per-class, per-level, per-node combiner params
    combiners = []
    width = n_in
    for lev in range(levels):
        half = width // 2
        combiners.append(torch.randn(n_classes, half, 3, generator=g))
        width = half

    def teacher(x):                                    # x: (N, n_in)
        N = x.shape[0]
        outs = []
        for cls in range(n_classes):
            v = x
            w = n_in
            for lev in range(levels):
                half = w // 2
                l = v[:, 0:2 * half:2]
                r = v[:, 1:2 * half:2]
                p = combiners[lev][cls].to(x.device)   # (half, 3)
                v = torch.tanh(l * p[:, 0] + r * p[:, 1] + (l * r) * p[:, 2])
                w = half
            outs.append(v[:, 0])
        return torch.stack(outs, 1)                     # (N, C)

    xtr = torch.randn(n_train, n_in, generator=g).to(device)
    xte = torch.randn(n_test, n_in, generator=g).to(device)
    with torch.no_grad():
        rtr = teacher(xtr)
        rte = teacher(xte)
        # standardise per-class roots on train stats -> roughly balanced argmax
        mu, sd = rtr.mean(0, keepdim=True), rtr.std(0, keepdim=True) + 1e-6
        ytr = ((rtr - mu) / sd).argmax(1)
        yte = ((rte - mu) / sd).argmax(1)
    return (_FastLoader(xtr, ytr, batch_size, True),
            _FastLoader(xte, yte, 1000, False), n_in, n_classes)


def make_tentmap(levels=8, n_in=4, n_train=50000, n_test=10000, seed=0,
                 batch_size=128, device="cpu"):
    """Telgarsky depth-separation task. Label = [tent^levels(x0) > 0.5], where
    tent(v)=1-|2v-1| self-composed `levels` times produces 2^levels oscillations.
    Compositional ReLU constructions have a depth-efficiency separation on this
    target, but ordinary random-initialized SGD is not guaranteed to discover
    them. The experiment is therefore only eligible for a scaling claim after
    BP demonstrates a held-out depth gain. Extra input dims are distractors."""
    g = torch.Generator(device="cpu").manual_seed(seed + 7)
    xtr = torch.rand(n_train, n_in, generator=g).to(device)
    xte = torch.rand(n_test, n_in, generator=g).to(device)

    def label(x):
        v = x[:, 0].clone()
        for _ in range(levels):
            v = 1.0 - (2.0 * v - 1.0).abs()
        return (v > 0.5).long()

    return (_FastLoader(xtr, label(xtr), batch_size, True),
            _FastLoader(xte, label(xte), 1000, False), n_in, 2)


def onehot(y, n_classes, device=None):
    oh = torch.zeros(y.shape[0], n_classes, device=y.device if device is None else device)
    oh.scatter_(1, y.view(-1, 1), 1.0)
    return oh