"""Real-ESRGAN x4plus (RRDBNet) in MLX: a port of basicsr's `RRDBNet` (xinntao/Real-ESRGAN, BSD-3-Clause).

    import mlx_model
    model = mlx_model.load("bf16")                       # a folder with config.json + model.safetensors
    out = mlx_model.upscale(model, rgb)                  # rgb: HxWx3 uint8 numpy -> (4H)x(4W)x3 uint8

Runs on its own next to its weights (only mlx and numpy). NHWC throughout; the weights are the original
`params_ema` with each conv kernel transposed to MLX's (out, kh, kw, in).
"""

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": "rrdbnet", "num_in_ch": 3, "num_out_ch": 3, "num_feat": 64, "num_block": 23,
          "num_grow_ch": 32, "scale": 4}
# Every layer is a convolution, which MLX can't quantize: only bf16 / fp16 (lab/zoo/real-esrgan/model.toml).
PRECISIONS = ["bf16", "fp16"]


def lrelu(x):
    return nn.leaky_relu(x, 0.2)


class RDB(nn.Module):
    def __init__(self, f: int, g: int):
        super().__init__()
        self.conv1 = nn.Conv2d(f, g, 3, padding=1)
        self.conv2 = nn.Conv2d(f + g, g, 3, padding=1)
        self.conv3 = nn.Conv2d(f + 2 * g, g, 3, padding=1)
        self.conv4 = nn.Conv2d(f + 3 * g, g, 3, padding=1)
        self.conv5 = nn.Conv2d(f + 4 * g, f, 3, padding=1)

    def __call__(self, x):
        x1 = lrelu(self.conv1(x))
        x2 = lrelu(self.conv2(mx.concatenate([x, x1], -1)))
        x3 = lrelu(self.conv3(mx.concatenate([x, x1, x2], -1)))
        x4 = lrelu(self.conv4(mx.concatenate([x, x1, x2, x3], -1)))
        x5 = self.conv5(mx.concatenate([x, x1, x2, x3, x4], -1))
        return x5 * 0.2 + x


class RRDB(nn.Module):
    def __init__(self, f: int, g: int):
        super().__init__()
        self.rdb1, self.rdb2, self.rdb3 = RDB(f, g), RDB(f, g), RDB(f, g)

    def __call__(self, x):
        return self.rdb3(self.rdb2(self.rdb1(x))) * 0.2 + x


def up2(x):
    """Nearest-neighbour x2 (NHWC)."""
    b, h, w, c = x.shape
    x = mx.broadcast_to(x[:, :, None, :, None, :], (b, h, 2, w, 2, c))
    return x.reshape(b, h * 2, w * 2, c)


class Model(nn.Module):
    def __init__(self, config: dict = CONFIG):
        super().__init__()
        f, g, cin, cout = config["num_feat"], config["num_grow_ch"], config["num_in_ch"], config["num_out_ch"]
        self.scale = config["scale"]
        assert self.scale == 4, "x4plus only (x2 needs basicsr's pixel_unshuffle input)"
        self.conv_first = nn.Conv2d(cin, f, 3, padding=1)
        self.body = [RRDB(f, g) for _ in range(config["num_block"])]
        self.conv_body = nn.Conv2d(f, f, 3, padding=1)
        self.conv_up1 = nn.Conv2d(f, f, 3, padding=1)
        self.conv_up2 = nn.Conv2d(f, f, 3, padding=1)
        self.conv_hr = nn.Conv2d(f, f, 3, padding=1)
        self.conv_last = nn.Conv2d(f, cout, 3, padding=1)

    def __call__(self, x):
        """x: (B, H, W, 3) in [0, 1] -> (B, 4H, 4W, 3)."""
        feat = self.conv_first(x)
        body = feat
        for blk in self.body:
            body = blk(body)
        feat = feat + self.conv_body(body)
        feat = lrelu(self.conv_up1(up2(feat)))
        feat = lrelu(self.conv_up2(up2(feat)))
        return self.conv_last(lrelu(self.conv_hr(feat)))


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


def upscale(model: Model, rgb: np.ndarray, tile: int = 256, pad: int = 16) -> np.ndarray:
    """HxWx3 uint8 -> (4H)x(4W)x3 uint8, in tiles of `tile` input pixels (overlap `pad`) to bound memory."""
    dt = model.conv_first.weight.dtype
    x = rgb.astype(np.float32) / 255.0
    h, w, _ = x.shape
    s = model.scale
    out = np.zeros((h * s, w * s, 3), np.float32)
    for y0 in range(0, h, tile):
        for x0 in range(0, w, tile):
            y1, x1 = min(y0 + tile, h), min(x0 + tile, w)
            ya, xa, yb, xb = max(y0 - pad, 0), max(x0 - pad, 0), min(y1 + pad, h), min(x1 + pad, w)
            t = mx.array(x[None, ya:yb, xa:xb]).astype(dt)
            o = np.array(model(t).astype(mx.float32))[0]
            out[y0 * s:y1 * s, x0 * s:x1 * s] = o[(y0 - ya) * s:(y1 - ya) * s, (x0 - xa) * s:(x1 - xa) * s]
    return (np.clip(out, 0, 1) * 255.0).round().astype(np.uint8)


def convert(source: Path) -> dict:
    """runs/source/RealESRGAN_x4plus.pth (params_ema) -> MLX names and layouts, float32."""
    import torch
    sd = torch.load(Path(source) / "RealESRGAN_x4plus.pth", map_location="cpu", weights_only=True)["params_ema"]
    out = {}
    for k, v in sd.items():
        a = v.float().numpy()
        out[k] = a.transpose(0, 2, 3, 1) if a.ndim == 4 else a
    return out