"""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. Students are WIDTH-LIMITED, so depth is the
    resource that matters: a shallow student cannot compose enough nonlinearity
    to match a deep teacher, a deep one can. This is the standard construction
    for a task where accuracy genuinely GROWS with model depth — unlike
    flattened-CIFAR MLPs, which are architecture-bottlenecked and depth-flat."""
    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 with PROVABLE depth-dependence. 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 depth-hard bit). Label = argmax over the C tree roots.
    A depth-d net can only realise ~d tree levels, so accuracy GROWS with depth
    up to `levels` — the regime where scaling the model actually helps, and where
    a rule with poor deep credit assignment (DFA) cannot follow."""
    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.
    A depth-d ReLU net can realise ~d tent compositions, resolving only ~2^d of
    the 2^levels teeth, so TEST ACCURACY PROVABLY RISES WITH DEPTH up to `levels`.
    This is the regime where scaling the model genuinely buys accuracy, and where
    a rule with poor deep credit assignment (DFA/FA) cannot build the composition
    and stalls. Extra input dims are distractors (must be ignored)."""
    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