"""encoder-things-eeg2: predict the EEG response of THINGS-EEG2 participants to an image.

This file is the whole model in plain PyTorch. ``modeling_encoder_things_eeg2.py`` wraps it for
``transformers.AutoModel``; without ``transformers``, call ``load`` below. It needs ``torch``,
``pillow`` and ``safetensors``, and ``huggingface-hub`` to download files.

Each variant combines one or two vision networks ("members"). A member runs an image to its feature
layer and maps the features to a fixed-size vector with a linear transform fitted on the THINGS-EEG2
training images. The members' vectors are concatenated, and a linear readout per participant maps
them to 63 channels x 80 time points: the response 0-0.79 s after image onset at 100 Hz, in
microvolts, on the recording's reference. By default the output averages the participants'
readouts, which is the prediction for a new participant.
"""

from __future__ import annotations

import json
from collections.abc import Sequence
from pathlib import Path

import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn

RESNET_LAYOUTS = {
    "resnet18": ("basic", (2, 2, 2, 2)),
    "resnet50": ("bottleneck", (3, 4, 6, 3)),
    "resnet152": ("bottleneck", (3, 8, 36, 3)),
}
STAGES = ("layer1", "layer2", "layer3", "layer4")


# Preprocessing ------------------------------------------------------------------------------------


def _to_tensor(image: Image.Image) -> torch.Tensor:
    """uint8 ``(3, height, width)`` from an RGB image."""

    width, height = image.size
    data = torch.frombuffer(bytearray(image.tobytes()), dtype=torch.uint8)
    return data.view(height, width, 3).permute(2, 0, 1)


def resize_center_crop(images: Sequence[Image.Image], resize: int, crop: int) -> torch.Tensor:
    """uint8 ``(images, 3, crop, crop)``: the shorter side resized to ``resize`` (Lanczos), then
    a center crop, as torchvision does for PIL images."""

    batch = []
    for image in images:
        image = image.convert("RGB")
        width, height = image.size
        if width <= height:
            size = (resize, int(resize * height / width))
        else:
            size = (int(resize * width / height), resize)
        image = image.resize(size, Image.Resampling.LANCZOS)
        left = int(round((size[0] - crop) / 2.0))
        top = int(round((size[1] - crop) / 2.0))
        batch.append(_to_tensor(image.crop((left, top, left + crop, top + crop))))
    return torch.stack(batch)


def normalize(pixels: torch.Tensor, mean: Sequence[float], std: Sequence[float]) -> torch.Tensor:
    mean = torch.tensor(mean).view(-1, 1, 1)
    std = torch.tensor(std).view(-1, 1, 1)
    return (pixels.float() / 255.0 - mean) / std


