"""encoder-tvsd: predict the multi-unit responses of TVSD electrodes in V1, V4 and IT to an image.

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

A variant is a set of members. A member is a vision network cut at one layer, whose output map is
turned into a feature vector: either the map's multiscale average pools, centered and scaled (or
their PCA scores), or the latents of an attention probe. Each area (V1, V4, IT) uses its own
members; a linear readout per monkey maps the concatenated member features of an area to that
monkey's electrodes in it.
"""

from __future__ import annotations

import json
import math
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 ``square`` pixels, 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["square"]
    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``; returns that block's
    output map."""

    def __init__(self, spec: dict):
        super().__init__()
        kind, depths = RESNET_LAYOUTS[spec["arch"]]
        block = BasicBlock if kind == "basic" else Bottleneck
        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) -> 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)
        return x


# 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 LayerScale(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * self.gamma


class VitBlock(nn.Module):
    def __init__(self, dim: int, num_heads: int, eps: float, layer_scale: bool):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim, eps=eps)
        self.attn = Attention(dim, num_heads)
        self.ls1 = LayerScale(dim) if layer_scale else nn.Identity()
        self.norm2 = nn.LayerNorm(dim, eps=eps)
        self.mlp = Mlp(dim, dim * 4)
        self.ls2 = LayerScale(dim) if layer_scale else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.ls1(self.attn(self.norm1(x)))
        return x + self.ls2(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 (ViT-S/16, DINOv2-B/14), cut after
    block ``last_block`` and before any final norm; returns the patch tokens as a map."""

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.patch_embed(x)
        x = torch.cat([self.cls_token.expand(len(x), -1, -1), x], 1) + self.pos_embed
        tokens = self.blocks(x)[:, 1:]
        return tokens.transpose(1, 2).reshape(len(x), -1, self.side, self.side)


