# SPDX-FileCopyrightText: 2026 Adaptive Machine Patterns Ltd
# SPDX-License-Identifier: CC-BY-NC-4.0
"""AMP Tadpole: a small vision-language-action model for the AMP Toy
Robotics Benchmark.

    in    images  uint8 (3, 3, 512, 512): the top, side and iso cameras
          state   (7,): 6 joint angles (rad) + gripper aperture (0 open, 1 closed)
          clause  one instruction clause, in canonical wording
    out   (8, 7): the next 8 joint states + gripper aperture

    images -> a CNN shared by the cameras -> 16 x 16 tokens per camera
    clause -> token ids (instructions.json) -> word embeddings
    state  -> one token
    all tokens + 8 learned action tokens -> 6-layer transformer
    each action token -> one step's change; the changes are added up

This file is a model package (see the benchmark's README): build() makes
the module, predict() runs it on one observation.  It needs only PyTorch.
"""

from __future__ import annotations

import json
import math
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn

HERE = Path(__file__).absolute().parent      # not resolve(): in the Hugging Face cache the files are links to blobs
GRID = 16             # tokens per camera side
CHUNK = 8             # steps per answer
STATE_DIM = 7
TOKEN_LEN = 48        # instruction slots
STATE_FEATURES = 31
INC_SCALE = [0.04] * 6 + [0.125]          # one step's change, per output unit
JOINT_LIMITS = [(0.1, 0.8 * math.pi), (math.pi / 2, 3 * math.pi / 2),
                (-0.7 * math.pi, 0.7 * math.pi), (-0.6 * math.pi, 0.6 * math.pi),
                (-math.pi, math.pi), (-math.pi, math.pi)]


class Tadpole(nn.Module):
    def __init__(self, vocab: int, d: int = 256, layers: int = 6, heads: int = 8):
        super().__init__()

        def block(i, o):
            return nn.Sequential(
                nn.Conv2d(i, o, 3, 2, 1), nn.GroupNorm(8, o), nn.GELU(),
                nn.Conv2d(o, o, 3, 1, 1), nn.GroupNorm(8, o), nn.GELU())
        self.cnn = nn.ModuleList([                                    # 512 -> 16
            nn.Sequential(nn.Conv2d(5, 48, 7, 4, 3), nn.GroupNorm(8, 48), nn.GELU()),
            block(48, 96), block(96, 192), block(192, d)])
        self.pos = nn.Parameter(torch.zeros(1, GRID * GRID, d))
        self.view = nn.Parameter(torch.zeros(3, 1, d))
        self.word = nn.Embedding(vocab, d)
        self.word_pos = nn.Parameter(torch.zeros(1, TOKEN_LEN, d))
        self.state = nn.Linear(STATE_FEATURES, d)
        self.query = nn.Parameter(torch.zeros(1, CHUNK, d))
        layer = nn.TransformerEncoderLayer(d, heads, 4 * d, dropout=0.0,
                                           batch_first=True, norm_first=True)
        self.fuse = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False)
        self.head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, STATE_DIM))
        self.register_buffer("lo", torch.tensor([a for a, _ in JOINT_LIMITS] + [0.0]), persistent=False)
        self.register_buffer("hi", torch.tensor([b for _, b in JOINT_LIMITS] + [1.0]), persistent=False)
        self.register_buffer("inc", torch.tensor(INC_SCALE), persistent=False)

    def state_features(self, state):
        """The 7 numbers scaled to -1..1, plus sines and cosines of the joint
        angles, so that small differences in pose are not lost."""
        q = state[:, :6]
        return torch.cat([2 * (state - self.lo) / (self.hi - self.lo) - 1,
                          torch.sin(2 * q), torch.cos(2 * q),
                          torch.sin(8 * q), torch.cos(8 * q)], dim=1)

    def forward(self, images, tokens, token_mask, state):
        """images uint8 (B, 3, 3, H, W); tokens (B, 48) ids from
        instructions.json; token_mask (B, 48) bool, True for real tokens;
        state (B, 7).  Returns (B, 8, 7): the next 8 states."""
        B, H = images.shape[0], images.shape[-1]
        x = images.float() / 255.0
        w = self.word(tokens) + self.word_pos
        ax = torch.linspace(-1, 1, H, device=images.device)
        yy, xx = torch.meshgrid(ax, ax, indexing="ij")
        f = torch.cat([x.flatten(0, 1) * 2 - 1,
                       xx.expand(B * 3, 1, H, H), yy.expand(B * 3, 1, H, H)], dim=1)
        for stage in self.cnn:
            f = stage(f)
        f = f.flatten(2).transpose(1, 2) + self.pos
        f = (f.view(B, 3, GRID * GRID, -1) + self.view).flatten(1, 2)
        s = self.state(self.state_features(state))[:, None]
        q = self.query.expand(B, -1, -1)
        h = torch.cat([f, w, s, q], dim=1)
        keep = torch.ones(B, f.shape[1], dtype=torch.bool, device=h.device)
        tail = torch.ones(B, 1 + CHUNK, dtype=torch.bool, device=h.device)
        h = self.fuse(h, src_key_padding_mask=~torch.cat([keep, token_mask, tail], dim=1))
        steps = self.head(h[:, -CHUNK:]) * self.inc
        return steps.cumsum(dim=1) + state[:, None]


_TABLE = None


def instructions() -> dict:
    """The token table: every canonical clause -> its token ids."""
    global _TABLE
    if _TABLE is None:
        _TABLE = json.loads((HERE / "instructions.json").read_text())
    return _TABLE


def encode(clause: str):
    """(ids, mask) for one canonical clause, each of length 48."""
    t = instructions()
    try:
        ids = t["instructions"][clause]
    except KeyError:
        raise ValueError(f"AMP Tadpole only reads canonical clauses "
                         f"(e.g. 'Pick up the lime star'); got {clause!r}") from None
    pad = t["pad"]
    return ids + [pad] * (TOKEN_LEN - len(ids)), [True] * len(ids) + [False] * (TOKEN_LEN - len(ids))


def build() -> Tadpole:
    return Tadpole(vocab=len(instructions()["vocab"]))


@torch.no_grad()
def predict(module: Tadpole, images, state, clause: str) -> np.ndarray:
    dev = next(module.parameters()).device
    ids, mask = encode(clause)
    out = module(torch.as_tensor(np.asarray(images)[None], device=dev),
                 torch.tensor([ids], device=dev), torch.tensor([mask], device=dev),
                 torch.as_tensor(np.asarray(state, dtype=np.float32)[None], device=dev))
    return out[0].cpu().numpy().astype(np.float64)