# examples/sam.py
"""Original SAM inference with verified native modules and reusable embeddings.

Download sam_native.zip and extract beside this file. Offline --smoke uses a
reduced, random-weight native architecture, not pretrained segmentation. The
optional --checkpoint profile strictly authenticates the official ViT-B weights.
The image encoder is frozen for this inference reference; this is not a claim
about original SAM training. Upstream source stays byte-for-byte unchanged.
"""

import argparse
import hashlib
import importlib
import json
import sys
import types
from pathlib import Path

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

REVISION = "dca509fe793f601edb92606367a655c15ac00fdf"
CHECKPOINT_URL = (
    "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
)
CHECKPOINT_SHA256 = (
    "ec2df62732614e57411cdcf32a23ffdf28910380d03139ee0f4fcbe91eb8c912"
)

# Independent authority for the complete nine-file pinned bundle, including LICENSE.
SOURCE_MANIFEST_SHA256 = (
    "2551ee19fe09c423ab5e3d38801acc34adaf972a7c587fff8f1c085096bbe618"
)


def sha256(path):
    """Hash a file without loading the full checkpoint into memory."""
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def native_modules():
    """Verify and import the bundled modeling package, bypassing optional APIs."""
    root = Path(__file__).resolve().parent / "sam_native"
    manifest_path = root / "source-manifest.json"
    if sha256(manifest_path) != SOURCE_MANIFEST_SHA256:
        raise ValueError(
            "Native source manifest does not match the pinned complete bundle"
        )
    manifest = json.loads(manifest_path.read_text())
    if manifest["revision"] != REVISION:
        raise ValueError("Unexpected native source revision")
    for relative, record in manifest["files"].items():
        if sha256(root / relative) != record["sha256"]:
            raise ValueError(f"Native source hash mismatch: {relative}")
    # Namespace adapter only: no upstream code modifications or optional imports.
    name = "_sam_verified_native"
    if name not in sys.modules:
        package = types.ModuleType(name)
        package.__path__ = [str(root / "segment_anything")]
        sys.modules[name] = package
    return importlib.import_module(f"{name}.modeling"), importlib.import_module(
        f"{name}.build_sam"
    )


def build_model(checkpoint=None):
    """Construct the reduced offline profile or strict official ViT-B model."""
    native, builder = native_modules()
    if checkpoint is not None:
        if sha256(checkpoint) != CHECKPOINT_SHA256:
            raise ValueError(
                "Checkpoint is not the authenticated official SAM ViT-B"
            )
        model = builder.build_sam_vit_b()
        state = torch.load(checkpoint, map_location="cpu", weights_only=True)
        model.load_state_dict(state, strict=True)
    else:
        model = native.Sam(
            image_encoder=native.ImageEncoderViT(
                img_size=64,
                patch_size=16,
                embed_dim=64,
                depth=2,
                num_heads=4,
                mlp_ratio=2,
                out_chans=32,
                qkv_bias=True,
                use_rel_pos=True,
                window_size=4,
                global_attn_indexes=(1,),
            ),
            prompt_encoder=native.PromptEncoder(32, (4, 4), (64, 64), 16),
            mask_decoder=native.MaskDecoder(
                transformer_dim=32,
                transformer=native.TwoWayTransformer(2, 32, 4, 64),
                iou_head_hidden_dim=32,
            ),
        )
    model.eval()
    model.image_encoder.requires_grad_(False)
    return model


def resized_shape(original_size, target):
    """Return official longest-side dimensions with round-half-up behavior."""
    if len(original_size) != 2 or min(original_size) <= 0 or target <= 0:
        raise ValueError("Image dimensions and target must be positive")
    height, width = original_size
    scale = target / max(height, width)
    result = (int(height * scale + 0.5), int(width * scale + 0.5))
    if min(result) < 1:
        raise ValueError("Image aspect ratio collapses the resized short side")
    return result


# region: preprocess
def preprocess(model, image):
    """PIL bilinear resize, RGB 0..255 normalization, then right/bottom padding."""
    image = image.convert("RGB")
    original_size = (image.height, image.width)
    target = model.image_encoder.img_size
    input_size = resized_shape(original_size, target)
    resized = image.resize(input_size[::-1], Image.Resampling.BILINEAR)
    pixels = torch.from_numpy(np.array(resized)).permute(2, 0, 1).float()
    normalized = (pixels - model.pixel_mean) / model.pixel_std
    padded = F.pad(
        normalized, (0, target - input_size[1], 0, target - input_size[0])
    )
    return padded.unsqueeze(0), input_size, original_size


# endregion: preprocess


# region: embedding
@torch.no_grad()
def encode_image(model, image):
    """Cache once; subsequent prompts reuse the same tensor without re-encoding."""
    padded, input_size, original_size = preprocess(model, image)
    embedding = model.image_encoder(padded)
    return {
        "embedding": embedding,
        "input_size": input_size,
        "original_size": original_size,
        "padded": padded,
    }


# endregion: embedding


