"""MobileSAM (ChaoningZhang/MobileSAM, Apache-2.0) in MLX: a port of `mobile_sam` (`build_sam_vit_t`: TinyViT-5M
image encoder + SAM's prompt encoder and mask decoder), the segment-anything half of MobileSAM v2.

    import mlx_model
    model = mlx_model.load("bf16")                        # a folder with config.json + model.safetensors
    emb = mlx_model.embed(model, rgb)                     # rgb HxWx3 uint8 -> image embedding (+ sizes)
    masks, scores, logits = mlx_model.predict(model, emb, points=[[x, y]], labels=[1], box=None, multimask=True)

`predict` follows SamPredictor.predict: point / box prompts in pixel coordinates of the original image, masks at
the original size (logits > 0), the IoU head's scores, and 256x256 low-res logits. Runs on its own next to its
weights (mlx, numpy, pillow). NHWC; every BatchNorm folded into its conv at conversion.
"""

from __future__ import annotations

import json
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
import numpy as np

CONFIG = {"model_type": "mobile_sam", "image_size": 1024, "embed_dims": [64, 128, 160, 320], "depths": [2, 2, 6, 2],
          "num_heads": [2, 4, 5, 10], "window_sizes": [7, 7, 14, 7], "prompt_dim": 256, "mask_in_chans": 16}
PRECISIONS = ["bf16", "fp16", "8bit", "6bit", "4bit"]
PIXEL_MEAN = np.array([123.675, 116.28, 103.53], np.float32)
PIXEL_STD = np.array([58.395, 57.12, 57.375], np.float32)
GROUP = 64


def quantizable(path: str, module) -> bool:
    return isinstance(module, nn.Linear) and module.weight.shape[-1] % GROUP == 0


def resize(x, size):
    """torch F.interpolate(mode="bilinear", align_corners=False), NHWC."""
    b, h, w, c = x.shape
    oh, ow = int(size[0]), int(size[1])
    if (oh, ow) == (h, w):
        return x

    def axis(n_in, n_out):
        s = np.clip((np.arange(n_out, dtype=np.float64) + 0.5) * (n_in / n_out) - 0.5, 0, None)
        i0 = np.minimum(np.floor(s).astype(np.int32), n_in - 1)
        return mx.array(i0), mx.array(np.minimum(i0 + 1, n_in - 1)), mx.array((s - i0).astype(np.float32))

    y0, y1, fy = axis(h, oh)
    x0, x1, fx = axis(w, ow)
    fy, fx = fy.astype(x.dtype)[None, :, None, None], fx.astype(x.dtype)[None, None, :, None]
    top, bot = x[:, y0], x[:, y1]
    top = top[:, :, x0] * (1 - fx) + top[:, :, x1] * fx
    bot = bot[:, :, x0] * (1 - fx) + bot[:, :, x1] * fx
    return top * (1 - fy) + bot * fy


class LayerNorm2d(nn.Module):
    """SAM's channel LayerNorm (eps 1e-6), statistics in float32."""

    def __init__(self, c, eps=1e-6):
        super().__init__()
        self.weight, self.bias, self.eps = mx.ones((c,)), mx.zeros((c,)), eps

    def __call__(self, x):
        y = x.astype(mx.float32)
        u = y.mean(-1, keepdims=True)
        s = ((y - u) ** 2).mean(-1, keepdims=True)
        return (self.weight * ((y - u) * mx.rsqrt(s + self.eps)).astype(x.dtype) + self.bias)


# ------------------------------------------------------------------------------------------------ TinyViT

def conv(cin, cout, k=1, s=1, p=0, g=1):
    return nn.Conv2d(cin, cout, k, stride=s, padding=p, groups=g)          # Conv2d_BN, BatchNorm folded


class MBConv(nn.Module):
    def __init__(self, c, ratio=4.0):
        super().__init__()
        h = int(c * ratio)
        self.conv1, self.conv2, self.conv3 = conv(c, h), conv(h, h, 3, 1, 1, h), conv(h, c)

    def __call__(self, x):
        y = nn.gelu(self.conv1(x))
        y = nn.gelu(self.conv2(y))
        return nn.gelu(self.conv3(y) + x)


