"""3AM inference engine for the Hugging Face Space.

3AM routing-fusion model: SAM2.1 Hiera-L whose memory attention is wrapped by
RoutingMemoryAttention (routing tokens + FiLM fusion of MUSt3R features), MUSt3R-512 run with
memory_mode='kv' keeping every frame's KV (no keyframe selection, no KD-tree, no retrieval),
and a SAM2 memory policy that keeps the previous frame plus the most recent positive masks.

Frames are letterboxed (never cropped) to a MUSt3R-supported aspect ratio; masks are mapped
back to the original resolution.
"""
from __future__ import annotations

import ast
import inspect
import math
import os
import textwrap
from types import MethodType

import numpy as np
import roma
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from PIL import Image
from torchvision import transforms as T

from sam2.build_sam import build_sam2_video_predictor
from sam2.utils.misc import fill_holes_in_mask_scores
from sam2.modeling.sam.transformer import RoPEAttention
from sam2.modeling.position_encoding import apply_rotary_enc
from sam2.modeling.sam2_base import SAM2Base

from must3r.model import ActivationType, apply_activation, load_model, get_pointmaps_activation
from must3r.tools.image import get_HW_resolution
import must3r.tools.path_to_dust3r  # noqa: F401  registers dust3r on sys.path
from dust3r.post_process import estimate_focal_knowing_depth

MUST3R_SIZE = 512
SAM2_IMAGE_SIZE = 1024
MUST3R_LAYERS = (0, 5, 8, 12)          # encoder output + decoder blocks 5, 8, 12
MEMORY_POLICY_NAME = 'previous_frame_then_recent_positive_masks'

IMAGE_TRANSFORM = T.Compose([
    T.ToTensor(),
    T.Resize((SAM2_IMAGE_SIZE, SAM2_IMAGE_SIZE), interpolation=T.InterpolationMode.NEAREST_EXACT),
    T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
])


# ----------------------------------------------------------------------------------------
# Geometry helpers (letterbox instead of crop)
# ----------------------------------------------------------------------------------------

def letterbox_geometry(width, height):
    target_h, target_w = get_HW_resolution(height, width, MUST3R_SIZE)
    padded_w = max(width, int(math.ceil(height * target_w / target_h)))
    padded_h = max(height, int(math.ceil(width * target_h / target_w)))
    return dict(width=width, height=height, padded_w=padded_w, padded_h=padded_h,
                left=(padded_w - width) // 2, top=(padded_h - height) // 2,
                target_w=target_w, target_h=target_h)


def pad_frame(rgb: np.ndarray, g: dict) -> np.ndarray:
    padded = np.zeros((g['padded_h'], g['padded_w'], 3), np.uint8)
    padded[g['top']:g['top'] + g['height'], g['left']:g['left'] + g['width']] = rgb
    return padded


def sam2_input(padded: np.ndarray) -> torch.Tensor:
    """[3, 1024, 1024] normalized SAM2 input from a padded RGB uint8 frame (nearest-exact)."""
    return IMAGE_TRANSFORM(Image.fromarray(padded))


def must3r_input(padded: np.ndarray, g: dict) -> torch.Tensor:
    """[3, target_h, target_w] in [-1, 1] for MUSt3R (nearest-exact)."""
    import cv2
    small = cv2.resize(padded, (g['target_w'], g['target_h']), interpolation=cv2.INTER_NEAREST_EXACT)
    return torch.from_numpy(small).permute(2, 0, 1).contiguous().float() / 255 * 2 - 1


def points_to_sam2(points, labels, g: dict, device):
    """Original-pixel clicks -> SAM2 1024x1024 coordinates on the padded frame."""
    coords = [[(x + g['left']) * SAM2_IMAGE_SIZE / g['padded_w'],
               (y + g['top']) * SAM2_IMAGE_SIZE / g['padded_h']] for x, y in points]
    return dict(point_coords=torch.tensor([coords], dtype=torch.float32, device=device),
                point_labels=torch.tensor([list(labels)], dtype=torch.int32, device=device))


def logits_to_original_mask(logits: torch.Tensor, g: dict) -> np.ndarray:
    """[1024, 1024] logits -> bool mask at the original frame resolution."""
    resized = F.interpolate(logits[None, None].float(), size=(g['padded_h'], g['padded_w']),
                            mode='bilinear', align_corners=False)[0, 0]
    y, x = g['top'], g['left']
    return (resized[y:y + g['height'], x:x + g['width']] > 0).cpu().numpy()


def original_mask_to_sam2(mask: np.ndarray, g: dict, device) -> torch.Tensor:
    """bool mask at original resolution -> [1, 1, 1024, 1024] bool mask input on the padded frame."""
    padded = np.zeros((g['padded_h'], g['padded_w']), np.float32)
    padded[g['top']:g['top'] + g['height'], g['left']:g['left'] + g['width']] = mask.astype(np.float32)
    t = torch.from_numpy(padded)[None, None]
    t = F.interpolate(t, size=(SAM2_IMAGE_SIZE, SAM2_IMAGE_SIZE), mode='nearest-exact')
    return (t >= 0.5).to(device)


