"""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 pickle
import struct
import hashlib
import numpy as np
import torch

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


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 _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