def qwen_patches(images: Sequence[Image.Image], preprocessing: dict, qwen: dict) -> torch.Tensor:
    """Qwen input ``(images, tokens, values)``.

    The image is center-cropped to a square and resized (Lanczos) to the size of the THINGS-EEG2
    image files, then processed as Qwen's image processor does: bicubic resize to the fixed token
    grid, normalization, and patching in 2 x 2 merge order with the frame repeated over the
    temporal patch.
    """

    side = preprocessing["stimulus_size"]
    patch, temporal, merge = qwen["patch_size"], qwen["temporal_patch_size"], qwen["merge_size"]
    _, grid_h, grid_w = qwen["grid"]
    batch = []
    for image in images:
        image = image.convert("RGB")
        short = min(image.size)
        left, top = (image.size[0] - short) // 2, (image.size[1] - short) // 2
        image = image.crop((left, top, left + short, top + short))
        image = image.resize((side, side), Image.Resampling.LANCZOS)
        # Resizing the uint8 image directly reproduces the torchvision resize of Qwen's processor.
        pixels = F.interpolate(
            _to_tensor(image)[None],
            size=(grid_h * patch, grid_w * patch),
            mode="bicubic",
            align_corners=False,
            antialias=True,
        )
        batch.append(normalize(pixels[0], preprocessing["mean"], preprocessing["std"]))
    pixels = torch.stack(batch)[:, None].repeat(1, temporal, 1, 1, 1)  # (images, temporal, 3, H, W)
    shape = (len(images), temporal, 3, grid_h // merge, merge, patch, grid_w // merge, merge, patch)
    patches = pixels.view(shape).permute(0, 3, 6, 4, 7, 2, 1, 5, 8)
    return patches.reshape(len(images), grid_h * grid_w, 3 * temporal * patch * patch)


# ResNet -------------------------------------------------------------------------------------------


class BatchNorm2d(nn.Module):
    """Batch norm for inference: running statistics and affine parameters, no training state."""

    def __init__(self, channels: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(channels))
        self.bias = nn.Parameter(torch.zeros(channels))
        self.register_buffer("running_mean", torch.zeros(channels))
        self.register_buffer("running_var", torch.ones(channels))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.batch_norm(
            x, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps
        )


def _downsample(inplanes: int, outplanes: int, stride: int) -> nn.Module | None:
    if stride == 1 and inplanes == outplanes:
        return None
    return nn.Sequential(
        nn.Conv2d(inplanes, outplanes, 1, stride, bias=False), BatchNorm2d(outplanes)
    )


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes: int, planes: int, stride: int):
        super().__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, 3, stride, 1, bias=False)
        self.bn1 = BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, 3, 1, 1, bias=False)
        self.bn2 = BatchNorm2d(planes)
        self.downsample = _downsample(inplanes, planes, stride)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        identity = x if self.downsample is None else self.downsample(x)
        out = F.relu(self.bn1(self.conv1(x)))
        return F.relu(self.bn2(self.conv2(out)) + identity)


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes: int, planes: int, stride: int):
        super().__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
        self.bn1 = BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, 3, stride, 1, bias=False)
        self.bn2 = BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, 1, bias=False)
        self.bn3 = BatchNorm2d(planes * 4)
        self.downsample = _downsample(inplanes, planes * 4, stride)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        identity = x if self.downsample is None else self.downsample(x)
        out = F.relu(self.bn1(self.conv1(x)))
        out = F.relu(self.bn2(self.conv2(out)))
        return F.relu(self.bn3(self.conv3(out)) + identity)


class ResNet(nn.Module):
    """torchvision's ResNet, cut after block ``last_block`` of ``stage``. Its feature sets are the
    output map (``full``) and its average pools over 1 x 1, 2 x 2 and 4 x 4 grids (``pool1``,
    ``pool2``, ``pool4``)."""

    def __init__(self, spec: dict):
        super().__init__()
        kind, depths = RESNET_LAYOUTS[spec["arch"]]
        block = BasicBlock if kind == "basic" else Bottleneck
        self.sets = spec["sets"]
        self.conv1 = nn.Conv2d(3, 64, 7, 2, 3, bias=False)
        self.bn1 = BatchNorm2d(64)
        self.stages = list(STAGES[: STAGES.index(spec["stage"]) + 1])
        inplanes = 64
        for index, name in enumerate(self.stages):
            planes = 64 * 2**index
            count = spec["last_block"] + 1 if name == spec["stage"] else depths[index]
            blocks = []
            for i in range(count):
                blocks.append(block(inplanes, planes, 2 if i == 0 and index > 0 else 1))
                inplanes = planes * block.expansion
            setattr(self, name, nn.Sequential(*blocks))

    def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
        x = F.max_pool2d(F.relu(self.bn1(self.conv1(x))), 3, 2, 1)
        for name in self.stages:
            x = getattr(self, name)(x)
        sets = {}
        for name in self.sets:
            pooled = (
                x if name == "full" else F.adaptive_avg_pool2d(x, int(name.removeprefix("pool")))
            )
            sets[name] = pooled.flatten(1)
        return sets


# ViT ----------------------------------------------------------------------------------------------