# ----------------------------------------------------------------------------------------
# Inlined 3AM modules (identical to the training / evaluation code)
# ----------------------------------------------------------------------------------------

def positional_encoding(positions, freqs, dim=1):
    assert 0 <= dim < positions.ndim
    bands = 2 ** torch.arange(freqs, dtype=positions.dtype, device=positions.device)
    expanded = positions.unsqueeze(dim + 1) * bands.view(-1, *([1] * (positions.ndim - dim - 1)))
    expanded = expanded.reshape(*positions.shape[:dim], -1, *positions.shape[dim + 1:])
    return torch.cat([torch.sin(expanded), torch.cos(expanded), positions], dim=dim)


@torch.autocast("cuda", dtype=torch.float32)
def postprocess_must3r_output(pointmaps, pointmaps_activation=ActivationType.NORM_EXP, compute_cam=True):
    out = {}
    C = pointmaps.shape[-1]
    out['pts3d'] = apply_activation(pointmaps[..., :3], activation=pointmaps_activation)
    if C >= 6:
        out['pts3d_local'] = apply_activation(pointmaps[..., 3:6], activation=pointmaps_activation)
    if C in (4, 7):
        out['conf'] = 1.0 + pointmaps[..., -1].exp()
    if not compute_cam:
        return out
    batch_dims = out['pts3d'].shape[:-3]
    nbd = len(batch_dims)
    H, W = out['conf'].shape[-2:]
    pp = torch.tensor((W / 2, H / 2), device=out['pts3d'].device)
    focal = estimate_focal_knowing_depth(
        out['pts3d_local'].reshape(math.prod(batch_dims), H, W, 3), pp, focal_mode='weiszfeld')
    out['focal'] = focal.reshape(*batch_dims)
    R, T_ = roma.rigid_points_registration(
        out['pts3d_local'].reshape(*batch_dims, -1, 3),
        out['pts3d'].reshape(*batch_dims, -1, 3),
        weights=out['conf'].reshape(*batch_dims, -1) - 1.0, compute_scaling=False)
    c2w = torch.eye(4, device=out['pts3d'].device).view(*([1] * nbd), 4, 4).repeat(*batch_dims, 1, 1)
    c2w[..., :3, :3] = R
    c2w[..., :3, 3] = T_.view(*batch_dims, 3)
    out['c2w'] = c2w
    ys, xs = torch.meshgrid(torch.arange(H, device=out['pts3d'].device),
                            torch.arange(W, device=out['pts3d'].device), indexing='ij')
    f = out['focal'].reshape(*batch_dims, 1, 1)
    x = (xs - pp[0]) / f
    y = (ys - pp[1]) / f
    d_cam = F.normalize(torch.stack([x, y, torch.ones_like(x)], dim=-1), dim=-1)
    d_world = torch.einsum('...ij,...hwj->...hwi', R, d_cam)
    o_world = c2w[..., :3, 3].view(*batch_dims, 1, 1, 3).expand(*batch_dims, H, W, 3)
    out['ray_origin'] = o_world
    out['ray_dir'] = d_world
    out['ray_plucker'] = torch.cat([torch.cross(o_world, d_world, dim=-1), d_world], dim=-1)
    return out


class LayerNorm2d(nn.Module):
    def __init__(self, num_channels: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(num_channels))
        self.bias = nn.Parameter(torch.zeros(num_channels))
        self.eps = eps

    def forward(self, x):
        u = x.mean(1, keepdim=True)
        s = (x - u).pow(2).mean(1, keepdim=True)
        x = (x - u) / torch.sqrt(s + self.eps)
        return self.weight[:, None, None] * x + self.bias[:, None, None]


