"""Resize every modality to 224x224 once and store as .npy (uint8) for fast training."""
import os

import numpy as np
import tifffile
import torch
import torch.nn.functional as F
from PIL import Image
from scipy import ndimage

from common import CACHE, IMG, ORION, load_all_tiles, marker_matrix

os.makedirs(CACHE, exist_ok=True)
df = load_all_tiles()
N = len(df)
he = np.lib.format.open_memmap(os.path.join(CACHE, "he.npy"), "w+", np.uint8, (N, 3, IMG, IMG))
iff = np.lib.format.open_memmap(os.path.join(CACHE, "if.npy"), "w+", np.uint8, (N, 17, IMG, IMG))
nuc = np.lib.format.open_memmap(os.path.join(CACHE, "nuclei.npy"), "w+", np.uint8, (N, 2, IMG, IMG))

for i, r in df.iterrows():
    h = Image.open(os.path.join(ORION, r.image_path)).convert("RGB").resize((IMG, IMG), Image.BICUBIC)
    he[i] = np.asarray(h).transpose(2, 0, 1)
    x = torch.from_numpy(tifffile.imread(os.path.join(ORION, r.target_path)).astype(np.float32)).permute(2, 0, 1)[None]
    iff[i] = F.interpolate(x, size=(IMG, IMG), mode="bilinear", antialias=True, align_corners=False)[0].round().clamp(0, 255).byte().numpy()
    lab = tifffile.imread(os.path.join(ORION, r.nuclei_path))
    # boundary = pixels whose label differs from a neighbour (separates touching nuclei)
    bnd = (ndimage.maximum_filter(lab, 3) != ndimage.minimum_filter(lab, 3)) & (lab > 0)
    m = torch.from_numpy(np.stack([lab > 0, bnd]).astype(np.float32))[None]
    nuc[i] = (F.interpolate(m, size=(IMG, IMG), mode="area")[0] * 255).round().byte().numpy()
    if i % 500 == 0:
        print(i, N, flush=True)

np.save(os.path.join(CACHE, "markers.npy"), marker_matrix(df))
he.flush(); iff.flush(); nuc.flush()
print("done")