class Attention(nn.Module):
    def __init__(self, dim: int, num_heads: int):
        super().__init__()
        self.num_heads = num_heads
        self.qkv = nn.Linear(dim, dim * 3)
        self.proj = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch, tokens, dim = x.shape
        qkv = self.qkv(x).reshape(batch, tokens, 3, self.num_heads, dim // self.num_heads)
        q, k, v = qkv.permute(2, 0, 3, 1, 4)
        x = F.scaled_dot_product_attention(q, k, v)
        return self.proj(x.transpose(1, 2).reshape(batch, tokens, dim))


class Mlp(nn.Module):
    def __init__(self, dim: int, hidden: int):
        super().__init__()
        self.fc1 = nn.Linear(dim, hidden)
        self.fc2 = nn.Linear(hidden, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc2(F.gelu(self.fc1(x)))


class VitBlock(nn.Module):
    def __init__(self, dim: int, num_heads: int, eps: float):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim, eps=eps)
        self.attn = Attention(dim, num_heads)
        self.norm2 = nn.LayerNorm(dim, eps=eps)
        self.mlp = Mlp(dim, dim * 4)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.norm1(x))
        return x + self.mlp(self.norm2(x))


class PatchEmbed(nn.Module):
    def __init__(self, patch: int, dim: int):
        super().__init__()
        self.proj = nn.Conv2d(3, dim, patch, patch)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.proj(x).flatten(2).transpose(1, 2)


class VisionTransformer(nn.Module):
    """timm's ViT with a class token and learned positions, cut after block ``depth - 1`` and
    before the final norm. Its feature set is the mean over patch tokens (``patch_mean``)."""

    def __init__(self, spec: dict):
        super().__init__()
        dim, patch = spec["embed_dim"], spec["patch_size"]
        tokens = (spec["preprocessing"]["crop"] // patch) ** 2 + 1
        self.patch_embed = PatchEmbed(patch, dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, tokens, dim))
        blocks = [VitBlock(dim, spec["num_heads"], spec["norm_eps"]) for _ in range(spec["depth"])]
        self.blocks = nn.Sequential(*blocks)

    def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
        x = self.patch_embed(x)
        x = torch.cat([self.cls_token.expand(len(x), -1, -1), x], 1) + self.pos_embed
        return {"patch_mean": self.blocks(x)[:, 1:].mean(1)}


# Qwen3-Omni vision encoder ------------------------------------------------------------------------


# The three classes below only hold weights under the names in Qwen's checkpoint; QwenBlock runs
# them.


class QwenPatchEmbed(nn.Module):
    def __init__(self, dim: int, temporal: int, patch: int):
        super().__init__()
        kernel = (temporal, patch, patch)
        self.proj = nn.Conv3d(3, dim, kernel, kernel)


class QwenAttention(nn.Module):
    def __init__(self, dim: int, num_heads: int):
        super().__init__()
        self.num_heads = num_heads
        self.qkv = nn.Linear(dim, dim * 3)
        self.proj = nn.Linear(dim, dim)


class QwenMlp(nn.Module):
    def __init__(self, dim: int, hidden: int):
        super().__init__()
        self.linear_fc1 = nn.Linear(dim, hidden)
        self.linear_fc2 = nn.Linear(hidden, dim)


class LoRA(nn.Module):
    """A low-rank update ``x @ A.T @ B.T`` of one linear layer. It stays separate from the
    bfloat16 layer, in its own (float32) dtype: merging it into bfloat16 weights would round most
    of it away."""

    def __init__(self, rank: int, inputs: int, outputs: int):
        super().__init__()
        self.A = nn.Parameter(torch.zeros(rank, inputs))
        self.B = nn.Parameter(torch.zeros(outputs, rank))

    def forward(self, layer: nn.Linear, x: torch.Tensor) -> torch.Tensor:
        return layer(x) + (x.to(self.A.dtype) @ self.A.T @ self.B.T).to(x.dtype)