class RoutingFusion(nn.Module):
    """Additive FiLM-gated fusion of MUSt3R decoder features + explicit 3D geometry into feat_2d."""
    def __init__(self, d_2d=256, d_3d=(1024, 768, 768, 768), d_route=256, d_hidden=384, explicit_3d_freqs=6):
        super().__init__()
        self.L = len(d_3d)
        self.d_2d = d_2d
        self.d_route = d_route
        self.explicit_3d_freqs = explicit_3d_freqs
        explicit_3d_in_ch = 3 * (2 * explicit_3d_freqs + 1) + 6
        self.explicit_3d_embedding = nn.Conv2d(explicit_3d_in_ch, d_2d, kernel_size=16, stride=16)
        self.proj = nn.ModuleList([
            nn.Sequential(nn.Conv2d(c, d_hidden, kernel_size=1, bias=True), LayerNorm2d(d_hidden), nn.GELU(),
                          nn.Conv2d(d_hidden, d_2d, kernel_size=3, padding=1, bias=True))
            for c in d_3d])
        self.mod_mlp = nn.ModuleList([
            nn.Sequential(nn.Linear(d_route, 2 * d_route), nn.GELU(), nn.Linear(2 * d_route, 2 * d_2d + 1))
            for _ in range(self.L + 1)])
        self.base_mod_mlp = nn.Sequential(nn.Linear(d_route, 2 * d_route), nn.GELU(), nn.Linear(2 * d_route, 2 * d_2d))
        nn.init.zeros_(self.base_mod_mlp[-1].weight)
        nn.init.zeros_(self.base_mod_mlp[-1].bias)

    def _film_add(self, feat, delta, mod, C, per_layer_alpha):
        gamma = 1.0 + mod[:, :C, None, None]
        beta = mod[:, C:2 * C, None, None]
        alpha = mod[:, 2 * C:, None, None]
        feat = feat + alpha * (delta * gamma + beta)
        per_layer_alpha.append(alpha.view(feat.shape[0]))
        return feat

    def forward(self, feat_2d, feat_3d_list, r, explicit_3d):
        B, C, H, W = feat_2d.shape
        assert r.shape[1] == self.L + 2
        base_mod = self.base_mod_mlp(r[:, self.L + 1, :]).to(dtype=feat_2d.dtype)
        feat = feat_2d * (1.0 + base_mod[:, :C, None, None]) + base_mod[:, C:2 * C, None, None]
        per_layer_alpha = []
        for l in range(self.L):
            delta = self.proj[l](feat_3d_list[l])
            delta = F.interpolate(delta, size=(H, W), mode='bilinear', align_corners=False)
            mod = self.mod_mlp[l](r[:, l, :]).to(dtype=feat.dtype)
            feat = self._film_add(feat, delta, mod, C, per_layer_alpha)
        pts3d = explicit_3d[:, :3]
        plucker = explicit_3d[:, 3:]
        pe_in = torch.cat([positional_encoding(pts3d, self.explicit_3d_freqs, dim=1), plucker], dim=1)
        delta = self.explicit_3d_embedding(pe_in.to(self.explicit_3d_embedding.weight.dtype))
        delta = F.interpolate(delta, size=(H, W), mode='bilinear', align_corners=False).to(feat.dtype)
        mod = self.mod_mlp[self.L](r[:, self.L, :]).to(dtype=feat.dtype)
        feat = self._film_add(feat, delta, mod, C, per_layer_alpha)
        alpha_stack = torch.stack(per_layer_alpha, dim=1)
        return feat, alpha_stack, None, base_mod.detach()


def _cross_attn_with_routing(ca, full_q, mem_k_in, mem_v_in, HW, num_k_exclude_rope=0):
    q = ca.q_proj(full_q)
    k = ca.k_proj(mem_k_in)
    v = ca.v_proj(mem_v_in)
    q = ca._separate_heads(q, ca.num_heads)
    k = ca._separate_heads(k, ca.num_heads)
    v = ca._separate_heads(v, ca.num_heads)
    if isinstance(ca, RoPEAttention):
        freqs_cis = ca.freqs_cis.to(q.device)
        assert freqs_cis.shape[0] == HW
        num_k_rope = k.size(-2) - num_k_exclude_rope
        q_spatial = q[:, :, :HW]
        q_routing = q[:, :, HW:]
        q_spatial, k_rope = apply_rotary_enc(q_spatial, k[:, :, :num_k_rope], freqs_cis=freqs_cis,
                                             repeat_freqs_k=ca.rope_k_repeat)
        q = torch.cat([q_spatial, q_routing], dim=2)
        k = torch.cat([k_rope, k[:, :, num_k_rope:]], dim=2)
    dropout_p = ca.dropout_p if ca.training else 0.0
    out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
    out = ca._recombine_heads(out)
    return ca.out_proj(out)


