"""ORMBG (Open Remove Background Model, an IS-Net from xuebinqin/DIS) in MLX: a port of schirrmacher/ormbg's
`ormbg/models/ormbg.py` (Apache-2.0).

    import mlx_model
    model = mlx_model.load("fp16")                     # a folder with config.json + model.safetensors
    alpha = mlx_model.matte(model, rgb)                # rgb: HxWx3 uint8 -> HxW uint8 alpha (255 = foreground)

Runs on its own next to its weights (mlx and numpy). NHWC; each BatchNorm is folded into the conv before it (exact
in fp32, done at conversion), so the weights are conv kernels and biases only.
"""

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": "ormbg", "input_size": 1024}
PRECISIONS = ["bf16", "fp16"]       # all convolutions: MLX can't quantize them (model.toml)


def resize_bilinear(x, size):
    """torch F.interpolate(mode="bilinear", align_corners=False) for NHWC."""
    b, h, w, c = x.shape
    oh, ow = size
    if (oh, ow) == (h, w):
        return x

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

    y0, y1, fy = axis(h, oh)
    x0, x1, fx = axis(w, ow)
    fy = fy.astype(x.dtype)[None, :, None, None]
    fx = 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


def maxpool2(x):
    """MaxPool2d(2, 2, ceil_mode=True), NHWC."""
    b, h, w, c = x.shape
    ph, pw = h % 2, w % 2
    if ph or pw:
        x = mx.pad(x, [(0, 0), (0, ph), (0, pw), (0, 0)], constant_values=-float("inf"))
        h, w = h + ph, w + pw
    return x.reshape(b, h // 2, 2, w // 2, 2, c).max(axis=(2, 4))


class REBNCONV(nn.Module):
    def __init__(self, cin, cout, d=1):
        super().__init__()
        self.conv_s1 = nn.Conv2d(cin, cout, 3, padding=d, dilation=d)

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


class RSU(nn.Module):
    """RSU-L (L = depth: 7, 6, 5, 4): an encoder-decoder with pooling; dilation 2 at the bottom."""

    def __init__(self, depth, cin, mid, cout):
        super().__init__()
        self.depth = depth
        self.rebnconvin = REBNCONV(cin, cout)
        setattr(self, "rebnconv1", REBNCONV(cout, mid))
        for i in range(2, depth):
            setattr(self, f"rebnconv{i}", REBNCONV(mid, mid))
        setattr(self, f"rebnconv{depth}", REBNCONV(mid, mid, d=2))
        for i in range(depth - 1, 1, -1):
            setattr(self, f"rebnconv{i}d", REBNCONV(mid * 2, mid))
        self.rebnconv1d = REBNCONV(mid * 2, cout)

    def __call__(self, x):
        hxin = self.rebnconvin(x)
        enc, hx = [], hxin
        for i in range(1, self.depth):
            hx = getattr(self, f"rebnconv{i}")(hx)
            enc.append(hx)
            if i < self.depth - 1:
                hx = maxpool2(hx)
        hx = getattr(self, f"rebnconv{self.depth}")(hx)
        for i in range(self.depth - 1, 0, -1):
            hx = getattr(self, f"rebnconv{i}d")(mx.concatenate([hx, enc[i - 1]], -1))
            if i > 1:
                hx = resize_bilinear(hx, enc[i - 2].shape[1:3])
        return hx + hxin


class RSU4F(nn.Module):
    """RSU-4F: no pooling, dilations 1, 2, 4, 8."""

    def __init__(self, cin, mid, cout):
        super().__init__()
        self.rebnconvin = REBNCONV(cin, cout)
        self.rebnconv1 = REBNCONV(cout, mid)
        self.rebnconv2 = REBNCONV(mid, mid, 2)
        self.rebnconv3 = REBNCONV(mid, mid, 4)
        self.rebnconv4 = REBNCONV(mid, mid, 8)
        self.rebnconv3d = REBNCONV(mid * 2, mid, 4)
        self.rebnconv2d = REBNCONV(mid * 2, mid, 2)
        self.rebnconv1d = REBNCONV(mid * 2, cout)

    def __call__(self, x):
        hxin = self.rebnconvin(x)
        hx1 = self.rebnconv1(hxin)
        hx2 = self.rebnconv2(hx1)
        hx3 = self.rebnconv3(hx2)
        hx4 = self.rebnconv4(hx3)
        hx3d = self.rebnconv3d(mx.concatenate([hx4, hx3], -1))
        hx2d = self.rebnconv2d(mx.concatenate([hx3d, hx2], -1))
        return self.rebnconv1d(mx.concatenate([hx2d, hx1], -1)) + hxin


class Model(nn.Module):
    def __init__(self, config: dict = CONFIG):
        super().__init__()
        self.size = config["input_size"]
        self.conv_in = nn.Conv2d(3, 64, 3, stride=2, padding=1)
        self.stage1, self.stage2 = RSU(7, 64, 32, 64), RSU(6, 64, 32, 128)
        self.stage3, self.stage4 = RSU(5, 128, 64, 256), RSU(4, 256, 128, 512)
        self.stage5, self.stage6 = RSU4F(512, 256, 512), RSU4F(512, 256, 512)
        self.stage5d, self.stage4d = RSU4F(1024, 256, 512), RSU(4, 1024, 128, 256)
        self.stage3d, self.stage2d, self.stage1d = RSU(5, 512, 64, 128), RSU(6, 256, 32, 64), RSU(7, 128, 16, 64)
        self.side1 = nn.Conv2d(64, 1, 3, padding=1)

    def __call__(self, x):
        """x: (B, S, S, 3) in [0, 1] -> (B, S, S, 1) foreground probability (side output 1, the one used)."""
        hxin = self.conv_in(x)
        hx1 = self.stage1(hxin)
        hx2 = self.stage2(maxpool2(hx1))
        hx3 = self.stage3(maxpool2(hx2))
        hx4 = self.stage4(maxpool2(hx3))
        hx5 = self.stage5(maxpool2(hx4))
        hx6 = self.stage6(maxpool2(hx5))
        up = lambda a, b: resize_bilinear(a, b.shape[1:3])
        hx5d = self.stage5d(mx.concatenate([up(hx6, hx5), hx5], -1))
        hx4d = self.stage4d(mx.concatenate([up(hx5d, hx4), hx4], -1))
        hx3d = self.stage3d(mx.concatenate([up(hx4d, hx3), hx3], -1))
        hx2d = self.stage2d(mx.concatenate([up(hx3d, hx2), hx2], -1))
        hx1d = self.stage1d(mx.concatenate([up(hx2d, hx1), hx1], -1))
        return mx.sigmoid(resize_bilinear(self.side1(hx1d), x.shape[1:3]))


def load(folder) -> Model:
    folder = Path(folder)
    model = Model(json.loads((folder / "config.json").read_text()))
    model.load_weights(str(folder / "model.safetensors"), strict=True)
    mx.eval(model.parameters())
    return model


def predict(model: Model, rgb: np.ndarray) -> np.ndarray:
    """The raw model output at the original size, float32 HxW in [0, 1] (as ormbg/inference.py, before min-max)."""
    dt = model.conv_in.weight.dtype
    h, w, _ = rgb.shape
    x = mx.array(rgb.astype(np.float32))[None]
    x = mx.floor(resize_bilinear(x, (model.size, model.size))) / 255.0      # the original rounds to uint8 first
    y = model(x.astype(dt)).astype(mx.float32)
    return np.array(resize_bilinear(y, (h, w)))[0, :, :, 0]


def matte(model: Model, rgb: np.ndarray) -> np.ndarray:
    """HxWx3 uint8 -> HxW uint8 alpha, min-max normalized like ormbg/inference.py."""
    p = predict(model, rgb)
    p = (p - p.min()) / max(p.max() - p.min(), 1e-8)
    return (p * 255).astype(np.uint8)


def convert(source: Path) -> dict:
    """runs/source/models/ormbg.pth -> conv weights with each BatchNorm folded in (fp32), MLX layout."""
    import torch
    sd = {k: v.double() for k, v in torch.load(Path(source) / "models" / "ormbg.pth", map_location="cpu",
                                                  weights_only=True).items()}
    out = {}
    for k, w in sd.items():
        if not k.endswith(".weight") or w.ndim != 4:
            continue
        base = k[:-len(".weight")]
        b = sd.get(base + ".bias", torch.zeros(w.shape[0], dtype=torch.float64))
        bn = base.replace("conv_s1", "bn_s1")
        if bn != base and bn + ".running_mean" in sd:
            g = sd[bn + ".weight"] / torch.sqrt(sd[bn + ".running_var"] + 1e-5)
            w = w * g[:, None, None, None]
            b = (b - sd[bn + ".running_mean"]) * g + sd[bn + ".bias"]
        if base.startswith("side") and base != "side1":
            continue                                    # side outputs 2-6 only matter for training
        out[base + ".weight"] = w.float().numpy().transpose(0, 2, 3, 1)
        out[base + ".bias"] = b.float().numpy()
    return out