class PatchMerging(nn.Module):
    def __init__(self, dim, out):
        super().__init__()
        s = 1 if out in (320, 448, 576) else 2
        self.conv1, self.conv2, self.conv3 = conv(dim, out), conv(out, out, 3, s, 1, out), conv(out, out)

    def __call__(self, x):
        return self.conv3(nn.gelu(self.conv2(nn.gelu(self.conv1(x)))))


class Attention(nn.Module):
    def __init__(self, dim, key_dim, heads, res):
        super().__init__()
        self.heads, self.key_dim, self.scale = heads, key_dim, key_dim ** -0.5
        self.norm = nn.LayerNorm(dim)
        self.qkv = nn.Linear(dim, 3 * key_dim * heads)
        self.proj = nn.Linear(key_dim * heads, dim)
        pts = [(i, j) for i in range(res) for j in range(res)]
        offs, idx = {}, []
        for p1 in pts:
            for p2 in pts:
                o = (abs(p1[0] - p2[0]), abs(p1[1] - p2[1]))
                offs.setdefault(o, len(offs))
                idx.append(offs[o])
        self.attention_biases = mx.zeros((heads, len(offs)))
        self._idx = mx.array(np.array(idx, np.int32).reshape(len(pts), len(pts)))

    def __call__(self, x):
        b, n, _ = x.shape
        qkv = self.qkv(self.norm(x)).reshape(b, n, self.heads, 3 * self.key_dim)
        q, k, v = mx.split(qkv, [self.key_dim, 2 * self.key_dim], axis=-1)
        q, k, v = (t.transpose(0, 2, 1, 3) for t in (q, k, v))
        a = (q @ k.transpose(0, 1, 3, 2)) * self.scale + self.attention_biases[:, self._idx].astype(q.dtype)
        a = mx.softmax(a, axis=-1, precise=True)
        return self.proj((a @ v).transpose(0, 2, 1, 3).reshape(b, n, -1))