class RoutingMemoryAttention(nn.Module):
    """Replaces SAM2's MemoryAttention: routing tokens ride inside every layer and see the memory bank."""
    def __init__(self, base, fusion: RoutingFusion, d_route=256, routing_num_heads=8, num_routing_tokens_per_layer=2):
        super().__init__()
        self.base = base
        self.fusion = fusion
        self.d_route = d_route
        self.L = fusion.L
        self.G = self.L + 2
        self.K = num_routing_tokens_per_layer
        self.N = self.K * self.G
        d_model = base.d_model
        kv_in_dim = base.layers[0].cross_attn_image.kv_in_dim
        self.routing_token = nn.Parameter(torch.zeros(1, self.N, d_model))
        self.r_in_proj = nn.Linear(d_route, d_model)
        self.r_out_proj = nn.Linear(self.K * d_model, d_route)
        self.r0 = nn.Parameter(torch.zeros(1, self.G, d_route))
        self._expected_obj_ptr_tokens = max(1, d_model // kv_in_dim)
        self.obj_ptr_to_r = nn.Linear(self._expected_obj_ptr_tokens * kv_in_dim, self.G * d_route)
        nn.init.zeros_(self.obj_ptr_to_r.weight)
        nn.init.zeros_(self.obj_ptr_to_r.bias)
        nn.init.zeros_(self.r_out_proj.weight)
        nn.init.zeros_(self.r_out_proj.bias)
        self._feat_2d_seq = None
        self._feat_3d_seq = None
        self._explicit_3d_seq = None
        self._frame_idx = 1
        self._r = None
        self._feats_for_memory_out = None
        self.history = {}

    @property
    def num_layers(self):
        return self.base.num_layers

    @property
    def d_model(self):
        return self.base.d_model

    def reset_state(self, feat_2d_seq, feat_3d_seq, explicit_3d_seq, B, device, dtype, feats_for_memory_out=None):
        self._feat_2d_seq = feat_2d_seq
        self._feat_3d_seq = feat_3d_seq
        self._explicit_3d_seq = explicit_3d_seq
        self._frame_idx = 1
        self._r = self.r0.expand(B, -1, -1).to(device=device, dtype=dtype)
        self._feats_for_memory_out = feats_for_memory_out
        self.history = {'r': [self._r.detach()], 'alpha': [], 'base_mod': []}

    def clear_state(self):
        self._feat_2d_seq = None
        self._feat_3d_seq = None
        self._explicit_3d_seq = None
        self._frame_idx = 1
        self._r = None
        self._feats_for_memory_out = None
        self.history = {}

    def forward(self, curr, memory, curr_pos=None, memory_pos=None, num_obj_ptr_tokens=0):
        curr_t = curr[0] if isinstance(curr, list) else curr
        curr_pos_t = curr_pos[0] if isinstance(curr_pos, list) else curr_pos
        base = self.base
        B_size = curr_t.shape[1]
        if self._frame_idx == 1 and num_obj_ptr_tokens > 0:
            assert num_obj_ptr_tokens == self._expected_obj_ptr_tokens
            obj_ptr_seq = memory[-num_obj_ptr_tokens:]
            obj_ptr_recon = obj_ptr_seq.transpose(0, 1).flatten(1, 2)
            r_delta = self.obj_ptr_to_r(obj_ptr_recon).view(B_size, self.G, self.d_route)
            self._r = self._r + r_delta.to(self._r.dtype)
        t = min(self._frame_idx, len(self._feat_2d_seq) - 1)
        feat_2d = self._feat_2d_seq[t].to(device=curr_t.device, dtype=curr_t.dtype)
        feat_3d = [f.to(device=curr_t.device, dtype=curr_t.dtype) for f in self._feat_3d_seq[t]]
        explicit_3d = self._explicit_3d_seq[t].to(device=curr_t.device, dtype=curr_t.dtype)
        fused, alpha, _, base_mod = self.fusion(feat_2d, feat_3d, self._r, explicit_3d)
        fused_flat = rearrange(fused, 'b c h w -> (h w) b c')
        HW, B, C = fused_flat.shape
        N, K, L, G = self.N, self.K, self.L, self.G
        if self._feats_for_memory_out is not None and t < len(self._feats_for_memory_out):
            self._feats_for_memory_out[t] = [fused_flat]
        r_inject = self.r_in_proj(self._r).to(curr_t.dtype)
        r_inject = r_inject.unsqueeze(2).expand(B, G, K, C).reshape(B, N, C)
        rt = self.routing_token.expand(1, N, C).repeat(B, 1, 1)
        rt = rt + r_inject
        spatial = fused_flat
        if base.pos_enc_at_input and curr_pos_t is not None:
            spatial = spatial + 0.1 * curr_pos_t
        assert base.batch_first
        spatial = spatial.transpose(0, 1)
        spatial_pos = curr_pos_t.transpose(0, 1)
        mem_b = memory.transpose(0, 1)
        mem_pos_b = memory_pos.transpose(0, 1)
        rt_pos = torch.zeros_like(rt)
        for layer in base.layers:
            tgt = torch.cat([spatial, rt], dim=1)
            qpos = torch.cat([spatial_pos, rt_pos], dim=1)
            tgt2 = layer.norm1(tgt)
            full_q = tgt2 + qpos if layer.pos_enc_at_attn else tgt2
            tgt2 = _cross_attn_with_routing(layer.self_attn, full_q, full_q, tgt2, HW=HW, num_k_exclude_rope=N)
            tgt = tgt + layer.dropout1(tgt2)
            tgt2 = layer.norm2(tgt)
            full_q = tgt2 + qpos if layer.pos_enc_at_cross_attn_queries else tgt2
            mem_k_in = mem_b + mem_pos_b if layer.pos_enc_at_cross_attn_keys else mem_b
            ca_out = _cross_attn_with_routing(layer.cross_attn_image, full_q, mem_k_in, mem_b, HW=HW,
                                              num_k_exclude_rope=num_obj_ptr_tokens)
            tgt = tgt + layer.dropout2(ca_out)
            tgt2 = layer.norm3(tgt)
            tgt2 = layer.linear2(layer.dropout(layer.activation(layer.linear1(tgt2))))
            tgt = tgt + layer.dropout3(tgt2)
            spatial, rt = tgt[:, :HW], tgt[:, HW:]
        spatial = base.norm(spatial)
        rt = base.norm(rt)
        spatial = spatial.transpose(0, 1)
        rt_per_layer = rt.reshape(B, G, K * C)
        new_r = self.r_out_proj(rt_per_layer).to(self._r.dtype)
        self._r = new_r
        self._frame_idx += 1
        self.history['r'].append(new_r.detach())
        self.history['alpha'].append(alpha)
        self.history['base_mod'].append(base_mod)
        return spatial


# ----------------------------------------------------------------------------------------
# SAM2 memory policy: previous frame, then the most recent nonempty masks
# ----------------------------------------------------------------------------------------

def _previous_visible_indices(outputs, frame_idx):
    previous = frame_idx - 1
    keys = [previous] if previous in outputs else []
    for k in sorted(outputs, reverse=True):
        if k >= previous:
            continue
        out = outputs[k]
        if '_memory_visible' not in out:
            out['_memory_visible'] = bool((out['pred_masks'] > 0).any())
        if out['_memory_visible']:
            keys.append(k)
    return keys


def _build_previous_visible_method(base_cls):
    """AST-patch SAM2Base._prepare_memory_conditioned_features so that only the memory *selection*
    changes (previous frame + recent positive masks, real pointer distances); encodings stay official."""
    original = base_cls._prepare_memory_conditioned_features
    tree = ast.parse(textwrap.dedent(inspect.getsource(original)))
    fn = tree.body[0]
    fn.name = '_prepare_previous_visible'
    fn.body.insert(1, ast.parse("assert not track_in_reverse, 'Harness uses forward-indexed passes'").body[0])
    fn.body.insert(2, ast.parse("visible_keys = _previous_visible_indices(output_dict['non_cond_frame_outputs'], frame_idx)").body[0])
    counts = {'spatial': 0, 'pointers': 0}

    class Patch(ast.NodeTransformer):
        def visit_Assign(self, node):
            self.generic_visit(node)
            if (len(node.targets) == 1 and isinstance(node.targets[0], ast.Name) and node.targets[0].id == 'out'
                    and isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Attribute)
                    and node.value.func.attr == 'get' and isinstance(node.value.func.value, ast.Subscript)
                    and node.value.args and isinstance(node.value.args[0], ast.Name)
                    and node.value.args[0].id == 'prev_frame_idx'):
                counts['spatial'] += 1
                override = ast.parse('prev_frame_idx = visible_keys[t_rel-1] if t_rel <= len(visible_keys) else -1').body[0]
                return [override, node]
            return node

        def visit_For(self, node):
            self.generic_visit(node)
            if isinstance(node.target, ast.Name) and node.target.id == 't_diff':
                counts['pointers'] += 1
                replacement = ast.parse('for t_diff, t in [(frame_idx-k,k) for k in visible_keys[:max_obj_ptrs_in_encoder-1]]:\n pass').body[0]
                assert isinstance(node.body[0], ast.Assign) and node.body[0].targets[0].id == 't'
                replacement.body = node.body[1:]
                return replacement
            return node

    tree = Patch().visit(tree)
    ast.fix_missing_locations(tree)
    assert counts == {'spatial': 1, 'pointers': 1}, counts
    namespace = dict(original.__globals__)
    namespace['_previous_visible_indices'] = _previous_visible_indices
    exec(compile(tree, '<previous-visible>', 'exec'), namespace)
    return namespace['_prepare_previous_visible']


