"""Predict AI Worker arm masks with the native SAM3 image model."""

import argparse
import hashlib
import json
from pathlib import Path

import numpy as np
import torch
from huggingface_hub import hf_hub_download
from PIL import Image


CONFIG_PATH = Path(__file__).with_name("inference_config.json")


def _activation(activation, linear, values):
    """Use the same unfused activation path as the fine-tuned checkpoint."""
    if activation in (torch.nn.GELU, torch.nn.functional.gelu):
        return torch.nn.functional.gelu(linear(values))
    if activation in (torch.nn.ReLU, torch.nn.functional.relu):
        return torch.nn.functional.relu(linear(values))
    raise ValueError(f"Unsupported activation: {activation}")


def load_predictor(checkpoint=None):
    """Load the release checkpoint and its fixed 504px inference settings."""
    from sam3.model import vitdet
    from sam3.model.sam3_image_processor import Sam3Processor
    from sam3.model_builder import build_sam3_image_model

    if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
        raise RuntimeError("A BF16-capable CUDA GPU is required. See README.md.")
    config = json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
    if checkpoint is None:
        checkpoint = hf_hub_download(config["model_id"], config["checkpoint"])
    checkpoint = Path(checkpoint)
    with checkpoint.open("rb") as stream:
        actual = hashlib.file_digest(stream, "sha256").hexdigest()
    if actual != config["checkpoint_sha256"]:
        raise ValueError("Checkpoint checksum mismatch. Download this release again.")

    torch.set_num_threads(1)
    vitdet.addmm_act = _activation
    model = build_sam3_image_model(
        checkpoint_path=str(checkpoint), load_from_HF=False,
        device="cpu", eval_mode=True, enable_inst_interactivity=False,
    )
    trunk = model.backbone.vision_backbone.trunk
    grid = config["resolution"] // trunk.patch_embed.proj.kernel_size[0]
    for block in trunk.blocks:
        attention = block.attn
        if block.window_size > 0 or attention.freqs_cis is None:
            continue
        attention.input_size = (grid, grid)
        scale = attention.rope_pt_size[0] / grid if attention.rope_interp else 1.0
        frequencies = attention.compute_cis(end_x=grid, end_y=grid, scale_pos=scale)
        attention.freqs_cis = frequencies
        if attention.use_rope_real:
            attention.freqs_cis_real = frequencies.real
            attention.freqs_cis_imag = frequencies.imag
    torch.set_num_threads(4)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.manual_seed(438)
    model = model.cuda().eval()
    processor = Sam3Processor(
        model, resolution=config["resolution"], device="cuda",
        confidence_threshold=config["confidence_threshold"],
    )
    return model, processor, config


@torch.inference_mode()
def predict_image(image, processor, config):
    """Return native-resolution instance masks, XYXY boxes and scores."""
    with torch.autocast("cuda", dtype=torch.bfloat16):
        state = processor.set_image(image.convert("RGB"))
        result = processor.set_text_prompt(prompt=config["prompt"], state=state)
    return (
        result["masks"][:, 0].cpu().numpy(),
        result["boxes"].float().cpu().numpy(),
        result["scores"].float().cpu().numpy(),
    )


def save_outputs(directory, masks, boxes, scores):
    """Save a white foreground union mask and all individual predictions."""
    directory = Path(directory)
    directory.mkdir(parents=True, exist_ok=True)
    Image.fromarray(masks.any(axis=0).astype(np.uint8) * 255).save(directory / "mask.png")
    np.savez_compressed(directory / "instances.npz", masks=masks, boxes_xyxy=boxes, scores=scores)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("image", type=Path, help="Input RGB image")
    parser.add_argument("--checkpoint", type=Path, help="Local model.pt; otherwise download it")
    parser.add_argument("--output-dir", type=Path, default=Path("prediction"))
    args = parser.parse_args()
    if args.image.resolve() in {(args.output_dir / name).resolve() for name in ["mask.png", "instances.npz"]}:
        parser.error("The output directory would overwrite the input image.")
    with Image.open(args.image) as source:
        image = source.convert("RGB")
    model, processor, config = load_predictor(args.checkpoint)
    masks, boxes, scores = predict_image(image, processor, config)
    save_outputs(args.output_dir, masks, boxes, scores)
    print(f"Saved {len(scores)} predicted instances to {args.output_dir}")


if __name__ == "__main__":
    main()