def transform_coords(coords, original_size, input_size, inverse=False):
    """Transform (x,y) using the actual rounded width and height independently."""
    if coords.shape[-1] != 2 or not torch.isfinite(coords).all():
        raise ValueError("Coordinates must be finite (x,y) pairs")
    if min(original_size) <= 0 or min(input_size) <= 0:
        raise ValueError("Coordinate frame dimensions must be positive")
    scale = coords.new_tensor(
        [input_size[1] / original_size[1], input_size[0] / original_size[0]]
    )
    return coords / scale if inverse else coords * scale


# region: sparse
def encode_sparse(encoder, points, labels, boxes=None):
    """Match native Fourier position plus point/corner type embeddings."""
    if (
        points.ndim != 3
        or points.shape[-1] != 2
        or labels.shape != points.shape[:2]
    ):
        raise ValueError("Expected points [B,N,2] and labels [B,N]")
    if not torch.isin(labels, labels.new_tensor([-1, 0, 1])).all():
        raise ValueError(
            "Point labels must be -1 padding, 0 background, or 1 foreground"
        )
    if not torch.isfinite(points).all():
        raise ValueError("Point coordinates must be finite")
    if boxes is None:
        points = torch.cat(
            [points, points.new_zeros(points.shape[0], 1, 2)], dim=1
        )
        labels = torch.cat(
            [labels, labels.new_full((labels.shape[0], 1), -1)], dim=1
        )
    sparse = encoder.pe_layer.forward_with_coords(
        points + 0.5, encoder.input_image_size
    )
    sparse[labels == -1] = encoder.not_a_point_embed.weight
    sparse[labels == 0] += encoder.point_embeddings[0].weight
    sparse[labels == 1] += encoder.point_embeddings[1].weight
    if boxes is not None:
        if (
            boxes.shape != (points.shape[0], 4)
            or not torch.isfinite(boxes).all()
        ):
            raise ValueError("Expected finite boxes [B,4]")
        if (boxes[:, 2:] < boxes[:, :2]).any():
            raise ValueError("Box maximum must follow its minimum")
        corners = encoder.pe_layer.forward_with_coords(
            boxes.reshape(-1, 2, 2) + 0.5, encoder.input_image_size
        )
        corners[:, 0] += encoder.point_embeddings[2].weight
        corners[:, 1] += encoder.point_embeddings[3].weight
        sparse = torch.cat([sparse, corners], dim=1)
    return sparse


# endregion: sparse


# region: dense
def encode_dense(encoder, batch_size, mask_logits=None):
    """Use low-resolution logits for refinement or the learned no-mask vector."""
    if batch_size <= 0:
        raise ValueError("Prompt batch must be positive")
    if mask_logits is not None:
        expected = (batch_size, 1, *encoder.mask_input_size)
        if (
            mask_logits.shape != expected
            or not torch.isfinite(mask_logits).all()
        ):
            raise ValueError(f"Expected finite mask logits {expected}")
        return encoder.mask_downscaling(mask_logits)
    return encoder.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(
        batch_size, -1, *encoder.image_embedding_size
    )


# endregion: dense


# region: two-way
def two_way(decoder, embedding, position, sparse, dense):
    """Expose native token setup and its two-way transformer without changing it."""
    if embedding.shape[0] != 1:
        raise ValueError("Reuse one image embedding per prompt batch")
    output = torch.cat([decoder.iou_token.weight, decoder.mask_tokens.weight])
    output = output.unsqueeze(0).expand(sparse.shape[0], -1, -1)
    tokens = torch.cat([output, sparse], dim=1)
    source = embedding.repeat_interleave(tokens.shape[0], dim=0) + dense
    position = position.repeat_interleave(tokens.shape[0], dim=0)
    # Each native layer: self-attn, token->image, token MLP, image->token.
    # Then native final token->image attention and LayerNorm.
    updated_tokens, updated_image = decoder.transformer(
        source, position, tokens
    )
    return updated_tokens, updated_image, source.shape


# endregion: two-way


# region: masks
def generate_masks(decoder, tokens, image_tokens, image_shape):
    """Hypernetwork vectors dot upsampled features; the IoU token regresses quality."""
    batch, channels, height, width = image_shape
    spatial = image_tokens.transpose(1, 2).reshape(
        batch, channels, height, width
    )
    upscaled = decoder.output_upscaling(spatial)
    mask_tokens = tokens[:, 1 : 1 + decoder.num_mask_tokens]
    vectors = torch.stack(
        [
            mlp(mask_tokens[:, i])
            for i, mlp in enumerate(decoder.output_hypernetworks_mlps)
        ],
        dim=1,
    )
    logits = (vectors @ upscaled.flatten(2)).reshape(
        batch, decoder.num_mask_tokens, *upscaled.shape[-2:]
    )
    scores = decoder.iou_prediction_head(tokens[:, 0])
    return logits, scores


# endregion: masks