def install_memory_policy(predictor):
    method = _build_previous_visible_method(SAM2Base)
    predictor._prepare_memory_conditioned_features = MethodType(method, predictor)
    predictor.memory_temporal_stride_for_eval = 1


class PositiveMemory:
    """Keep useful (nonempty) outputs across arbitrarily long gaps of empty masks."""
    def __init__(self, capacity):
        self.capacity = capacity
        self.outputs = {}

    def restore(self, noncond):
        noncond.update(self.outputs)

    def remember(self, index, output):
        output['_memory_visible'] = bool((output['pred_masks'] > 0).any())
        if output['_memory_visible']:
            self.outputs[index] = output
        while len(self.outputs) > self.capacity:
            del self.outputs[min(self.outputs)]


# ----------------------------------------------------------------------------------------
# Model construction
# ----------------------------------------------------------------------------------------

import contextlib


@contextlib.contextmanager
def cpu_construction():
    """Build models without touching CUDA. SAM2 moves RoPE tables / position caches to 'cuda' at
    construction when torch.cuda.is_available(); on ZeroGPU that happens at import time and the
    resulting placeholder tensors cannot be deep-copied next to CPU tensors. Forward passes move
    those tables to the input device anyway."""
    original = torch.cuda.is_available
    torch.cuda.is_available = lambda: False
    try:
        yield
    finally:
        torch.cuda.is_available = original


