"""Exports the models of the demos (examples/models), one ONNX file each.

Every model takes `image` (1, 3, S, S), RGB in [0, 1] (S = 224 unless SIZES says otherwise; the
model's own normalization is part of the graph), and returns:
  features     (1, C, H, W)  the last feature map, before pooling (ViTs: the patch
                             tokens after the final norm, as a grid)
  cell_logits  (1, K, H, W)  classifiers only: the classification head applied to every cell (the
                             dense readout); for models that pool their features
                             (CNNs, Swin, MaxViT) the mean over the cells equals
                             the model's logits; for ViTs, which classify from
                             the class token, it is an approximation
  logits       (1, K)        only for models whose prediction is not the mean of the cells (CLIP
                             zero-shot: attention pooling or a class token): the model's own logits
Weights are stored as float16 and cast to float32 when the model is loaded, which
halves the download without changing the computation. Each model is the largest
of its family whose file stays under 100 MB, except where a specific model is the point
(CLIP RN50x4: the model of the typographic attack of Goh et al. 2021).

Writes NAME.onnx to the current folder and checks it against PyTorch on the .jpg images in $SAMPLES.
Run one model per process, with a memory cap:

    SAMPLES=path/to/images systemd-run --user --scope -p MemoryMax=12G uv run --no-project \
        --with torch --with torchvision --with timm --with onnx --with onnxruntime --with pillow \
        --with open_clip_torch --with safetensors --with huggingface_hub \
        python export_models.py NAME

The files are published on Hugging Face (karay/demo-models), not in the blog repo.
"""
import os
import sys
from pathlib import Path

import numpy as np
import onnx
import onnxruntime as ort
import timm
import torch
from onnx import helper, numpy_helper
from PIL import Image
from torch import nn

OUT = Path.cwd()                                  # the .onnx files are written to the current folder
# the check images: any folder of .jpg files (default: the blog's examples, next to this file's repo)
SAMPLES = Path(os.environ.get('SAMPLES', Path(__file__).resolve().parents[2] / 'assets' / 'img' / 'samples'))


# Feature maps: (model, image) -> (B, H, W, C) cells

def cnn_cells(model, x):
    return model.forward_features(x).permute(0, 2, 3, 1)


def swin_cells(model, x):
    return model.forward_features(x)                      # timm Swin is already (B, H, W, C)


def vit_cells(model, x):
    tokens = model.forward_features(x)[:, model.num_prefix_tokens:]   # drop the class token
    return tokens.reshape(tokens.shape[0], *model.patch_embed.grid_size, -1)


# Dense readouts: model -> function from (B, H, W, C) cells to (B, H, W, K) cell logits

def linear_head(model):
    """ResNet, Swin, ViT: the head is one Linear layer, applied to every cell."""
    return model.get_classifier()


def norm_head(model):
    """ConvNeXt: LayerNorm after pooling, then a Linear layer."""
    return lambda cells: model.head.fc(pooled_layernorm(cells, model.head.norm))


def norm_tanh_head(model):
    """MaxViT: LayerNorm after pooling, Linear + tanh (linearized at the image's own
    pre-activation, so that the mean of the cells is still tanh of it), then a Linear layer."""
    def readout(cells):
        a = model.head.pre_logits.fc(pooled_layernorm(cells, model.head.norm))
        a_bar = a.mean(dim=(1, 2), keepdim=True)
        t = torch.tanh(a_bar)
        return model.head.fc(t + (1 - t ** 2) * (a - a_bar))
    return readout


def pooled_layernorm(cells, ln):
    """Per-cell share of ``ln(mean of the cells)``: each cell centred on its own mean and
    divided by the spread of the pooled vector, then the affine."""
    pooled = cells.mean(dim=(1, 2))
    sigma = torch.sqrt(pooled.var(dim=-1, unbiased=False) + ln.eps)
    return (cells - cells.mean(dim=-1, keepdim=True)) / sigma[:, None, None, None] * ln.weight + ln.bias


# Models: () -> the pretrained model, with `forward_features` and `pretrained_cfg` (mean, std)

def timm_model(checkpoint, **kwargs):
    return lambda: timm.create_model(checkpoint, pretrained=True, **kwargs)