class Mlp(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        self.norm, self.fc1, self.fc2 = nn.LayerNorm(dim), nn.Linear(dim, hidden), nn.Linear(hidden, dim)

    def __call__(self, x):
        return self.fc2(nn.gelu(self.fc1(self.norm(x))))


class TinyViTBlock(nn.Module):
    def __init__(self, dim, heads, ws):
        super().__init__()
        self.ws = ws
        self.attn = Attention(dim, dim // heads, heads, ws)
        self.mlp = Mlp(dim, dim * 4)
        self.local_conv = conv(dim, dim, 3, 1, 1, dim)

    def __call__(self, x):
        """x: (B, H, W, C)."""
        b, h, w, c = x.shape
        ws = self.ws
        if h == ws and w == ws:
            y = self.attn(x.reshape(b, h * w, c)).reshape(b, h, w, c)
        else:
            pb, pr = (ws - h % ws) % ws, (ws - w % ws) % ws
            y = mx.pad(x, [(0, 0), (0, pb), (0, pr), (0, 0)]) if pb or pr else x
            ph, pw = h + pb, w + pr
            y = y.reshape(b, ph // ws, ws, pw // ws, ws, c).transpose(0, 1, 3, 2, 4, 5).reshape(-1, ws * ws, c)
            y = self.attn(y)
            y = y.reshape(b, ph // ws, pw // ws, ws, ws, c).transpose(0, 1, 3, 2, 4, 5).reshape(b, ph, pw, c)
            y = y[:, :h, :w]
        x = self.local_conv(x + y)
        return x + self.mlp(x.reshape(b, h * w, c)).reshape(b, h, w, c)


class Layer(nn.Module):
    def __init__(self, i, dims, depth, heads, ws):
        super().__init__()
        d = dims[i]
        self.blocks = [MBConv(d) for _ in range(depth)] if i == 0 else [TinyViTBlock(d, heads, ws) for _ in range(depth)]
        if i < len(dims) - 1:
            self.downsample = PatchMerging(d, dims[i + 1])

    def __call__(self, x):
        for b in self.blocks:
            x = b(x)
        return self.downsample(x) if "downsample" in self else x


class TinyViT(nn.Module):
    def __init__(self, c):
        super().__init__()
        d = c["embed_dims"]
        self.patch_embed = [conv(3, d[0] // 2, 3, 2, 1), conv(d[0] // 2, d[0], 3, 2, 1)]
        self.layers = [Layer(i, d, c["depths"][i], c["num_heads"][i], c["window_sizes"][i]) for i in range(4)]
        self.neck_conv1 = nn.Conv2d(d[-1], 256, 1, bias=False)
        self.neck_norm1 = LayerNorm2d(256)
        self.neck_conv2 = nn.Conv2d(256, 256, 3, padding=1, bias=False)
        self.neck_norm2 = LayerNorm2d(256)

    def __call__(self, x):
        x = self.patch_embed[1](nn.gelu(self.patch_embed[0](x)))
        for layer in self.layers:
            x = layer(x)
        return self.neck_norm2(self.neck_conv2(self.neck_norm1(self.neck_conv1(x))))


# ------------------------------------------------------------------------------------------------ SAM decoder

class SAMAttention(nn.Module):
    def __init__(self, dim, heads, down=1):
        super().__init__()
        inner = dim // down
        self.heads = heads
        self.q_proj, self.k_proj, self.v_proj = nn.Linear(dim, inner), nn.Linear(dim, inner), nn.Linear(dim, inner)
        self.out_proj = nn.Linear(inner, dim)

    def __call__(self, q, k, v):
        q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v)
        b, nq, c = q.shape
        sp = lambda t: t.reshape(b, t.shape[1], self.heads, c // self.heads).transpose(0, 2, 1, 3)
        q, k, v = sp(q), sp(k), sp(v)
        a = mx.softmax((q @ k.transpose(0, 1, 3, 2)) * (c // self.heads) ** -0.5, axis=-1, precise=True)
        return self.out_proj((a @ v).transpose(0, 2, 1, 3).reshape(b, nq, c))


class MLPBlock(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        self.lin1, self.lin2 = nn.Linear(dim, hidden), nn.Linear(hidden, dim)

    def __call__(self, x):
        return self.lin2(nn.relu(self.lin1(x)))


class TwoWayBlock(nn.Module):
    def __init__(self, dim, heads, mlp_dim, skip_pe):
        super().__init__()
        self.skip_pe = skip_pe
        self.self_attn = SAMAttention(dim, heads)
        self.cross_attn_token_to_image = SAMAttention(dim, heads, 2)
        self.cross_attn_image_to_token = SAMAttention(dim, heads, 2)
        self.mlp = MLPBlock(dim, mlp_dim)
        self.norm1, self.norm2, self.norm3, self.norm4 = (nn.LayerNorm(dim) for _ in range(4))

    def __call__(self, q, k, qpe, kpe):
        if self.skip_pe:
            q = self.self_attn(q, q, q)
        else:
            t = q + qpe
            q = q + self.self_attn(t, t, q)
        q = self.norm1(q)
        q = self.norm2(q + self.cross_attn_token_to_image(q + qpe, k + kpe, k))
        q = self.norm3(q + self.mlp(q))
        k = self.norm4(k + self.cross_attn_image_to_token(k + kpe, q + qpe, q))
        return q, k


class TwoWayTransformer(nn.Module):
    def __init__(self, dim=256, heads=8, mlp_dim=2048, depth=2):
        super().__init__()
        self.layers = [TwoWayBlock(dim, heads, mlp_dim, i == 0) for i in range(depth)]
        self.final_attn_token_to_image = SAMAttention(dim, heads, 2)
        self.norm_final_attn = nn.LayerNorm(dim)

    def __call__(self, img, pe, tokens):
        b, h, w, c = img.shape
        k, kpe, q = img.reshape(b, h * w, c), pe.reshape(b, h * w, c), tokens
        for layer in self.layers:
            q, k = layer(q, k, tokens, kpe)
        q = self.norm_final_attn(q + self.final_attn_token_to_image(q + tokens, k + kpe, k))
        return q, k


class MLP(nn.Module):
    def __init__(self, din, dh, dout, n):
        super().__init__()
        dims = [din] + [dh] * (n - 1)
        self.layers = [nn.Linear(a, b) for a, b in zip(dims, [dh] * (n - 1) + [dout])]

    def __call__(self, x):
        for i, l in enumerate(self.layers):
            x = nn.relu(l(x)) if i < len(self.layers) - 1 else l(x)
        return x


class PromptEncoder(nn.Module):
    def __init__(self, dim, emb_hw, img_size, mask_in):
        super().__init__()
        self.emb_hw, self.img_size = emb_hw, img_size
        self.positional_encoding_gaussian_matrix = mx.zeros((2, dim // 2))
        self.point_embeddings = [mx.zeros((1, dim)) for _ in range(4)]
        self.not_a_point_embed = mx.zeros((1, dim))
        self.no_mask_embed = mx.zeros((1, dim))
        self.mask_conv1 = nn.Conv2d(1, mask_in // 4, 2, stride=2)
        self.mask_norm1 = LayerNorm2d(mask_in // 4)
        self.mask_conv2 = nn.Conv2d(mask_in // 4, mask_in, 2, stride=2)
        self.mask_norm2 = LayerNorm2d(mask_in)
        self.mask_conv3 = nn.Conv2d(mask_in, dim, 1)

    def pe(self, coords):                                   # coords in [0, 1], float32
        c = (2 * coords - 1) @ self.positional_encoding_gaussian_matrix.astype(mx.float32)
        c = 2 * np.pi * c
        return mx.concatenate([mx.sin(c), mx.cos(c)], -1)

    def dense_pe(self):
        h, w = self.emb_hw
        ys = (mx.arange(h, dtype=mx.float32) + 0.5) / h
        xs = (mx.arange(w, dtype=mx.float32) + 0.5) / w
        grid = mx.stack([mx.broadcast_to(xs[None, :], (h, w)), mx.broadcast_to(ys[:, None], (h, w))], -1)
        return self.pe(grid)[None]                          # (1, h, w, C)

    def __call__(self, points, labels, box, mask):
        dt = self.no_mask_embed.dtype
        sparse = []
        s = mx.array([self.img_size, self.img_size], dtype=mx.float32)
        if points is not None:
            pts, lab = points + 0.5, labels
            if box is None:
                pts = mx.concatenate([pts, mx.zeros((pts.shape[0], 1, 2))], 1)
                lab = mx.concatenate([lab, -mx.ones((lab.shape[0], 1))], 1)
            e = self.pe(pts / s)
            l = lab[..., None]
            e = mx.where(l == -1, 0.0, e)
            e = e + (l == -1) * self.not_a_point_embed.astype(mx.float32)
            e = e + (l == 0) * self.point_embeddings[0].astype(mx.float32)
            e = e + (l == 1) * self.point_embeddings[1].astype(mx.float32)
            sparse.append(e)
        if box is not None:
            e = self.pe((box + 0.5).reshape(-1, 2, 2) / s)
            e = e + mx.concatenate([self.point_embeddings[2], self.point_embeddings[3]], 0).astype(mx.float32)[None]
            sparse.append(e)
        sparse = mx.concatenate(sparse, 1).astype(dt)
        if mask is not None:
            m = nn.gelu(self.mask_norm1(self.mask_conv1(mask)))
            dense = self.mask_conv3(nn.gelu(self.mask_norm2(self.mask_conv2(m))))
        else:
            h, w = self.emb_hw
            dense = mx.broadcast_to(self.no_mask_embed.reshape(1, 1, 1, -1), (sparse.shape[0], h, w, self.no_mask_embed.shape[-1]))
        return sparse, dense


class MaskDecoder(nn.Module):
    def __init__(self, dim=256, n_mask=3):
        super().__init__()
        self.n = n_mask + 1
        self.transformer = TwoWayTransformer(dim)
        self.iou_token = mx.zeros((1, dim))
        self.mask_tokens = mx.zeros((self.n, dim))
        self.up1 = nn.ConvTranspose2d(dim, dim // 4, 2, stride=2)
        self.up_norm = LayerNorm2d(dim // 4)
        self.up2 = nn.ConvTranspose2d(dim // 4, dim // 8, 2, stride=2)
        self.output_hypernetworks_mlps = [MLP(dim, dim, dim // 8, 3) for _ in range(self.n)]
        self.iou_prediction_head = MLP(dim, 256, self.n, 3)

    def __call__(self, img, pe, sparse, dense):
        b = sparse.shape[0]
        out = mx.concatenate([self.iou_token, self.mask_tokens], 0)
        tokens = mx.concatenate([mx.broadcast_to(out[None], (b,) + out.shape), sparse], 1)
        src = mx.broadcast_to(img, (b,) + img.shape[1:]) + dense
        pos = mx.broadcast_to(pe, (b,) + pe.shape[1:]).astype(src.dtype)
        hs, src = self.transformer(src, pos, tokens)
        _, h, w, c = img.shape
        up = nn.gelu(self.up2(nn.gelu(self.up_norm(self.up1(src.reshape(b, h, w, c))))))
        hyper = mx.stack([self.output_hypernetworks_mlps[i](hs[:, 1 + i]) for i in range(self.n)], 1)
        bb, uh, uw, uc = up.shape
        masks = (hyper @ up.reshape(b, uh * uw, uc).transpose(0, 2, 1)).reshape(b, self.n, uh, uw)
        return masks, self.iou_prediction_head(hs[:, 0])


class Model(nn.Module):
    def __init__(self, config: dict = CONFIG):
        super().__init__()
        self.size = config["image_size"]
        e = self.size // 16
        self.image_encoder = TinyViT(config)
        self.prompt_encoder = PromptEncoder(config["prompt_dim"], (e, e), self.size, config["mask_in_chans"])
        self.mask_decoder = MaskDecoder(config["prompt_dim"])


# ------------------------------------------------------------------------------------------------ predictor

def load(folder) -> Model:
    folder = Path(folder)
    config = json.loads((folder / "config.json").read_text())
    model = Model(config)
    q = config.get("quantization")
    if q:
        nn.quantize(model, group_size=q["group_size"], bits=q["bits"], class_predicate=quantizable)
    model.load_weights(str(folder / "model.safetensors"), strict=True)
    mx.eval(model.parameters())
    return model


def dtype(model):
    return model.mask_decoder.iou_token.dtype


def embed(model: Model, rgb: np.ndarray) -> dict:
    """SamPredictor.set_image: longest side to 1024 (PIL bilinear), normalize, pad bottom-right, encode."""
    from PIL import Image
    h, w = rgb.shape[:2]
    s = model.size / max(h, w)
    nh, nw = int(h * s + 0.5), int(w * s + 0.5)
    x = np.asarray(Image.fromarray(rgb).resize((nw, nh), Image.BILINEAR), np.float32)
    x = (x - PIXEL_MEAN) / PIXEL_STD
    x = np.pad(x, ((0, model.size - nh), (0, model.size - nw), (0, 0)))
    feats = model.image_encoder(mx.array(x[None]).astype(dtype(model)))
    mx.eval(feats)
    return {"features": feats, "original_size": (h, w), "input_size": (nh, nw)}


def predict(model: Model, emb: dict, points=None, labels=None, box=None, multimask=True):
    """SamPredictor.predict: prompts in original-image pixels -> (masks bool (N, H, W), scores (N,), low-res
    logits (N, 256, 256))."""
    h, w = emb["original_size"]
    nh, nw = emb["input_size"]
    sx, sy = nw / w, nh / h
    pts = lab = bx = None
    if points is not None:
        pts = mx.array(np.asarray(points, np.float32)[None] * np.array([sx, sy], np.float32))
        lab = mx.array(np.asarray(labels, np.float32)[None])
    if box is not None:
        bx = mx.array(np.asarray(box, np.float32).reshape(1, 4) * np.array([sx, sy, sx, sy], np.float32))
    sparse, dense = model.prompt_encoder(pts, lab, bx, None)
    pe = model.prompt_encoder.dense_pe()
    logits, iou = model.mask_decoder(emb["features"], pe, sparse, dense.astype(sparse.dtype))
    sl = slice(1, None) if multimask else slice(0, 1)
    logits, iou = logits[:, sl].astype(mx.float32), iou[:, sl].astype(mx.float32)
    low = logits
    m = resize(logits.transpose(0, 2, 3, 1), (model.size, model.size))[:, :nh, :nw]
    m = resize(m, (h, w)).transpose(0, 3, 1, 2)
    return np.array(m[0] > 0), np.array(iou[0]), np.array(low[0]), np.array(m[0])


# ------------------------------------------------------------------------------------------------ conversion

def _fold(sd, pre):
    """Conv2d_BN (`pre`.c + `pre`.bn) -> one conv's (weight, bias) in MLX layout."""
    w = sd[pre + ".c.weight"].astype(np.float64)
    g = sd[pre + ".bn.weight"] / np.sqrt(sd[pre + ".bn.running_var"].astype(np.float64) + 1e-5)
    b = sd[pre + ".bn.bias"] - sd[pre + ".bn.running_mean"] * g
    return (w * g[:, None, None, None]).astype(np.float32).transpose(0, 2, 3, 1), b.astype(np.float32)


def convert(source: Path) -> dict:
    """runs/source/weights/mobile_sam.pt (a plain state dict) -> MLX names: BatchNorms folded, Sequentials renamed,
    the classification head (norm_head / head) dropped."""
    import torch
    sd = {k: v.float().numpy() for k, v in torch.load(Path(source) / "weights" / "mobile_sam.pt", map_location="cpu",
                                                         weights_only=True).items() if v.is_floating_point()}
    out = {}
    folded = sorted({k[:-len(".c.weight")] for k in sd if k.endswith(".c.weight")})
    for pre in folded:
        dst = pre.replace("patch_embed.seq.0", "patch_embed.0").replace("patch_embed.seq.2", "patch_embed.1")
        out[dst + ".weight"], out[dst + ".bias"] = _fold(sd, pre)
    ren = {"image_encoder.neck.0.": "image_encoder.neck_conv1.", "image_encoder.neck.1.": "image_encoder.neck_norm1.",
           "image_encoder.neck.2.": "image_encoder.neck_conv2.", "image_encoder.neck.3.": "image_encoder.neck_norm2.",
           "prompt_encoder.pe_layer.": "prompt_encoder.",
           "prompt_encoder.mask_downscaling.0.": "prompt_encoder.mask_conv1.",
           "prompt_encoder.mask_downscaling.1.": "prompt_encoder.mask_norm1.",
           "prompt_encoder.mask_downscaling.3.": "prompt_encoder.mask_conv2.",
           "prompt_encoder.mask_downscaling.4.": "prompt_encoder.mask_norm2.",
           "prompt_encoder.mask_downscaling.6.": "prompt_encoder.mask_conv3.",
           "mask_decoder.output_upscaling.0.": "mask_decoder.up1.",
           "mask_decoder.output_upscaling.1.": "mask_decoder.up_norm.",
           "mask_decoder.output_upscaling.3.": "mask_decoder.up2."}
    for k, v in sd.items():
        if any(k.startswith(p + ".") for p in folded) or ".head." in k or ".norm_head." in k:
            continue
        n = k
        for a, b in ren.items():
            if n.startswith(a):
                n = b + n[len(a):]
        if n.endswith(".weight") and (n.startswith("prompt_encoder.point_embeddings") or
                                      n.split(".")[-2] in ("not_a_point_embed", "no_mask_embed", "iou_token", "mask_tokens")):
            n = n[:-len(".weight")]
        if v.ndim == 4:
            v = v.transpose(1, 2, 3, 0) if ".up1." in n or ".up2." in n else v.transpose(0, 2, 3, 1)
        out[n] = v
    return out