def build_3am_predictor(checkpoint_path, device='cpu'):
    with cpu_construction():
        predictor = build_sam2_video_predictor('configs/sam2.1/sam2.1_hiera_l.yaml', device='cpu',
                                               apply_postprocessing=True).eval()
    d_route = predictor.memory_attention.d_model
    fusion = RoutingFusion(d_2d=256, d_3d=(1024, 768, 768, 768), d_route=d_route)
    predictor.memory_attention = RoutingMemoryAttention(base=predictor.memory_attention, fusion=fusion,
                                                        d_route=d_route, num_routing_tokens_per_layer=2).to(device).eval()
    weights = torch.load(checkpoint_path, map_location='cpu', mmap=True, weights_only=False)
    predictor.load_state_dict(weights.get('model', weights), strict=True)
    del weights
    install_memory_policy(predictor)
    return predictor.to(device)


def build_original_predictor(checkpoint_path, device='cpu'):
    with cpu_construction():
        predictor = build_sam2_video_predictor('configs/sam2.1/sam2.1_hiera_l.yaml', checkpoint_path,
                                               device='cpu', apply_postprocessing=True).eval()
    return predictor.to(device)


def build_must3r(weights_path, device='cpu'):
    from must3r.model.blocks.attention import has_xformers, toggle_memory_efficient_attention
    # xformers is off unless MUST3R_XFORMERS=1: on the Space the wheel has no CUDA ops for the runtime
    # GPU, and PyTorch SDPA gives the same features (verified against the reference implementation).
    toggle_memory_efficient_attention(enabled=bool(has_xformers and os.environ.get('MUST3R_XFORMERS', '0') == '1'))
    with cpu_construction():
        model = load_model(weights_path, device='cpu', img_size=MUST3R_SIZE, memory_mode='kv', verbose=False)
    model = [m.eval().to(device) for m in model]
    assert model[0].patch_size == 16
    assert get_pointmaps_activation(model[1]) == ActivationType.NORM_EXP
    return model


# ----------------------------------------------------------------------------------------
# Feature extraction
# ----------------------------------------------------------------------------------------

@torch.no_grad()
@torch.autocast(device_type='cuda', dtype=torch.bfloat16)
def get_image_feature(predictor, images):
    backbone_out = predictor.forward_image(images)
    backbone_out = {'backbone_fpn': backbone_out['backbone_fpn'].copy(),
                    'vision_pos_enc': backbone_out['vision_pos_enc'].copy()}
    backbone_out, vision_feats, vision_pos_embeds, feat_sizes = predictor._prepare_backbone_features(backbone_out)
    return vision_feats, vision_pos_embeds, feat_sizes


@torch.no_grad()
def run_must3r(must3r, padded_frames, g, device='cuda', progress=None):
    """MUSt3R over all frames, every frame kept in KV memory (no keyframe selection, no KD-tree).
    Returns must3r_feats: 4 x [N, C, h, w] (fp16, cpu) and explicit: [N, 9, target_h, target_w] (fp16, cpu)."""
    from must3r.engine.inference import inference_video_multi_ar
    encoder, decoder = must3r
    decoder.all_feats = []
    n = len(padded_frames)
    imgs = [must3r_input(p, g).to(device) for p in padded_frames]
    ts = [torch.tensor([g['target_h'], g['target_w']], dtype=torch.int32, device=device) for _ in range(n)]
    mem_batches = [min(2, n)] + [1] * max(0, n - 2)
    pointmaps_activation = get_pointmaps_activation(decoder, verbose=False)
    try:
        with torch.autocast('cuda', dtype=torch.bfloat16):
            inference_video_multi_ar(
                encoder, decoder, imgs, ts, mem_batches, max_bs=1, verbose=False, device=device,
                preserve_gpu_mem=False,   # keep the KV memory on the GPU (a CPU round trip per frame is very slow)
                post_process_function=lambda x: {'pts3d': x},   # geometry is recomputed below from layer 12
                viser_server=None, num_refinements_iterations=0, local_context_size=0,
                is_keyframe_function=lambda *_: True, scene_state=None,
                scene_state_update_function=lambda res, ss: ss)
        assert len(decoder.all_feats) == n, (len(decoder.all_feats), n)
        h, w = g['target_h'] // 16, g['target_w'] // 16
        per_view = decoder.all_feats
        feats = [rearrange(torch.stack([pv[k] for pv in per_view], 0).squeeze(1), 'b (h w) c -> b c h w', h=h, w=w).to(torch.float16)
                 for k in MUST3R_LAYERS]
        explicit = []
        shape = torch.tensor([[g['target_h'], g['target_w']]], dtype=torch.int32, device=device)
        with torch.autocast('cuda', dtype=torch.bfloat16):
            for i in range(n):
                feat_flat = per_view[i][12].to(device)
                if feat_flat.ndim == 2:
                    feat_flat = feat_flat[None]
                pm = decoder._compute_prediction_head(shape[:, None], 1, 1, [feat_flat], norm=True).squeeze()
                out = postprocess_must3r_output(pm[None], pointmaps_activation=pointmaps_activation, compute_cam=True)
                explicit.append(torch.cat((out['pts3d'], out['ray_plucker']), dim=-1).squeeze(0).permute(2, 0, 1).to(torch.float16).cpu())
                if progress is not None and (i % 10 == 0 or i == n - 1):
                    progress(i + 1, n)
        explicit = torch.stack(explicit, 0)
    finally:
        decoder.all_feats = []
        torch.cuda.empty_cache()
    return feats, explicit