# region: ranking
def select_masks(logits, scores, multimask=True):
    """Keep the dedicated single token or all three ambiguity tokens, then rank."""
    if (
        logits.ndim != 4
        or logits.shape[1] != 4
        or scores.shape != logits.shape[:2]
    ):
        raise ValueError(
            "Expected four native masks and matching predicted IoUs"
        )
    if not torch.isfinite(logits).all() or not torch.isfinite(scores).all():
        raise ValueError("Mask logits and quality predictions must be finite")
    selected = slice(1, 4) if multimask else slice(0, 1)
    logits, scores = logits[:, selected], scores[:, selected]
    # Explicit application policy; native multimask API does not pick a winner.
    best = scores.argmax(dim=1)
    return logits, scores, best


# endregion: ranking


# region: postprocess
def postprocess(logits, target, input_size, original_size):
    """Upscale logits to padded square, crop padding, resize back, then threshold."""
    if logits.ndim != 4 or min(logits.shape) <= 0:
        raise ValueError("Expected nonempty logits [B,K,H,W]")
    if (
        min(input_size) <= 0
        or max(input_size) > target
        or min(original_size) <= 0
    ):
        raise ValueError("Invalid input or original dimensions")
    square = F.interpolate(
        logits, (target, target), mode="bilinear", align_corners=False
    )
    cropped = square[..., : input_size[0], : input_size[1]]
    full = F.interpolate(
        cropped, original_size, mode="bilinear", align_corners=False
    )
    return full, full > 0.0


# endregion: postprocess


@torch.no_grad()
def predict(
    model, cached, points, labels, boxes=None, mask_logits=None, multimask=True
):
    """Decode original-pixel prompts using an existing image embedding."""
    points = transform_coords(
        points.float(), cached["original_size"], cached["input_size"]
    )
    if boxes is not None:
        boxes = transform_coords(
            boxes.float().reshape(-1, 2, 2),
            cached["original_size"],
            cached["input_size"],
        ).reshape(-1, 4)
    sparse = encode_sparse(model.prompt_encoder, points, labels, boxes)
    dense = encode_dense(model.prompt_encoder, points.shape[0], mask_logits)
    tokens, image, shape = two_way(
        model.mask_decoder,
        cached["embedding"],
        model.prompt_encoder.get_dense_pe(),
        sparse,
        dense,
    )
    logits, scores = generate_masks(model.mask_decoder, tokens, image, shape)
    logits, scores, best = select_masks(logits, scores, multimask)
    full, masks = postprocess(
        logits,
        model.image_encoder.img_size,
        cached["input_size"],
        cached["original_size"],
    )
    return {
        "masks": masks,
        "scores": scores,
        "best": best,
        "low_res_logits": logits,
        "full_logits": full,
        "sparse": sparse,
    }


def smoke():
    """Exercise the native reduced profile on synthetic pixels without network."""
    torch.set_num_threads(2)
    torch.manual_seed(7)
    model = build_model()
    pixels = np.arange(39 * 57 * 3, dtype=np.uint8).reshape(39, 57, 3)
    cached = encode_image(model, Image.fromarray(pixels))
    points = torch.tensor([[[20.0, 15.0]]])
    result = predict(model, cached, points, torch.ones(1, 1, dtype=torch.long))
    refined = predict(
        model,
        cached,
        points,
        torch.zeros(1, 1, dtype=torch.long),
        mask_logits=result["low_res_logits"][:, :1],
        multimask=False,
    )
    boxed = predict(
        model,
        cached,
        points,
        torch.ones(1, 1, dtype=torch.long),
        boxes=torch.tensor([[2.0, 3.0, 40.0, 30.0]]),
    )
    return {
        "profile": "reduced native random weights; mechanics only",
        "sourceRevision": REVISION,
        "embeddingShape": list(cached["embedding"].shape),
        "maskShape": list(result["masks"].shape),
        "predictedIoU": result["scores"].tolist(),
        "bestIndex": result["best"].tolist(),
        "singleShape": list(refined["masks"].shape),
        "boxSparseShape": list(boxed["sparse"].shape),
        "encoderFrozen": True,
    }


def main():
    """Run a portable mechanics check or explicit official checkpoint inference."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--checkpoint", type=Path)
    parser.add_argument("--image", type=Path)
    parser.add_argument("--point", nargs=2, type=float, metavar=("X", "Y"))
    parser.add_argument("--output", type=Path, default=Path("sam-output"))
    args = parser.parse_args()
    if args.smoke:
        print(json.dumps(smoke(), indent=2))
        return
    if args.checkpoint is None or args.image is None or args.point is None:
        parser.error(
            "Use --smoke or supply --checkpoint, --image, and --point X Y"
        )
    torch.set_num_threads(2)
    torch.manual_seed(7)
    model = build_model(args.checkpoint)
    cached = encode_image(model, Image.open(args.image))
    result = predict(
        model,
        cached,
        torch.tensor([[args.point]]),
        torch.ones(1, 1, dtype=torch.long),
    )
    args.output.mkdir(parents=True, exist_ok=True)
    for index, mask in enumerate(result["masks"][0]):
        Image.fromarray(mask.numpy().astype(np.uint8) * 255).save(
            args.output / f"mask-{index}.png"
        )
    print(
        json.dumps(
            {
                "predictedIoU": result["scores"].tolist(),
                "bestIndex": result["best"].tolist(),
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()