def dino_resnet50():
    """DINOv1 ResNet-50: the torchvision trunk with DINO's weights (the checkpoint has no classifier)."""
    import torchvision
    m = torchvision.models.resnet50(weights=None)
    url = 'https://dl.fbaipublicfiles.com/dino/dino_resnet50_pretrain/dino_resnet50_pretrain.pth'
    m.load_state_dict(torch.hub.load_state_dict_from_url(url, map_location='cpu'), strict=False)
    trunk = nn.Sequential(*list(m.children())[:-2])
    trunk.forward_features = trunk.forward
    trunk.pretrained_cfg = {'mean': (0.485, 0.456, 0.406), 'std': (0.229, 0.224, 0.225)}
    return trunk


# Pathology models (inputs: H&E tiles at about 0.5 um/px, except the tumour detector: 0.39 um/px at 224 px)

def kaiko_vits16():
    """Kaiko ViT-S/16 (Kaiko non-commercial licence), normalized with mean = std = 0.5."""
    m = torch.hub.load('kaiko-ai/towards_large_pathology_fms', 'vits16', trust_repo=True)
    m.pretrained_cfg = {'mean': (0.5, 0.5, 0.5), 'std': (0.5, 0.5, 0.5)}
    return m


def ctranspath():
    """CTransPath (GPL-3.0): a Swin-T with a convolutional stem, timm port by 1aurent."""
    from timm.layers.helpers import to_2tuple

    class ConvStem(nn.Module):   # from the model card (adapted from TransPath, ctran.py)
        def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=768, norm_layer=None, **kwargs):
            super().__init__()
            self.img_size, self.patch_size = to_2tuple(img_size), to_2tuple(patch_size)
            self.grid_size = (self.img_size[0] // self.patch_size[0], self.img_size[1] // self.patch_size[1])
            self.num_patches = self.grid_size[0] * self.grid_size[1]
            stem, input_dim, output_dim = [], 3, embed_dim // 8
            for _ in range(2):
                stem += [nn.Conv2d(input_dim, output_dim, 3, 2, 1, bias=False), nn.BatchNorm2d(output_dim), nn.ReLU(inplace=True)]
                input_dim, output_dim = output_dim, output_dim * 2
            stem.append(nn.Conv2d(input_dim, embed_dim, kernel_size=1))
            self.proj = nn.Sequential(*stem)
            self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

        def forward(self, x):
            return self.norm(self.proj(x).permute(0, 2, 3, 1))

    m = timm.create_model('hf-hub:1aurent/swin_tiny_patch4_window7_224.CTransPath', embed_layer=ConvStem, pretrained=True)
    cfg = timm.data.resolve_model_data_config(m)
    m.pretrained_cfg = {'mean': cfg['mean'], 'std': cfg['std']}
    return m


class TorchvisionClassifier(nn.Module):
    """A torchvision ResNet as trunk + fc, with the timm-style methods used here."""

    def __init__(self, resnet, mean, std):
        super().__init__()
        self.trunk, self.fc = nn.Sequential(*list(resnet.children())[:-2]), resnet.fc
        self.pretrained_cfg = {'mean': tuple(mean), 'std': tuple(std)}

    def forward_features(self, x):
        return self.trunk(x)

    def get_classifier(self):
        return self.fc

    def forward(self, x):
        return self.fc(self.trunk(x).mean(dim=(2, 3)))


def tumour_detector():
    """Breast tumour detector, ResNet-34 on TCGA-BRCA (Le et al. 2020; CC-BY-4.0): classes [Other, Tumor];
    trained on 350 px tiles at 0.25 um/px resized to 224 px."""
    import json
    import torchvision
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file
    repo = 'kaczmarj/breast-tumor-resnet34.tcga-brca'
    cfg = json.load(open(hf_hub_download(repo, 'config.json')))
    norm = next(t for t in cfg['transform'] if t['name'] == 'Normalize')['arguments']
    resnet = torchvision.models.resnet34(num_classes=2)
    resnet.load_state_dict(load_file(hf_hub_download(repo, 'model.safetensors')))
    return TorchvisionClassifier(resnet, norm['mean'], norm['std'])


# CLIP models with a zero-shot head. Their prediction is not the mean of the cells (attention pooling,
# class token), so they also return `logits`; cell_logits is the per-location readout (as in MaskCLIP,
# Zhou et al. 2022): each location's embedding through the same head.

TISSUE_PROMPTS = {'tumour': 'tumor cells', 'stroma': 'stroma', 'lymphocytes': 'lymphocytes',
                  'necrosis': 'necrosis', 'fat': 'adipose tissue'}
TISSUE_TEMPLATES = ['a histopathology image of {}.', 'an H&E image of {}.']


def text_head(clip, tokenizer, names, templates, cache=None):
    """Zero-shot classifier: the mean of the normalized text embeddings of the templates, per class. Computed
    on the GPU if there is one (80 templates x 1000 classes take hours on a CPU) and cached in `cache`."""
    if cache and Path(cache).exists():
        return torch.load(cache)
    dev = 'cuda' if torch.cuda.is_available() else 'cpu'
    clip = clip.to(dev)
    with torch.no_grad():
        rows = []
        for name in names:
            t = clip.encode_text(tokenizer([tpl.format(name) if isinstance(tpl, str) else tpl(name) for tpl in templates]).to(dev))
            t = t / t.norm(dim=-1, keepdim=True)
            rows.append(t.mean(0) / t.mean(0).norm())
    head = torch.stack(rows).cpu()
    clip.cpu()
    if cache:
        torch.save(head, cache)
    return head


class ClipRN(nn.Module):
    """OpenAI CLIP ResNet (attention pooling) with its zero-shot ImageNet classifier."""

    def __init__(self, name='RN50x4'):
        super().__init__()
        import open_clip
        from open_clip.zero_shot_metadata import IMAGENET_CLASSNAMES, OPENAI_IMAGENET_TEMPLATES
        clip, _, pre = open_clip.create_model_and_transforms(name, pretrained='openai')
        clip.eval()
        self.visual = clip.visual
        self.register_buffer('text', text_head(clip, open_clip.get_tokenizer(name), IMAGENET_CLASSNAMES, OPENAI_IMAGENET_TEMPLATES,
                                                 Path(__file__).resolve().parent / f'zeroshot_{name}_imagenet.pt'))
        self.scale = float(clip.logit_scale.exp())
        n = pre.transforms[-1]
        self.pretrained_cfg = {'mean': tuple(n.mean), 'std': tuple(n.std)}

    def forward_features(self, x):
        v = self.visual
        return v.layer4(v.layer3(v.layer2(v.layer1(v.stem(x)))))

    def classify(self, emb):
        return self.scale * (emb / emb.norm(dim=-1, keepdim=True)) @ self.text.T

    def cell_logits(self, cells):            # (B, H, W, C) -> (B, H, W, K): c_proj(v_proj(x + pos)) per cell
        pool = self.visual.attnpool
        B, H, W, C = cells.shape
        x = cells + pool.positional_embedding[1:].reshape(H, W, C)
        return self.classify(pool.c_proj(pool.v_proj(x)))

    def forward(self, x):                    # the model's own zero-shot logits
        return self.classify(self.visual.attnpool(self.forward_features(x)))


class ClipViT(nn.Module):
    """An open_clip ViT (class token) with a zero-shot head over tissue classes (QuiltNet-B-32, MIT)."""

    def __init__(self, hub='hf-hub:wisdomik/QuiltNet-B-32'):
        super().__init__()
        import open_clip
        clip, _, pre = open_clip.create_model_and_transforms(hub)
        clip.eval()
        self.visual = clip.visual
        self.register_buffer('text', text_head(clip, open_clip.get_tokenizer(hub), list(TISSUE_PROMPTS.values()), TISSUE_TEMPLATES))
        self.scale = float(clip.logit_scale.exp())
        n = pre.transforms[-1]
        self.pretrained_cfg = {'mean': tuple(n.mean), 'std': tuple(n.std)}

    def tokens(self, x):                     # all tokens after ln_post: (B, 1 + HW, C)
        v = self.visual
        x = v.conv1(x)
        B, C, H, W = x.shape
        x = x.reshape(B, C, H * W).permute(0, 2, 1)
        x = torch.cat([v.class_embedding.expand(B, 1, C), x], dim=1) + v.positional_embedding
        x = v.transformer(v.ln_pre(x))
        return v.ln_post(x), H, W

    def forward_features(self, x):           # (B, H, W, C), patch tokens only
        t, H, W = self.tokens(x)
        return t[:, 1:].reshape(t.shape[0], H, W, -1)

    def classify(self, emb):
        return self.scale * (emb / emb.norm(dim=-1, keepdim=True)) @ self.text.T

    def cell_logits(self, cells):
        return self.classify(cells @ self.visual.proj)

    def forward(self, x):
        t, _, _ = self.tokens(x)
        return self.classify(t[:, 0] @ self.visual.proj)


class HibouB(nn.Module):
    """Hibou-B (HistAI, Apache-2.0; DINOv2 ViT-B/14 with 4 registers; hibou_port.py) with a tissue head trained
    on its mean patch token (hibou_b_tissue_head.pt, from gap_mil_revision/tmp/kuleuven_talk/deck/hibou_head.py):
    classes tumour, stroma, lymphocytes, necrosis, fat; BCSS tiles of 112 um at 0.5 um/px; no lobular case in
    training. The head is linear on the mean of the patch tokens, so the mean of the cell logits is its logits."""

    def __init__(self):
        super().__init__()
        sys.path.insert(0, str(Path(__file__).resolve().parent))
        from hibou_port import hibou_b
        self.vit = hibou_b()
        saved = torch.load(Path(__file__).resolve().parent / 'hibou_b_tissue_head.pt')
        self.head = nn.Linear(self.vit.embed_dim, len(saved['classes']))
        self.head.load_state_dict(saved['state_dict'])
        self.pretrained_cfg = {'mean': tuple(saved['mean']), 'std': tuple(saved['std'])}

    def forward_features(self, x):           # (B, H, W, C): patch tokens after the final norm
        t = self.vit.forward_features(x)[:, self.vit.num_prefix_tokens:]
        return t.reshape(t.shape[0], *self.vit.patch_embed.grid_size, -1)

    def get_classifier(self):
        return self.head

    def forward(self, x):
        return self.head(self.forward_features(x).mean(dim=(1, 2)))


def bhwc_cells(model, x):
    return model.forward_features(x)


def clip_cells(model, x):
    return model.forward_features(x) if isinstance(model, ClipViT) else model.forward_features(x).permute(0, 2, 3, 1)


def zero_shot_head(model):
    return model.cell_logits


MODELS = {   # file name -> (model, feature map, readout; None: self-supervised, no classifier)
    # classifiers (ImageNet)
    'resnet50':        (timm_model('resnet50.tv2_in1k'), cnn_cells, linear_head),
    'convnext_small':  (timm_model('convnext_small.fb_in1k'), cnn_cells, norm_head),
    'swin_small':      (timm_model('swin_small_patch4_window7_224.ms_in1k'), swin_cells, linear_head),
    'maxvit_tiny':     (timm_model('maxvit_tiny_tf_224.in1k'), cnn_cells, norm_tanh_head),
    'vit_small':       (timm_model('vit_small_patch16_224.augreg_in21k_ft_in1k'), vit_cells, linear_head),
    'deit_small':      (timm_model('deit_small_patch16_224.fb_in1k'), vit_cells, linear_head),
    'deit3_medium':    (timm_model('deit3_medium_patch16_224.fb_in1k'), vit_cells, linear_head),
    # self-supervised and image-text models
    'dino_resnet50':         (dino_resnet50, cnn_cells, None),
    'dino_vit_small':        (timm_model('vit_small_patch16_224.dino'), vit_cells, None),
    'dinov2_vit_small':      (timm_model('vit_small_patch14_dinov2.lvd142m', img_size=224), vit_cells, None),
    'dinov2_reg_vit_small':  (timm_model('vit_small_patch14_reg4_dinov2.lvd142m', img_size=224), vit_cells, None),
    'dinov3_vit_small_plus': (timm_model('vit_small_plus_patch16_dinov3.lvd1689m', img_size=224), vit_cells, None),
    'dinov3_convnext_small': (timm_model('convnext_small.dinov3_lvd1689m'), cnn_cells, None),
    'clip_resnet101':        (timm_model('resnet101_clip_gap.openai'), cnn_cells, None),
    'tinyclip_vit_medium':   (timm_model('vit_medium_patch16_clip_224.tinyclip_yfcc15m'), vit_cells, None),
    'fcmae_convnextv2_tiny': (timm_model('convnextv2_tiny.fcmae'), cnn_cells, None),
    # CLIP with a zero-shot head (outputs also `logits`)
    'clip_rn50x4':           (ClipRN, clip_cells, zero_shot_head),
    # pathology
    'tumour_detector':       (tumour_detector, cnn_cells, linear_head),
    'kaiko_vit_small':       (kaiko_vits16, vit_cells, None),
    'ctranspath':            (ctranspath, swin_cells, None),
    'hibou_b':               (HibouB, bhwc_cells, linear_head),
    'quiltnet_b32':          (ClipViT, clip_cells, zero_shot_head),
}
SIZES = {'clip_rn50x4': 288}          # input sizes other than 224
IMAGE_LOGITS = {'clip_rn50x4', 'quiltnet_b32'}   # prediction is not the mean of the cell logits


class DenseReadout(nn.Module):
    """Normalizes the image and returns the model's feature map and, for a classifier, its dense readout."""

    def __init__(self, model, cells, head, image_logits=False):
        super().__init__()
        self.model, self.cells, self.head, self.image_logits = model, cells, head and head(model), image_logits
        cfg = model.pretrained_cfg
        self.register_buffer('mean', torch.tensor(cfg['mean'])[:, None, None])
        self.register_buffer('std', torch.tensor(cfg['std'])[:, None, None])

    def forward(self, image):
        cells = self.cells(self.model, (image - self.mean) / self.std)
        if self.head is None:
            return cells.permute(0, 3, 1, 2)
        out = (cells.permute(0, 3, 1, 2), self.head(cells).permute(0, 3, 1, 2))
        if self.image_logits:
            out += (self.model((image - self.mean) / self.std),)
        return out


def weights_to_float16(model):
    """Store large float32 weights as float16, each followed by a Cast back to float32."""
    graph = model.graph
    casts = []
    for init in list(graph.initializer):
        if init.data_type != onnx.TensorProto.FLOAT or np.prod(init.dims) < 1024:
            continue
        half = numpy_helper.from_array(numpy_helper.to_array(init).astype(np.float16), init.name + '_fp16')
        graph.initializer.remove(init)
        graph.initializer.append(half)
        casts.append(helper.make_node('Cast', [half.name], [init.name], to=onnx.TensorProto.FLOAT))
    nodes = casts + list(graph.node)
    del graph.node[:]
    graph.node.extend(nodes)
    return model


def export(name):
    size = SIZES.get(name, 224)
    create, cells, head = MODELS[name]
    model = create().eval()
    dense = DenseReadout(model, cells, head, name in IMAGE_LOGITS).eval()
    outputs = ['features'] if head is None else ['features', 'cell_logits'] + (['logits'] if name in IMAGE_LOGITS else [])
    path = OUT / f'{name}.onnx'
    torch.onnx.export(dense, torch.rand(1, 3, size, size), path, input_names=['image'],
                      output_names=outputs, opset_version=17, dynamo=False)
    onnx_model = weights_to_float16(onnx.load(path))
    onnx.checker.check_model(onnx_model)
    onnx.save(onnx_model, path)

    # check on the sample images: the file against PyTorch (the last output: cell logits or features),
    # and for a classifier the readout against the model's logits
    session = ort.InferenceSession(path)
    file_error, readout_error, same_top1, corr = 0, 0, 0, 1
    images = sorted(SAMPLES.glob('*.jpg'))
    for p in images:
        x = torch.from_numpy(np.asarray(Image.open(p).convert('RGB').resize((size, size), Image.BICUBIC)))
        x = x.permute(2, 0, 1)[None].float() / 255
        got_all = session.run(None, {'image': x.numpy()})
        got = got_all[min(1, len(got_all) - 1)]          # cell logits, or features without a head
        with torch.no_grad():
            out = dense(x)
            expected = (out if head is None else out[1]).numpy()
        file_error = max(file_error, np.abs(got - expected).max() / np.abs(expected).max())
        corr = min(corr, np.corrcoef(got.ravel(), expected.ravel())[0, 1])
        if head is not None:
            with torch.no_grad():
                logits = model((x - dense.mean) / dense.std).numpy()
            if name in IMAGE_LOGITS:   # the file's logits against the model; the cells' mean is not the prediction
                readout_error = max(readout_error, np.abs(got_all[2] - logits).max())
                same_top1 += got_all[2].argmax() == logits.argmax()
            else:
                readout_error = max(readout_error, np.abs(got.mean((2, 3)) - logits).max())
                same_top1 += got.mean((2, 3)).argmax() == logits.argmax()
    report = f'{path.name}: {path.stat().st_size / 2**20:.1f} MB, {outputs[min(1, len(outputs) - 1)]} {got.shape[1:]}, ' \
             f'max |file - pytorch| / max = {file_error:.4f}, min correlation {corr:.5f}'
    if head is not None:
        what = '|file logits - logits|' if name in IMAGE_LOGITS else '|mean(cell_logits) - logits|'
        report += f', max {what} = {readout_error:.4f}, same top-1 {same_top1}/{len(images)}'
    print(report, flush=True)


if __name__ == '__main__':
    for name in sys.argv[1:] or MODELS:
        export(name)