"""Export official PyTorch super-resolution models to ONNX."""

import numpy as np


OPSET_VERSION = 17


def _require_build_dependencies():
    try:
        import onnx
        import onnxruntime
        import torch
    except ImportError as error:
        raise RuntimeError(
            "Model build dependencies are missing. Run with: "
            "uv sync"
        ) from error
    return onnx, onnxruntime, torch


def _load_weights(torch, checkpoint_path, key):
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    return checkpoint[key]


def _validate_export(onnx, onnxruntime, torch, model, sample, destination):
    onnx.checker.check_model(onnx.load(destination))
    expected = model(sample).detach().numpy()
    session = onnxruntime.InferenceSession(
        str(destination),
        providers=["CPUExecutionProvider"],
    )
    actual = session.run(None, {"input": sample.numpy()})[0]
    np.testing.assert_allclose(actual, expected, rtol=1e-4, atol=1e-5)


def _export(torch, model, sample, destination, dynamic_axes=None):
    model.eval()
    torch.onnx.export(
        model,
        sample,
        destination,
        export_params=True,
        opset_version=OPSET_VERSION,
        do_constant_folding=True,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes=dynamic_axes,
        dynamo=False,
    )


def _restore_realesrgan_scope_names(onnx, destination):
    model = onnx.load(destination)
    for node in model.graph.node:
        node.name = node.name.replace("/body.", "/body/body.")
        for values in (node.input, node.output):
            for index, value in enumerate(values):
                values[index] = value.replace("/body.", "/body/body.")
    onnx.save(model, destination)


def export_realesrgan_wdn_x4(checkpoint_path, destination):
    onnx, onnxruntime, torch = _require_build_dependencies()
    nn = torch.nn

    class SRVGGNetCompact(nn.Module):
        def __init__(self):
            super().__init__()
            self.upscale = 4
            self.body = nn.ModuleList([nn.Conv2d(3, 64, 3, 1, 1), nn.PReLU(64)])
            for _ in range(32):
                self.body.extend([nn.Conv2d(64, 64, 3, 1, 1), nn.PReLU(64)])
            self.body.append(nn.Conv2d(64, 3 * self.upscale**2, 3, 1, 1))
            self.upsampler = nn.PixelShuffle(self.upscale)

        def forward(self, value):
            result = value
            for layer in self.body:
                result = layer(result)
            result = self.upsampler(result)
            return result + torch.nn.functional.interpolate(
                value,
                scale_factor=self.upscale,
                mode="nearest",
            )

    model = SRVGGNetCompact()
    model.load_state_dict(_load_weights(torch, checkpoint_path, "params"), strict=True)
    sample = torch.rand(1, 3, 16, 16)
    dynamic_axes = {
        "input": {2: "height", 3: "width"},
        "output": {2: "height4", 3: "width4"},
    }
    _export(torch, model, sample, destination, dynamic_axes)
    _restore_realesrgan_scope_names(onnx, destination)
    _validate_export(onnx, onnxruntime, torch, model, sample, destination)


def export_safmn_real_x4(checkpoint_path, destination):
    onnx, onnxruntime, torch = _require_build_dependencies()
    nn = torch.nn
    functional = torch.nn.functional

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

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

    class SAFM(nn.Module):
        def __init__(self, channels):
            super().__init__()
            chunk_channels = channels // 4
            self.mfr = nn.ModuleList(
                nn.Conv2d(chunk_channels, chunk_channels, 3, 1, 1, groups=chunk_channels)
                for _ in range(4)
            )
            self.aggr = nn.Conv2d(channels, channels, 1)
            self.act = nn.GELU()

        def forward(self, value):
            height, width = value.shape[-2:]
            outputs = []
            for level, chunk in enumerate(value.chunk(4, dim=1)):
                if level:
                    size = (height // 2**level, width // 2**level)
                    chunk = functional.adaptive_max_pool2d(chunk, size)
                chunk = self.mfr[level](chunk)
                if level:
                    chunk = functional.interpolate(
                        chunk,
                        size=(height, width),
                        mode="nearest",
                    )
                outputs.append(chunk)
            return self.act(self.aggr(torch.cat(outputs, dim=1))) * value

    class CCM(nn.Module):
        def __init__(self, channels):
            super().__init__()
            self.ccm = nn.Sequential(
                nn.Conv2d(channels, channels * 2, 3, 1, 1),
                nn.GELU(),
                nn.Conv2d(channels * 2, channels, 1),
            )

        def forward(self, value):
            return self.ccm(value)

    class AttBlock(nn.Module):
        def __init__(self, channels):
            super().__init__()
            self.norm1 = LayerNorm(channels)
            self.norm2 = LayerNorm(channels)
            self.safm = SAFM(channels)
            self.ccm = CCM(channels)

        def forward(self, value):
            value = self.safm(self.norm1(value)) + value
            return self.ccm(self.norm2(value)) + value

    class SAFMN(nn.Module):
        def __init__(self):
            super().__init__()
            self.to_feat = nn.Conv2d(3, 128, 3, 1, 1)
            self.feats = nn.Sequential(*(AttBlock(128) for _ in range(16)))
            self.to_img = nn.Sequential(
                nn.Conv2d(128, 3 * 4**2, 3, 1, 1),
                nn.PixelShuffle(4),
            )

        def forward(self, value):
            value = self.to_feat(value)
            return self.to_img(self.feats(value) + value)

    model = SAFMN()
    model.load_state_dict(
        _load_weights(torch, checkpoint_path, "params_ema"),
        strict=True,
    )
    sample = torch.rand(1, 3, 256, 256)
    _export(torch, model, sample, destination)
    _validate_export(onnx, onnxruntime, torch, model, sample, destination)


def export_model(exporter, checkpoint_path, destination):
    exporters = {
        "realesrgan_wdn_x4": export_realesrgan_wdn_x4,
        "safmn_real_x4": export_safmn_real_x4,
    }
    try:
        export = exporters[exporter]
    except KeyError as error:
        raise ValueError(f"Unknown model exporter: {exporter}") from error
    export(checkpoint_path, destination)