"""NAFNet (megvii-research, MIT) rebuilt in PyTorch to load the mlx-community NHWC safetensors."""
import json
import numpy as np
import torch
import torch.nn as nn
from safetensors.numpy import load_file


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

    def forward(self, x):
        mu = x.mean(1, keepdim=True)
        var = (x - mu).pow(2).mean(1, keepdim=True)
        y = (x - mu) / torch.sqrt(var + self.eps)
        return self.weight.view(1, -1, 1, 1) * y + self.bias.view(1, -1, 1, 1)


class SimpleGate(nn.Module):
    def forward(self, x):
        a, b = x.chunk(2, dim=1)
        return a * b


class NAFBlock(nn.Module):
    def __init__(self, c):
        super().__init__()
        dw = c * 2
        self.conv1 = nn.Conv2d(c, dw, 1)
        self.conv2 = nn.Conv2d(dw, dw, 3, padding=1, groups=dw)
        self.conv3 = nn.Conv2d(dw // 2, c, 1)
        # Simplified channel attention; global average pooling (the tile size equals the TLC window).
        self.sca = nn.Conv2d(dw // 2, dw // 2, 1)
        self.sg = SimpleGate()
        self.conv4 = nn.Conv2d(c, 2 * c, 1)
        self.conv5 = nn.Conv2d(c, c, 1)
        self.norm1 = LayerNorm2d(c)
        self.norm2 = LayerNorm2d(c)
        self.beta = nn.Parameter(torch.zeros(1, c, 1, 1))
        self.gamma = nn.Parameter(torch.zeros(1, c, 1, 1))

    def forward(self, inp):
        x = self.sg(self.conv2(self.conv1(self.norm1(inp))))
        x = x * self.sca(x.mean(dim=(2, 3), keepdim=True))
        y = inp + self.conv3(x) * self.beta
        x = self.conv5(self.sg(self.conv4(self.norm2(y))))
        return y + x * self.gamma


class NAFNet(nn.Module):
    def __init__(self, img_channel=3, width=64, middle_blk_num=1, enc_blk_nums=(1, 1, 1, 28), dec_blk_nums=(1, 1, 1, 1), **_):
        super().__init__()
        self.intro = nn.Conv2d(img_channel, width, 3, padding=1)
        self.ending = nn.Conv2d(width, img_channel, 3, padding=1)
        self.encoders, self.decoders, self.ups, self.downs = nn.ModuleList(), nn.ModuleList(), nn.ModuleList(), nn.ModuleList()
        chan = width
        for num in enc_blk_nums:
            self.encoders.append(nn.Sequential(*[NAFBlock(chan) for _ in range(num)]))
            self.downs.append(nn.Conv2d(chan, 2 * chan, 2, 2))
            chan *= 2
        self.middle_blks = nn.Sequential(*[NAFBlock(chan) for _ in range(middle_blk_num)])
        for num in dec_blk_nums:
            self.ups.append(nn.Sequential(nn.Conv2d(chan, chan * 2, 1, bias=False), nn.PixelShuffle(2)))
            chan //= 2
            self.decoders.append(nn.Sequential(*[NAFBlock(chan) for _ in range(num)]))

    def forward(self, inp):
        x = self.intro(inp)
        skips = []
        for enc, down in zip(self.encoders, self.downs):
            x = enc(x)
            skips.append(x)
            x = down(x)
        x = self.middle_blks(x)
        for dec, up, skip in zip(self.decoders, self.ups, skips[::-1]):
            x = dec(up(x) + skip)
        return self.ending(x) + inp


def load(weights, config=None):
    cfg = config or {"width": 64, "middle_blk_num": 1, "enc_blk_nums": [1, 1, 1, 28], "dec_blk_nums": [1, 1, 1, 1]}
    model = NAFNet(**cfg)
    raw = load_file(weights)
    state = {}
    target = model.state_dict()
    for key, value in raw.items():
        name = key
        if key.startswith("ups.") and key.count(".") == 2:
            name = key.replace(".weight", ".0.weight")  # ups.N is Sequential(conv, PixelShuffle)
        if value.ndim == 4 and name.endswith("weight"):
            value = np.transpose(value, (0, 3, 1, 2))  # MLX (O,kH,kW,I) -> PyTorch (O,I,kH,kW)
        elif name.endswith(("beta", "gamma")):
            value = np.transpose(value, (0, 3, 1, 2))
        state[name] = torch.from_numpy(np.ascontiguousarray(value))
    missing = set(target) - set(state)
    unused = set(state) - set(target)
    bad = [k for k in target if k in state and tuple(target[k].shape) != tuple(state[k].shape)]
    if missing or unused or bad:
        raise RuntimeError(f"missing={sorted(missing)[:8]} unused={sorted(unused)[:8]} shape={bad[:8]}")
    model.load_state_dict(state, strict=True)
    return model.eval()