def _rotate(t: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Rotary position embedding, computed in float32 as in transformers' Qwen3-Omni encoder."""

    x = t.float()
    half = x.shape[-1] // 2
    return (x * cos + torch.cat([-x[..., half:], x[..., :half]], -1) * sin).to(t.dtype)


class QwenBlock(nn.Module):
    def __init__(self, dim: int, num_heads: int, hidden: int):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
        self.attn = QwenAttention(dim, num_heads)
        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
        self.mlp = QwenMlp(dim, hidden)

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        lora: dict[str, LoRA] | None = None,
    ) -> torch.Tensor:
        """``x`` ``(images, tokens, dim)``, attention within each image; ``lora`` maps layer names
        (``qkv``, ``proj``, ``linear_fc1``, ``linear_fc2``) to their updates."""

        def linear(name: str, layer: nn.Linear, inputs: torch.Tensor) -> torch.Tensor:
            return lora[name](layer, inputs) if lora and name in lora else layer(inputs)

        batch, tokens, dim = x.shape
        heads = self.attn.num_heads
        qkv = linear("qkv", self.attn.qkv, self.norm1(x)).reshape(batch, tokens, 3, heads, -1)
        q, k, v = qkv.unbind(2)
        q, k = _rotate(q, cos, sin), _rotate(k, cos, sin)
        scale = (dim // heads) ** -0.5
        attended = F.scaled_dot_product_attention(
            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), scale=scale
        )
        x = x + linear("proj", self.attn.proj, attended.transpose(1, 2).reshape(batch, tokens, dim))
        hidden = F.gelu(
            linear("linear_fc1", self.mlp.linear_fc1, self.norm2(x)), approximate="tanh"
        )
        return x + linear("linear_fc2", self.mlp.linear_fc2, hidden)