# ----------------------------------------------------------------------------------------
# Tracking (evaluation tracker + positive-memory bank)
# ----------------------------------------------------------------------------------------

class Tracker3AM(nn.Module):
    def __init__(self, predictor, device='cuda'):
        super().__init__()
        self.predictor = predictor
        self.device = device

    @torch.no_grad()
    @torch.autocast(device_type='cuda', dtype=torch.bfloat16)
    def initialize(self, vision_feats_all, vision_pos_all, feat_sizes, must3r_feats, explicit_3d,
                   processing_order, mask_inputs=None, point_inputs=None):
        device, bf16 = self.device, torch.bfloat16
        order = list(processing_order)
        N = len(order)
        feat_2d_seq = [rearrange(vision_feats_all[i][-1], '(h w) b c -> b c h w', h=64, w=64).contiguous() for i in order]
        feat_3d_seq = [[must3r_feats[layer][i:i + 1] for layer in range(4)] for i in order]
        explicit_3d_seq = [explicit_3d[i:i + 1] for i in order]
        r0 = self.predictor.memory_attention.r0.expand(1, -1, -1).to(device=device, dtype=bf16)
        fused_0, _, _, _ = self.predictor.memory_attention.fusion(
            feat_2d_seq[0].to(device=device, dtype=bf16), [f.to(device=device, dtype=bf16) for f in feat_3d_seq[0]],
            r0, explicit_3d_seq[0].to(device=device, dtype=bf16))
        feats_for_memory = [None] * N
        feats_for_memory[0] = [rearrange(fused_0, 'b c h w -> (h w) b c')]
        self.predictor.memory_attention.reset_state(feat_2d_seq=feat_2d_seq, feat_3d_seq=feat_3d_seq,
                                                    explicit_3d_seq=explicit_3d_seq, B=1, device=device, dtype=bf16,
                                                    feats_for_memory_out=feats_for_memory)
        self.processing_order, self.N = order, N
        self.mask_inputs = mask_inputs.to(device) if mask_inputs is not None else None
        self.point_inputs = point_inputs
        self.fused_0 = fused_0
        self.feats_for_memory = feats_for_memory
        self.vision_feats_all, self.vision_pos_all, self.feat_sizes = vision_feats_all, vision_pos_all, feat_sizes
        self.output_dict = {'cond_frame_outputs': {}, 'non_cond_frame_outputs': {}}
        self.current_idx = 0
        self.positive_memory = PositiveMemory(max(32, self.predictor.num_maskmem, self.predictor.max_obj_ptrs_in_encoder))

    @torch.no_grad()
    @torch.autocast(device_type='cuda', dtype=torch.bfloat16)
    def step(self):
        device = self.device
        t = self.current_idx
        orig_idx = self.processing_order[t]
        self.positive_memory.restore(self.output_dict['non_cond_frame_outputs'])
        vf = [v.to(device) for v in self.vision_feats_all[orig_idx]]
        vpe = [v.to(device) for v in self.vision_pos_all[orig_idx]]
        if t == 0:
            vf[-1] = rearrange(self.fused_0, 'b c h w -> (h w) b c').contiguous()
        current_out = self.predictor.track_step(
            frame_idx=t, is_init_cond_frame=(t == 0), current_vision_feats=vf, current_vision_pos_embeds=vpe,
            feat_sizes=self.feat_sizes,
            point_inputs=self.point_inputs if t == 0 else None,
            mask_inputs=self.mask_inputs if t == 0 else None,
            output_dict=self.output_dict, num_frames=self.N, track_in_reverse=False,
            run_mem_encoder=False, prev_sam_mask_logits=None)
        memory_high_res = current_out['pred_masks_high_res']
        if self.predictor.fill_hole_area > 0:
            current_out['pred_masks'] = fill_holes_in_mask_scores(current_out['pred_masks'], self.predictor.fill_hole_area)
        current_out['pred_masks_high_res'] = F.interpolate(current_out['pred_masks'],
                                                           size=(self.predictor.image_size, self.predictor.image_size),
                                                           mode='bilinear', align_corners=False)
        if t == 0:
            memory_high_res = current_out['pred_masks_high_res']
        mem_vf = list(vf)
        mem_vf[-1] = self.feats_for_memory[t][0]
        maskmem_features, maskmem_pos_enc = self.predictor._encode_new_memory(
            current_vision_feats=mem_vf, feat_sizes=self.feat_sizes, pred_masks_high_res=memory_high_res,
            object_score_logits=current_out['object_score_logits'], is_mask_from_pts=(t == 0))
        current_out['maskmem_features'] = maskmem_features.to(torch.bfloat16)
        current_out['maskmem_pos_enc'] = maskmem_pos_enc
        logits = current_out['pred_masks_high_res'].squeeze().float().cpu()
        bucket = 'cond_frame_outputs' if t == 0 else 'non_cond_frame_outputs'
        self.output_dict[bucket][t] = current_out
        self.feats_for_memory[t] = None
        noncond = self.output_dict['non_cond_frame_outputs']
        if t > 0:
            self.positive_memory.remember(t, noncond[t])
        for index in list(noncond):
            noncond[index].pop('pred_masks_high_res', None)
            if index < t - 32:
                del noncond[index]
        for values in self.predictor.memory_attention.history.values():
            values.clear()
        self.current_idx += 1
        return logits


