"""Original SAM prompt-encoder/mask-decoder teaching subsystem.

Accepts cached 2x2 image features (8 channels), normalized coordinates and prompt
labels. No image encoder, mask-input encoder, pretrained weights or image resizing.
Two width-eight two-way blocks with one attention head and three mask tokens.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class PositionEncoding(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("frequencies", torch.randn(2, 4))

    def forward(self, coordinates):
        angle = (2 * coordinates - 1) @ self.frequencies * (2 * math.pi)
        return torch.cat([angle.sin(), angle.cos()], -1)


class PromptEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.position = PositionEncoding()
        self.type_embedding = nn.Embedding(4, 8)

    def forward(self, coordinates, labels):
        # 0 background, 1 foreground, 2 top-left box corner, 3 bottom-right.
        return self.position(coordinates) + self.type_embedding(labels)


class Attention(nn.Module):
    def __init__(self, width=4):
        super().__init__()
        self.width = width
        self.query = nn.Linear(8, width)
        self.key = nn.Linear(8, width)
        self.value = nn.Linear(8, width)
        self.softmax = nn.Softmax(-1)
        self.project = nn.Linear(width, 8)

    def forward(self, q, k, v):
        weights = self.softmax(
            self.query(q) @ self.key(k).transpose(-1, -2) / (self.width ** 0.5)
        )
        return self.project(weights @ self.value(v))


class TwoWayBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.self_attention = Attention(8)
        self.self_norm = nn.LayerNorm(8)
        self.token_to_image = Attention()
        self.token_norm = nn.LayerNorm(8)
        self.up = nn.Linear(8, 32)
        self.relu = nn.ReLU()
        self.down = nn.Linear(32, 8)
        self.mlp_norm = nn.LayerNorm(8)
        self.image_to_token = Attention()
        self.image_norm = nn.LayerNorm(8)

    def forward(self, tokens, image, token_positions, image_positions):
        q = tokens + token_positions
        tokens = self.self_norm(tokens + self.self_attention(q, q, tokens))
        tokens = self.token_norm(
            tokens
            + self.token_to_image(
                tokens + token_positions, image + image_positions, image
            )
        )
        tokens = self.mlp_norm(tokens + self.down(self.relu(self.up(tokens))))
        image = self.image_norm(
            image
            + self.image_to_token(
                image + image_positions, tokens + token_positions, tokens
            )
        )
        return tokens, image


class TwoWayDecoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.first = TwoWayBlock()
        self.second = TwoWayBlock()
        self.final_attention = Attention()
        self.final_norm = nn.LayerNorm(8)

    def forward(self, tokens, image, image_positions):
        positions = tokens
        tokens, image = self.first(tokens, image, positions, image_positions)
        tokens, image = self.second(tokens, image, positions, image_positions)
        tokens = self.final_norm(
            tokens
            + self.final_attention(tokens + positions, image + image_positions, image)
        )
        return tokens, image


class Upscale(nn.Module):
    def __init__(self):
        super().__init__()
        self.first = nn.ConvTranspose2d(8, 4, 2, stride=2)
        self.norm = nn.LayerNorm(4)
        self.second = nn.ConvTranspose2d(4, 2, 2, stride=2)

    def forward(self, image):
        image = self.first(image)
        image = F.gelu(self.norm(image.permute(0, 2, 3, 1)).permute(0, 3, 1, 2))
        return F.gelu(self.second(image))


class Hypernetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.first = nn.Linear(8, 8)
        self.second = nn.Linear(8, 8)
        self.project = nn.Linear(8, 2)

    def forward(self, token):
        return self.project(F.relu(self.second(F.relu(self.first(token)))))


class MaskHeads(nn.Module):
    def __init__(self):
        super().__init__()
        self.mask0 = Hypernetwork()
        self.mask1 = Hypernetwork()
        self.mask2 = Hypernetwork()

    def forward(self, tokens, image):
        weights = torch.stack(
            [
                self.mask0(tokens[:, 0]),
                self.mask1(tokens[:, 1]),
                self.mask2(tokens[:, 2]),
            ],
            1,
        )
        return (weights @ image.flatten(2)).reshape(-1, 3, 8, 8)


class QualityHead(nn.Module):
    def __init__(self):
        super().__init__()
        self.first = nn.Linear(8, 8)
        self.second = nn.Linear(8, 8)
        self.scores = nn.Linear(8, 3)

    def forward(self, token):
        return self.scores(F.relu(self.second(F.relu(self.first(token)))))


class PromptableMaskDecoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.prompt_encoder = PromptEncoder()
        self.output_tokens = nn.Parameter(torch.randn(1, 4, 8) * 0.02)
        self.no_mask = nn.Parameter(torch.randn(1, 1, 8) * 0.02)
        y, x = torch.meshgrid(
            torch.tensor([0.25, 0.75]), torch.tensor([0.25, 0.75]), indexing="ij"
        )
        self.register_buffer("grid", torch.stack([x, y], -1).reshape(1, 4, 2))
        self.decoder = TwoWayDecoder()
        self.upscale = Upscale()
        self.masks = MaskHeads()
        self.quality = QualityHead()

    def forward(self, features, coordinates, labels):
        prompts = self.prompt_encoder(coordinates, labels)
        tokens = torch.cat(
            [self.output_tokens.expand(features.shape[0], -1, -1), prompts], 1
        )
        image = features.flatten(2).transpose(1, 2) + self.no_mask
        positions = self.prompt_encoder.position(self.grid)
        tokens, image = self.decoder(tokens, image, positions)
        upscaled = self.upscale(image.transpose(1, 2).reshape(-1, 8, 2, 2))
        masks = self.masks(tokens[:, 1:4], upscaled)
        scores = self.quality(tokens[:, 0])
        return torch.cat([masks.flatten(1), scores], -1)


def setup():
    torch.manual_seed(67)
    return (
        PromptableMaskDecoder().double(),
        torch.randn(1, 8, 2, 2, dtype=torch.float64),
    )


def verify():
    model, features = setup()
    coords = torch.tensor([[[0.25, 0.25], [0.75, 0.75]]], dtype=torch.float64)
    labels = torch.tensor([[1, 0]])
    original = features.clone()
    base = model(features, coords, labels)
    torch.testing.assert_close(model(features, coords.flip(1), labels.flip(1)), base)
    assert not torch.allclose(model(features, coords, torch.tensor([[0, 0]])), base)
    shifted = coords.clone()
    shifted[:, 0, 0] = 0.6
    assert not torch.allclose(model(features, shifted, labels), base)
    torch.testing.assert_close(features, original, atol=0, rtol=0)
    tokens = torch.randn(1, 3, 8, dtype=torch.float64)
    image = torch.randn(1, 2, 8, 8, dtype=torch.float64)
    weights = torch.stack(
        [
            model.masks.mask0(tokens[:, 0]),
            model.masks.mask1(tokens[:, 1]),
            model.masks.mask2(tokens[:, 2]),
        ],
        1,
    )
    explicit = (weights[:, :, :, None, None] * image[:, None]).sum(2)
    torch.testing.assert_close(model.masks(tokens, image), explicit)
    coords.requires_grad_()
    model(features, coords, labels)[:, :192].square().sum().backward()
    assert coords.grad.norm() > 0
    assert model.decoder.first.image_to_token.query.weight.grad.norm() > 0
    return [
        "Reordering sparse prompts preserves mask and quality outputs",
        "Moving a point or changing foreground/background type changes outputs without modifying cached image features",
        "Hypernetwork masks equal explicit per-pixel feature-weight dot products",
        "Mask gradients reach prompt coordinates and the image-to-token attention path",
    ]


def experiment():
    model, features = setup()
    cases = []
    with torch.no_grad():
        for identity, label, coordinates, types in [
            (
                "points",
                "Foreground + background points",
                [[0.25, 0.25], [0.75, 0.75]],
                [1, 0],
            ),
            ("move", "Move the foreground point", [[0.6, 0.25], [0.75, 0.75]], [1, 0]),
            (
                "type",
                "Mark both points as background",
                [[0.25, 0.25], [0.75, 0.75]],
                [0, 0],
            ),
            ("box", "Two box corners", [[0.2, 0.2], [0.8, 0.8]], [2, 3]),
        ]:
            coords = torch.tensor([coordinates], dtype=torch.float64)
            labels = torch.tensor([types])
            out = model(features, coords, labels)
            masks = out[:, :192].reshape(3, 8, 8)
            cases.append(
                dict(
                    id=identity,
                    label=label,
                    target="prompt_encoder.type_embedding",
                    note="The same cached 2 × 2 image embedding is reused for each prompt. Coordinates are normalized (x,y) in [0,1]. Type IDs are 0 = background, 1 = foreground, 2 = top-left box corner, 3 = bottom-right. Three mask logits express the multi-mask interface; these untrained outputs do not segment a real object and quality scores are not calibrated IoUs.",
                    matrices=[
                        dict(
                            label=f"Candidate {i+1} · mask logits",
                            values=masks[i].tolist(),
                        )
                        for i in range(3)
                    ],
                    vectors=[
                        dict(
                            label="Prompt coordinates · x₀, y₀, x₁, y₁",
                            values=coords.flatten().tolist(),
                        ),
                        dict(label="Prompt type IDs", values=types),
                        dict(
                            label="Predicted quality scores · untrained",
                            values=out[0, 192:].tolist(),
                        ),
                    ],
                    metrics=[
                        dict(label="Image encoder calls per changed prompt", value=0),
                        dict(label="Candidate masks", value=3),
                    ],
                )
            )
    return dict(
        kind="matrices",
        title="Reuse the image; change the prompt.",
        description="Recorded prompt-to-mask computation from cached synthetic image features. Explore point location, prompt type and a box while the two-way decoder updates image and token features.",
        controlLabel="Segmentation prompt",
        cases=cases,
    )