class QwenVision(nn.Module):
    """The first ``depth`` blocks of the Qwen3-Omni vision encoder, for images of one token grid."""

    def __init__(self, qwen: dict):
        super().__init__()
        dim = qwen["hidden_size"]
        self.qwen = qwen
        self.patch_embed = QwenPatchEmbed(dim, qwen["temporal_patch_size"], qwen["patch_size"])
        self.pos_embed = nn.Embedding(qwen["num_position_embeddings"], dim)
        self.blocks = nn.ModuleList(
            QwenBlock(dim, qwen["num_heads"], qwen["intermediate_size"])
            for _ in range(qwen["depth"])
        )

    def set_positions(self) -> None:
        """Position embeddings and rotary tables of the token grid, computed as in transformers'
        Qwen3-Omni encoder (after the encoder is cast to bfloat16). Call after loading weights."""

        dtype, device = self.pos_embed.weight.dtype, self.pos_embed.weight.device
        _, h, w = self.qwen["grid"]
        merge = self.qwen["merge_size"]
        side = int(self.qwen["num_position_embeddings"] ** 0.5)

        # bilinear interpolation of the learned position table onto the grid
        h_idx, w_idx = torch.linspace(0, side - 1, h), torch.linspace(0, side - 1, w)
        h0, w0 = h_idx.int(), w_idx.int()
        h1, w1 = (h0 + 1).clip(max=side - 1), (w0 + 1).clip(max=side - 1)
        dh, dw = h_idx - h0, w_idx - w0
        corners = [
            (h0, w0, (1 - dh)[:, None] * (1 - dw)[None]),
            (h0, w1, (1 - dh)[:, None] * dw[None]),
            (h1, w0, dh[:, None] * (1 - dw)[None]),
            (h1, w1, dh[:, None] * dw[None]),
        ]
        position = None
        for rows, cols, weight in corners:
            index = (rows[:, None] * side + cols[None]).flatten().to(device)
            term = self.pos_embed(index) * weight.flatten().to(device, dtype)[:, None]
            position = term if position is None else position + term
        position = position.view(h // merge, merge, w // merge, merge, -1).permute(0, 2, 1, 3, 4)

        # 2-D rotary angles of each token (row and column), tokens in merge order
        rotary_dim = self.qwen["hidden_size"] // self.qwen["num_heads"] // 2
        steps = torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim
        inv_freq = (1.0 / (10000.0**steps)).to(dtype)
        table = torch.outer(torch.arange(max(h, w), dtype=dtype), inv_freq)
        offsets = torch.arange(merge)
        rows = torch.arange(h // merge)[:, None, None, None] * merge + offsets[None, None, :, None]
        cols = torch.arange(w // merge)[None, :, None, None] * merge + offsets[None, None, None, :]
        rows = rows.expand(h // merge, w // merge, merge, merge).reshape(-1)
        cols = cols.expand(h // merge, w // merge, merge, merge).reshape(-1)
        angles = torch.cat([table[rows], table[cols]], 1)
        angles = torch.cat([angles, angles], 1).to(device)

        self.register_buffer("position", position.reshape(h * w, -1), persistent=False)
        self.register_buffer("cos", angles.cos().float()[:, None], persistent=False)
        self.register_buffer("sin", angles.sin().float()[:, None], persistent=False)

    def embed(self, patches: torch.Tensor) -> torch.Tensor:
        batch, tokens, _ = patches.shape
        proj = self.patch_embed.proj
        x = patches.to(proj.weight.dtype).reshape(-1, 3, *proj.kernel_size)
        return proj(x).reshape(batch, tokens, -1) + self.position


def load_qwen_vision(qwen: dict, device: torch.device) -> QwenVision:
    """The Qwen3-Omni vision blocks a variant uses, read in bfloat16 from Qwen's repository at a
    pinned revision. Only the safetensors shards holding those weights are downloaded."""

    from huggingface_hub import hf_hub_download
    from safetensors import safe_open

    repo, revision, prefix = qwen["repo"], qwen["revision"], "thinker.visual."
    index = Path(hf_hub_download(repo, "model.safetensors.index.json", revision=revision))
    weight_map = json.loads(index.read_text())["weight_map"]
    model = QwenVision(qwen)
    wanted = set(model.state_dict())
    shards = {shard for key, shard in weight_map.items() if key.removeprefix(prefix) in wanted}
    state = {}
    for shard in sorted(shards):
        with safe_open(hf_hub_download(repo, shard, revision=revision), "pt") as handle:
            for key in handle.keys():
                if key.startswith(prefix) and key.removeprefix(prefix) in wanted:
                    state[key.removeprefix(prefix)] = handle.get_tensor(key)
    missing = wanted - set(state)
    if missing:
        raise RuntimeError(f"{repo}@{revision} lacks vision weights {sorted(missing)[:5]}")
    model.load_state_dict(state)
    model = model.to(device=device, dtype=torch.bfloat16)
    model.set_positions()
    return model.eval().requires_grad_(False)


# Members ------------------------------------------------------------------------------------------


class Member(nn.Module):
    """One vision network (a copy per fine-tuning seed) and the linear map from its feature sets
    to the member's part of the readout input.

    PCA-compressed members: ``sum over (seed, set) of features @ weight[seed, set] + bias``.
    Others: ``(features - center) * scale`` for each (seed, set), concatenated.
    """

    def __init__(self, spec: dict):
        super().__init__()
        self.spec = spec
        self.dims = spec["dims"]
        kind = spec["kind"]
        if kind == "resnet":
            self.networks = nn.ModuleDict({seed: ResNet(spec) for seed in spec["seeds"]})
        elif kind == "vit":
            self.networks = nn.ModuleDict({seed: VisionTransformer(spec) for seed in spec["seeds"]})
        elif kind == "qwen_vision":
            self.lora = nn.ModuleDict({seed: self._qwen_lora(spec) for seed in spec["seeds"]})
            # The frozen Qwen blocks are not ours to redistribute: they load from Qwen's repository
            # on first use and are kept out of this module's state dict.
            self._qwen: list[QwenVision] = []
        else:
            raise ValueError(f"unknown member kind {kind!r}")
        pca = spec["pca_components"] is not None
        self.bias = nn.Parameter(torch.zeros(self.dims)) if pca else None
        self.weight = nn.ParameterDict()
        self.center = nn.ParameterDict()
        self.scale = nn.ParameterDict()
        for seed in spec["seeds"]:
            for name in spec["sets"]:
                key, size = f"{seed}__{name}", spec["feature_dims"][name]
                if pca:
                    self.weight[key] = nn.Parameter(torch.zeros(size, self.dims))
                else:
                    self.center[key] = nn.Parameter(torch.zeros(size))
                    self.scale[key] = nn.Parameter(torch.zeros(size))

    @staticmethod
    def _qwen_lora(spec: dict) -> nn.ModuleDict:
        qwen, rank = spec["qwen"], spec["lora_rank"]
        dim, hidden = qwen["hidden_size"], qwen["intermediate_size"]
        shapes = {
            "qkv": (dim, 3 * dim),
            "proj": (dim, dim),
            "linear_fc1": (dim, hidden),
            "linear_fc2": (hidden, dim),
        }
        return nn.ModuleDict(
            {
                f"{block}_{name}": LoRA(rank, *shape)
                for block in range(spec["first_block"], qwen["depth"])
                for name, shape in shapes.items()
            }
        )

    def _qwen_feature_sets(self, patches: torch.Tensor) -> dict[str, dict[str, torch.Tensor]]:
        if not self._qwen:
            self._qwen.append(load_qwen_vision(self.spec["qwen"], patches.device))
        qwen = self._qwen[0] = self._qwen[0].to(patches.device)
        first, (start, last) = self.spec["first_block"], self.spec["window"]
        shared = qwen.embed(patches)
        for index in range(first):
            shared = qwen.blocks[index](shared, qwen.cos, qwen.sin)
        out = {}
        for seed, lora in self.lora.items():
            x, sets = shared, {}
            for index in range(first, last + 1):
                block_lora = {
                    name: lora[f"{index}_{name}"]
                    for name in ("qkv", "proj", "linear_fc1", "linear_fc2")
                }
                x = qwen.blocks[index](x, qwen.cos, qwen.sin, block_lora)
                if index >= start:
                    sets[f"block{index}"] = x.float().mean(1)
            out[seed] = sets
        return out

    def forward(self, pixels: torch.Tensor) -> torch.Tensor:
        if self.spec["kind"] == "qwen_vision":
            feature_sets = self._qwen_feature_sets(pixels)
        else:
            feature_sets = {seed: network(pixels) for seed, network in self.networks.items()}
        total, parts = None, []
        for seed, sets in feature_sets.items():
            for name in self.spec["sets"]:
                key = f"{seed}__{name}"
                if self.bias is not None:
                    weight = self.weight[key]
                    contribution = sets[name].to(weight.dtype) @ weight
                    total = contribution if total is None else total + contribution
                else:
                    center, scale = self.center[key], self.scale[key]
                    parts.append((sets[name].to(scale.dtype) - center) * scale)
        return total + self.bias if self.bias is not None else torch.cat(parts, 1)


class EncoderThingsEEG2(nn.Module):
    """Images -> predicted THINGS-EEG2 EEG ``(images, 63 channels, 80 time points)`` in microvolts.

    ``config`` is one variant's config (``variant_config`` of the combined ``config.json``): its
    ``members``, ``subjects``, ``channels`` and ``time_points``.
    """

    def __init__(self, config: dict):
        super().__init__()
        self.variant = config.get("variant", "")
        self.subjects = list(config["subjects"])
        self.channels = list(config["channels"])
        self.times = list(config["time_points"])
        self.members = nn.ModuleDict({spec["name"]: Member(spec) for spec in config["members"]})
        dims = sum(member.dims for member in self.members.values())
        units = len(self.channels) * len(self.times)
        self.readout_weight = nn.Parameter(torch.zeros(len(self.subjects), dims, units))
        self.readout_bias = nn.Parameter(torch.zeros(len(self.subjects), units))
        self.requires_grad_(False)
        self.eval()

    def inputs(self, images: Sequence[Image.Image]) -> dict[str, torch.Tensor]:
        """Each member's input. Members that resize the same way (ResNet and ViT) share the
        resized crops, so each image is resized once."""

        device, dtype = self.readout_weight.device, self.readout_weight.dtype
        crops, inputs = {}, {}
        for name, member in self.members.items():
            preprocessing = member.spec["preprocessing"]
            if member.spec["kind"] == "qwen_vision":
                pixels = qwen_patches(images, preprocessing, member.spec["qwen"])
            else:
                size = (preprocessing["resize"], preprocessing["crop"])
                if size not in crops:
                    crops[size] = resize_center_crop(images, *size)
                pixels = normalize(crops[size], preprocessing["mean"], preprocessing["std"])
            inputs[name] = pixels.to(device, dtype)
        return inputs

    @torch.no_grad()
    def features(self, images: Sequence[Image.Image]) -> torch.Tensor:
        """The readout input ``(images, dims)``: the members' vectors, concatenated."""

        inputs = self.inputs(images)
        return torch.cat([member(inputs[name]) for name, member in self.members.items()], 1)

    @torch.no_grad()
    def forward(
        self,
        images: Image.Image | Sequence[Image.Image],
        subjects: str | Sequence[str] | None = None,
        batch_size: int = 64,
    ) -> torch.Tensor:
        """Predicted EEG for ``images`` (PIL images).

        ``subjects=None`` returns ``(images, channels, time)``, the mean over the participants'
        readouts. ``subjects="all"`` or a list such as ``["sub-03"]`` returns
        ``(images, len(subjects), channels, time)``.
        """

        if isinstance(images, Image.Image):
            images = [images]
        if subjects is None:
            weight = self.readout_weight.mean(0, keepdim=True)
            bias = self.readout_bias.mean(0, keepdim=True)
        else:
            if subjects == "all":
                names = self.subjects
            elif isinstance(subjects, str):
                names = [subjects]
            else:
                names = list(subjects)
            unknown = sorted(set(names) - set(self.subjects))
            if unknown or not names:
                raise ValueError(f"unknown subjects {unknown or names}; available: {self.subjects}")
            device = self.readout_weight.device
            index = torch.tensor([self.subjects.index(name) for name in names], device=device)
            weight, bias = self.readout_weight[index], self.readout_bias[index]
        outputs = []
        for start in range(0, len(images), batch_size):
            features = self.features(images[start : start + batch_size])
            outputs.append(torch.einsum("id,sdu->isu", features.to(weight.dtype), weight) + bias)
        shape = (len(images), len(weight), len(self.channels), len(self.times))
        out = torch.cat(outputs).reshape(shape).float()  # float32 also under autocast
        return out[:, 0] if subjects is None else out


def variant_config(config: dict, variant: str | None = None) -> dict:
    """One variant's config from the combined ``config.json``: the shared fields plus that
    variant's section (``members``, ``scores``). ``None`` selects ``default_variant``."""

    name = variant or config["default_variant"]
    if name not in config["variants"]:
        raise ValueError(f"unknown variant {name!r}; available: {', '.join(config['variants'])}")
    shared = {k: v for k, v in config.items() if k not in ("variants", "default_variant")}
    return {**shared, **config["variants"][name], "variant": name}


def weight_names(config: dict, variant: str | None = None) -> tuple[str, str]:
    """A variant's weight file and shard index, named as ``transformers``' ``variant`` argument
    names them: ``model.safetensors`` for the default variant, ``model.<variant>.safetensors``
    (shards ``model.<variant>-0000N-of-0000M.safetensors``) for the others."""

    name = variant or config["default_variant"]
    if name == config["default_variant"]:
        return "model.safetensors", "model.safetensors.index.json"
    return f"model.{name}.safetensors", f"model.safetensors.index.{name}.json"


def load(
    repo: str | Path = "epfl-neuroai/encoder-things-eeg2",
    variant: str | None = None,
    revision: str | None = None,
    device: str | torch.device = "cpu",
) -> EncoderThingsEEG2:
    """Build a variant (default: the repository's default variant) from a local copy of the
    repository or from the Hub (pin ``revision``). Only that variant's weights are downloaded."""

    from safetensors.torch import load_file

    folder = Path(repo)
    if folder.is_dir():
        config = json.loads((folder / "config.json").read_text())
        selected = variant_config(config, variant)  # refuses an unknown variant
    else:
        from huggingface_hub import hf_hub_download, snapshot_download

        path = Path(hf_hub_download(str(repo), "config.json", revision=revision))
        config = json.loads(path.read_text())
        selected = variant_config(config, variant)  # refuses an unknown variant before downloading
        single, index = weight_names(config, variant)
        stem = single.removesuffix(".safetensors")
        patterns = [single, index, f"{stem}-*.safetensors"]
        folder = Path(snapshot_download(str(repo), revision=revision, allow_patterns=patterns))
    single, index = weight_names(config, variant)
    if (folder / index).exists():  # large variants are split into shards
        state = {}
        for shard in sorted(set(json.loads((folder / index).read_text())["weight_map"].values())):
            state.update(load_file(folder / shard))
    else:
        state = load_file(folder / single)
    model = EncoderThingsEEG2(selected)
    model.load_state_dict(state, assign=True)  # use the loaded tensors; no second copy in memory
    return model.requires_grad_(False).to(device)


__all__ = ["EncoderThingsEEG2", "load", "variant_config", "weight_names"]