@torch.no_grad()
def encode_sam2_frames(predictor, padded_frames, device='cuda', batch=2):
    """SAM2 image-encoder features for every frame (kept on the device)."""
    vision_feats_all, vision_pos_all, feat_sizes = [], None, None
    for start in range(0, len(padded_frames), batch):
        images = torch.stack([sam2_input(p) for p in padded_frames[start:start + batch]]).to(device)
        vf, vpe, fs = get_image_feature(predictor, images)
        for k in range(images.shape[0]):
            vision_feats_all.append([v[:, k:k + 1] for v in vf])
        if vision_pos_all is None:
            vision_pos_all = [v[:, 0:1] for v in vpe]
            feat_sizes = fs
    # positional embeddings are identical for every frame; expose them per frame like the evaluator
    return vision_feats_all, [vision_pos_all] * len(vision_feats_all), feat_sizes


@torch.no_grad()
def track(predictor, padded_frames, g, must3r_feats, explicit, start_idx, first_mask=None, points=None, labels=None,
          device='cuda', progress=None):
    """Bidirectional tracking from start_idx. Prompt: first_mask (bool, original resolution) or clicks.
    Returns {frame_idx: bool mask at original resolution}."""
    n = len(padded_frames)
    vision_feats_all, vision_pos_all, feat_sizes = encode_sam2_frames(predictor, padded_frames, device)
    feats_dev = [f.to(device) for f in must3r_feats]
    explicit_dev = explicit.to(device)
    mask_inputs = original_mask_to_sam2(first_mask, g, device) if first_mask is not None else None
    point_inputs = points_to_sam2(points, labels, g, device) if points else None
    tracker = Tracker3AM(predictor, device)
    masks, done = {}, 0
    orders = [list(range(start_idx, n))] + ([list(range(start_idx, -1, -1))] if start_idx > 0 else [])
    total = sum(map(len, orders))
    try:
        for order in orders:
            tracker.initialize(vision_feats_all, vision_pos_all, feat_sizes, feats_dev, explicit_dev, order,
                               mask_inputs=mask_inputs, point_inputs=point_inputs)
            for i in order:
                logits = tracker.step()
                if i not in masks:
                    masks[i] = logits_to_original_mask(logits, g)
                done += 1
                if progress is not None and (done % 10 == 0 or done == total):
                    progress(done, total)
            predictor.memory_attention.clear_state()
    finally:
        predictor.memory_attention.clear_state()
        del vision_feats_all, feats_dev, explicit_dev
        torch.cuda.empty_cache()
    return masks


@torch.no_grad()
@torch.autocast(device_type='cuda', dtype=torch.bfloat16)
def single_frame_mask(original_predictor, padded_frame, g, points, labels, device='cuda'):
    """Original SAM2.1 prompt-to-mask on one frame (points: original pixels, labels 1/0/2/3)."""
    image = sam2_input(padded_frame)[None].to(device)
    vf, vpe, fs = get_image_feature(original_predictor, image)
    prompts = points_to_sam2(points, labels, g, device)
    out = original_predictor.track_step(
        frame_idx=0, is_init_cond_frame=True, current_vision_feats=vf, current_vision_pos_embeds=vpe, feat_sizes=fs,
        point_inputs=prompts, mask_inputs=None, output_dict={'cond_frame_outputs': {}, 'non_cond_frame_outputs': {}},
        num_frames=1, track_in_reverse=False, run_mem_encoder=False, prev_sam_mask_logits=None)
    logits = F.interpolate(out['pred_masks'], size=(SAM2_IMAGE_SIZE, SAM2_IMAGE_SIZE), mode='bilinear', align_corners=False)
    return logits_to_original_mask(logits.squeeze().float().cpu(), g)