# Qwen3 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 encoders."""

    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 a Qwen3 vision encoder (Qwen3-VL or Qwen3-Omni), 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 vision encoders (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


_QWEN_CACHE: dict[tuple, QwenVision] = {}


def load_qwen_vision(qwen: dict, device: torch.device) -> QwenVision:
    """The Qwen3 vision blocks a member uses, read in bfloat16 from Qwen's repository at a pinned
    revision (only the safetensors files holding them are downloaded) and shared by the members
    that use the same tower."""

    key = (qwen["repo"], qwen["revision"], qwen["depth"], tuple(qwen["grid"]), str(device))
    if key in _QWEN_CACHE:
        return _QWEN_CACHE[key]
    from huggingface_hub import hf_hub_download
    from safetensors import safe_open

    repo, revision, prefix = qwen["repo"], qwen["revision"], qwen["prefix"]
    model = QwenVision(qwen)
    wanted = set(model.state_dict())
    try:
        index = Path(hf_hub_download(repo, "model.safetensors.index.json", revision=revision))
        weight_map = json.loads(index.read_text())["weight_map"]
        files = sorted({f for k, f in weight_map.items() if k.removeprefix(prefix) in wanted})
    except Exception:  # a single-file checkpoint has no index
        files = ["model.safetensors"]
    state = {}
    for name in files:
        with safe_open(hf_hub_download(repo, name, revision=revision), "pt") as handle:
            for key_name in handle.keys():
                if key_name.startswith(prefix) and key_name.removeprefix(prefix) in wanted:
                    state[key_name.removeprefix(prefix)] = handle.get_tensor(key_name)
    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()
    model = model.eval().requires_grad_(False)
    _QWEN_CACHE[key] = model
    return model


# Attention probe ----------------------------------------------------------------------------------


class AttentionProbe(nn.Module):
    """Learned queries attending to a map's tokens (one cross-attention layer); returns the
    normalized latents, flattened."""

    def __init__(self, channels: int, width: int, heads: int, latents: int):
        super().__init__()
        self.width = width
        self.project = nn.Linear(channels, width)
        self.token_norm = nn.LayerNorm(width)
        self.queries = nn.Parameter(torch.zeros(latents, width))
        self.query_norm = nn.LayerNorm(width)
        self.cross = nn.MultiheadAttention(width, heads, batch_first=True)
        self.final_norm = nn.LayerNorm(width)

    def codes(self, tokens: int, device: torch.device) -> torch.Tensor:
        """Sinusoidal position codes of the tokens (in row order), ``(tokens, width)``."""

        position = torch.arange(tokens, dtype=torch.float32, device=device)[:, None]
        steps = torch.arange(0, self.width, 2, dtype=torch.float32, device=device)
        frequency = torch.exp(steps * (-math.log(10000.0) / self.width))
        codes = torch.zeros(tokens, self.width, device=device)
        codes[:, 0::2], codes[:, 1::2] = torch.sin(position * frequency), torch.cos(position * frequency)
        return codes

    def forward(self, maps: torch.Tensor) -> torch.Tensor:
        tokens = maps.flatten(2).transpose(1, 2)
        codes = self.codes(tokens.shape[1], maps.device).to(tokens.dtype)
        x = self.token_norm(self.project(tokens)) + codes
        q = self.queries.expand(len(x), -1, -1)
        attended, _ = self.cross(self.query_norm(q), x, x, need_weights=False)
        return self.final_norm(q + attended).flatten(1)


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


def multiscale_sets(maps: torch.Tensor, names: Sequence[str]) -> dict[str, torch.Tensor]:
    """The named feature sets of a map ``(batch, C, H, W)``: ``full`` (the map) and ``pool<s>``
    (its average pool over an s x s grid), each flattened."""

    sets = {}
    for name in names:
        pooled = maps if name == "full" else F.adaptive_avg_pool2d(maps, int(name.removeprefix("pool")))
        sets[name] = pooled.flatten(1)
    return sets


class Member(nn.Module):
    """One fine-tuned vision network and the map from its output to the member's features."""

    def __init__(self, spec: dict):
        super().__init__()
        self.spec = spec
        kind = spec["kind"]
        if kind == "resnet":
            self.network = ResNet(spec)
        elif kind == "vit":
            self.network = VisionTransformer(spec)
        elif kind == "qwen_vision":
            self.lora = self._qwen_lora(spec)
            # 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}")
        if spec["feature"] == "probe":
            probe = spec["probe"]
            self.probe = AttentionProbe(spec["channels"], probe["width"], probe["heads"], probe["latents"])
        else:
            # Each set is centered and scaled; then either projected on principal components and
            # summed, or (without PCA) concatenated as it is.
            self.center = nn.ParameterDict()
            self.pca = spec.get("pca_components") is not None
            self.components = nn.ParameterDict() if self.pca else None
            self.scale = None if self.pca else nn.ParameterDict()
            for name in spec["sets"]:
                size = spec["feature_dims"][name]
                self.center[name] = nn.Parameter(torch.zeros(size))
                if self.pca:
                    self.components[name] = nn.Parameter(torch.zeros(size, spec["dims"]))
                else:
                    self.scale[name] = nn.Parameter(torch.zeros(()))

    @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"], spec["last_block"] + 1)
                for name, shape in shapes.items()
            }
        )

    def _qwen_map(self, patches: torch.Tensor) -> torch.Tensor:
        spec = self.spec
        qwen = load_qwen_vision(spec["qwen"], patches.device)
        first, last = spec["first_block"], spec["last_block"]
        x = qwen.embed(patches)
        for index in range(last + 1):
            block_lora = None
            if index >= first:
                block_lora = {
                    name: self.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 spec["post_norm"]:
            x = qwen.blocks[last + 1].norm1(x)
        # tokens are in 2 x 2 merge order; lay them out on the patch grid
        _, height, width = spec["qwen"]["grid"]
        batch, _, hidden = x.shape
        x = x.float().reshape(batch, height // 2, width // 2, 2, 2, hidden)
        return x.permute(0, 5, 1, 3, 2, 4).reshape(batch, hidden, height, width)

    def forward(self, pixels: torch.Tensor) -> torch.Tensor:
        if self.spec["kind"] == "qwen_vision":
            maps = self._qwen_map(pixels)
        else:
            maps = self.network(pixels)
        if self.spec["feature"] == "probe":
            return self.probe(maps.to(self.probe.project.weight.dtype))
        total, parts = None, []
        for name, values in multiscale_sets(maps, self.spec["sets"]).items():
            center = self.center[name]
            values = values.to(center.dtype) - center
            if self.pca:
                part = values @ self.components[name]
                total = part if total is None else total + part
            else:
                parts.append(values * self.scale[name])
        return total if self.pca else torch.cat(parts, 1)


class EncoderTVSD(nn.Module):
    """Images -> predicted TVSD multi-unit responses, ``{monkey: {area: (images, electrodes)}}``.

    ``config`` is one variant's config (``variant_config`` of the combined ``config.json``): its
    ``members``, and per area the members it reads and each monkey's electrode count."""

    def __init__(self, config: dict):
        super().__init__()
        self.variant = config.get("variant", "")
        self.monkeys = list(config["monkeys"])
        self.areas = {area: dict(spec) for area, spec in config["areas"].items()}
        self.members = nn.ModuleDict({spec["name"]: Member(spec) for spec in config["members"]})
        self.readout = nn.ModuleDict()
        for area, spec in self.areas.items():
            dims = sum(self.members[name].spec["dims"] for name in spec["members"])
            heads = nn.ParameterDict()
            for monkey in self.monkeys:
                units = spec["electrodes"][monkey]
                heads[f"{monkey}_weight"] = nn.Parameter(torch.zeros(dims, units))
                heads[f"{monkey}_bias"] = nn.Parameter(torch.zeros(units))
            self.readout[area] = heads
        self.requires_grad_(False)
        self.eval()

    def inputs(self, images: Sequence[Image.Image], members: Sequence[str]) -> dict[str, torch.Tensor]:
        """Each member's input; members that resize the same way share the resized crops."""

        device = next(iter(self.readout.parameters())).device
        crops, inputs = {}, {}
        for name in members:
            spec = self.members[name].spec
            preprocessing = spec["preprocessing"]
            if spec["kind"] == "qwen_vision":
                pixels = qwen_patches(images, preprocessing, 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)
        return inputs

    @torch.no_grad()
    def features(self, images: Image.Image | Sequence[Image.Image], area: str, batch_size: int = 32) -> torch.Tensor:
        """The readout input of ``area`` ``(images, dims)``: its members' features, concatenated.
        For other animals or datasets, fit a linear readout on these."""

        if isinstance(images, Image.Image):
            images = [images]
        if area not in self.areas:
            raise ValueError(f"unknown area {area!r}; available: {list(self.areas)}")
        members = self.areas[area]["members"]
        out = []
        for start in range(0, len(images), batch_size):
            inputs = self.inputs(images[start : start + batch_size], members)
            out.append(torch.cat([self.members[m](inputs[m]).float() for m in members], 1))
        return torch.cat(out)

    @torch.no_grad()
    def forward(
        self,
        images: Image.Image | Sequence[Image.Image],
        areas: str | Sequence[str] | None = None,
        batch_size: int = 32,
    ) -> dict[str, dict[str, torch.Tensor]]:
        """Predicted responses for ``images`` (PIL images): ``{monkey: {area: (images,
        electrodes)}}`` for ``areas`` (default: all of ``V1``, ``V4``, ``IT``)."""

        if isinstance(images, Image.Image):
            images = [images]
        names = list(self.areas) if areas is None else ([areas] if isinstance(areas, str) else list(areas))
        unknown = sorted(set(names) - set(self.areas))
        if unknown or not names:
            raise ValueError(f"unknown areas {unknown or names}; available: {list(self.areas)}")
        needed = list(dict.fromkeys(m for area in names for m in self.areas[area]["members"]))
        out = {monkey: {area: [] for area in names} for monkey in self.monkeys}
        for start in range(0, len(images), batch_size):
            batch = images[start : start + batch_size]
            inputs = self.inputs(batch, needed)
            features = {name: self.members[name](inputs[name]) for name in needed}
            for area in names:
                x = torch.cat([features[m] for m in self.areas[area]["members"]], 1)
                heads = self.readout[area]
                for monkey in self.monkeys:
                    weight = heads[f"{monkey}_weight"]
                    out[monkey][area].append((x.to(weight.dtype) @ weight + heads[f"{monkey}_bias"]).float())
        return {m: {a: torch.cat(v) for a, v in areas_out.items()} for m, areas_out in out.items()}


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``, ``areas``, ``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-tvsd",
    variant: str | None = None,
    revision: str | None = None,
    device: str | torch.device = "cpu",
) -> EncoderTVSD:
    """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 = EncoderTVSD(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__ = ["EncoderTVSD", "load", "variant_config", "weight_names"]