"""Frozen ImageBind-Huge vision embeddings of every H&E tile under each of the 8 dihedral transforms.
Output cache/he_imagebind.npy : (N, 8, 1024) float16, L2-normalised (ImageBind postprocessor)."""
import os

import numpy as np
import torch

from common import CACHE, dihedral, load_imagebind_vision

MEAN = torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1)
STD = torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1)

he = np.load(os.path.join(CACHE, "he.npy"), mmap_mode="r")
model, MT = load_imagebind_vision()
out = np.zeros((len(he), 8, 1024), np.float16)
bs = 64
with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
    for s in range(0, len(he), bs):
        x = (torch.from_numpy(np.ascontiguousarray(he[s:s + bs])).float() / 255 - MEAN) / STD
        for d in range(8):
            e = model({MT.VISION: dihedral(x, d).cuda()})[MT.VISION]
            out[s:s + bs, d] = e.float().cpu().numpy()
        print(s, len(he), flush=True)
np.save(os.path.join(CACHE, "he_imagebind.npy"), out)
print("saved", out.shape)