# KAGGLE: {"gpu": true, "internet": true, "datasets": [], "kernels": ["alekstop/02-imgstats"], "acc": "NvidiaTeslaT4"}

# %% [markdown]
# # 03 — Замороженные эмбеддинги бэкбонов (единственная обязательная GPU-трата)
#
# Читает кэш 320px из ноутбука `02-imgstats` (а не 28.6 ГБ исходников), гоняет
# 6 ракурсов каждого объекта через несколько предобученных бэкбонов и сохраняет:
#
# * `emb_<model>.npy` — (N, 6, D) fp16, пулинг всего изображения;
# * `tok_<primary>.npy` — (N, 6, 16, D) fp16, патч-токены, спулленные в сетку 4x4.
#   Они нужны для cross-view трансформера (Слой 2) и для карт внимания
#   в ноутбуке интерпретируемости.
#
# Домен здесь не natural images, а карты нормалей, поэтому берём несколько
# принципиально разных бэкбонов: self-supervised ViT (DINO), vision-language ViT
# (SigLIP) и свёрточный (ConvNeXt-V2) — последний заметно чувствительнее к
# высокочастотной текстуре.

# %%
import os, sys, io, json, time, math, tarfile, random, warnings
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from PIL import Image

warnings.filterwarnings('ignore')
SEED = 42
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.backends.cudnn.benchmark = True
T0 = time.time()

WORK = Path(os.environ.get('MQC_WORK', '/kaggle/working'))   # в Colab задаётся окружением
IN = Path(os.environ.get('MQC_IN', '/kaggle/input'))
DEV = 'cuda' if torch.cuda.is_available() else 'cpu'
NGPU = torch.cuda.device_count()
print('torch', torch.__version__, '| GPU:', [torch.cuda.get_device_name(i) for i in range(NGPU)])

# Колёса torch cu12.8 собраны начиная с sm_70; на P100 (sm_60) любое обращение к
# GPU падает с 'no kernel image is available'. Проверяем сразу, до скачивания
# весов, чтобы не терять время сессии.
CAP = torch.cuda.get_device_capability(0) if DEV == 'cuda' else (0, 0)
print('compute capability: %d.%d' % CAP)
if DEV == 'cuda' and CAP[0] < 7:
    raise RuntimeError('GPU sm_%d%d is too old for this torch build; '
                       'push with machine_shape=NvidiaTeslaT4' % CAP)
USE_HALF = CAP >= (7, 0)

CACHE_DIR = next(p.parent for p in IN.rglob('img320_train.tar'))
print('кэш:', CACHE_DIR, sorted(p.name for p in CACHE_DIR.iterdir())[:8])

# %%
# --- индекс tar-архивов: имя -> (offset, size), чтение без распаковки ---
class TarIndex:
    def __init__(self, path):
        self.path = str(path)
        self.idx = {}
        with tarfile.open(path) as tf:
            for m in tf:
                if m.isfile():
                    self.idx[Path(m.name).stem] = (m.offset_data, m.size)
        self.fh = None

    def read(self, key):
        if self.fh is None:
            self.fh = open(self.path, 'rb')
        off, size = self.idx[key]
        self.fh.seek(off)
        return self.fh.read(size)


t = time.time()
TARS = [TarIndex(CACHE_DIR / 'img320_train.tar'), TarIndex(CACHE_DIR / 'img320_test.tar')]
KEY2TAR = {}
for ti, tar in enumerate(TARS):
    for k in tar.idx:
        KEY2TAR[k] = ti
print('объектов в кэше:', len(KEY2TAR), '| индексация %.1f с' % (time.time() - t))

ids_train = list(TARS[0].idx)
ids_test = list(TARS[1].idx)
IDS = ids_train + ids_test
CACHE, NCOL, NROW, NVIEW = 320, 3, 2, 6

# %%
# --- датасет: один объект -> 6 ракурсов ---
# EMB<<<
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as Fn


class ViewDS(Dataset):
    def __init__(self, ids, size, mean, std):
        self.ids = ids
        self.size = size
        self.mean = torch.tensor(mean).view(3, 1, 1)
        self.std = torch.tensor(std).view(3, 1, 1)

    def __len__(self):
        return len(self.ids)

    def __getitem__(self, i):
        iid = self.ids[i]
        blob = TARS[KEY2TAR[iid]].read(iid)
        a = np.asarray(Image.open(io.BytesIO(blob)).convert('RGB'))
        vs = np.stack([a[r * CACHE:(r + 1) * CACHE, c * CACHE:(c + 1) * CACHE]
                       for r in range(NROW) for c in range(NCOL)])
        x = torch.from_numpy(vs).permute(0, 3, 1, 2).float() / 255.0
        if self.size != CACHE:
            x = Fn.interpolate(x, size=(self.size, self.size), mode='bicubic',
                               align_corners=False, antialias=True).clamp(0, 1)
        return (x - self.mean) / self.std


def make_loader(size, mean, std, bs):
    return DataLoader(ViewDS(IDS, size, mean, std), batch_size=bs, shuffle=False,
                      num_workers=2, pin_memory=True, persistent_workers=False)
# EMB>>>

# %%
# --- бэкбоны ---
# EMB<<<
import timm

MODELS = [
    # (имя в выводе, спека timm, разрешение, нужны ли патч-токены)
    ('dinov2b', 'vit_base_patch14_dinov2.lvd142m', 224, True),
    ('siglip2b', 'vit_base_patch16_siglip_256.webli', 256, False),
    ('cnxv2t', 'convnextv2_tiny.fcmae_ft_in22k_in1k', 256, False),
]


def build(spec, size):
    """Создать бэкбон под наше разрешение.

    У ViT из timm разрешение зашито в веса (у DINOv2 это 518) и patch_embed
    проверяет его строго. Пробрасываем img_size — timm сам интерполирует
    позиционные эмбеддинги при загрузке; для свёрточных моделей такого аргумента
    нет, поэтому падаем обратно на обычное создание.
    """
    try:
        return timm.create_model(spec, pretrained=True, num_classes=0,
                                 img_size=size, dynamic_img_size=True)
    except TypeError:
        return timm.create_model(spec, pretrained=True, num_classes=0)


# EMB>>>
# DINOv3 закрыт лицензией Meta и требует токена — пробуем, но не падаем без него
try:
    build('vit_base_patch16_dinov3.lvd1689m', 224)
    MODELS.insert(0, ('dinov3b', 'vit_base_patch16_dinov3.lvd1689m', 224, True))
    print('DINOv3 dostupen')
except Exception as e:
    print('DINOv3 nedostupen:', str(e)[:150])


# EMB<<<
@torch.no_grad()
def extract(spec, size, want_tokens, bs=32):
    m = build(spec, size)
    cfg = m.pretrained_cfg
    m = m.eval().to(DEV)
    m = m.half() if USE_HALF else m.float()
    loader = make_loader(size, cfg['mean'], cfg['std'], bs)

    embs, toks = [], []
    t = time.time()
    # однострочный прогресс, если есть tqdm; иначе редкие печати с ETA
    every = max(1, min(40, len(loader) // 10))
    try:
        from tqdm.auto import tqdm
        it = enumerate(tqdm(loader, desc=spec.split('.')[0], unit='батч'))
        every = 0
    except ImportError:
        print('  %s @%d: батчей %d по %d' % (spec.split('.')[0], size, len(loader), bs),
              flush=True)
        it = enumerate(loader)
    for bi, x in it:
        B = x.shape[0]
        x = x.view(B * NVIEW, 3, size, size).to(DEV, non_blocking=True)
        x = x.half() if USE_HALF else x.float()
        if want_tokens:
            feats = m.forward_features(x)                       # (B*6, 1+P, D)
            pooled = m.forward_head(feats, pre_logits=True)
            npre = getattr(m, 'num_prefix_tokens', 1)
            pt = feats[:, npre:]
            g = int(round(math.sqrt(pt.shape[1])))
            assert g * g == pt.shape[1], 'ne kvadratnaya setka tokenov: %d' % pt.shape[1]
            pt = pt.transpose(1, 2).reshape(pt.shape[0], -1, g, g)
            pt = Fn.adaptive_avg_pool2d(pt.float(), 4).flatten(2).transpose(1, 2)
            toks.append(pt.reshape(B, NVIEW, 16, -1).cpu().numpy().astype(np.float16))
        else:
            pooled = m(x)
        embs.append(pooled.reshape(B, NVIEW, -1).float().cpu().numpy().astype(np.float16))
        if every and (bi % every == 0 or bi + 1 == len(loader)):
            done = min((bi + 1) * bs, len(IDS))
            el = time.time() - t
            print('  %5d/%d  %.1f мин  ETA %.1f мин' % (done, len(IDS), el / 60,
                  el / done * (len(IDS) - done) / 60), flush=True)
    del m
    torch.cuda.empty_cache()
    return np.concatenate(embs), (np.concatenate(toks) if toks else None)

# EMB>>>

meta = {'ids': IDS, 'n_train': len(ids_train), 'models': {}}
for name, spec, size, want_tok in MODELS:
    t = time.time()
    try:
        E, TK = extract(spec, size, want_tok)
    except Exception as e:
        # эмбеддинги независимы: пропускаем сломавшийся бэкбон, а не сессию
        print('SKIP %s: %s' % (name, str(e)[:200]), flush=True)
        torch.cuda.empty_cache()
        continue
    np.save(WORK / ('emb_%s.npy' % name), E)
    if TK is not None:
        np.save(WORK / ('tok_%s.npy' % name), TK)
    meta['models'][name] = dict(spec=spec, size=size, dim=int(E.shape[-1]),
                                tokens=TK is not None, minutes=round((time.time() - t) / 60, 2))
    json.dump(meta, open(WORK / 'emb_meta.json', 'w'), indent=2)
    print('%s: %s  %.1f min' % (name, E.shape, (time.time() - t) / 60), flush=True)

json.dump(meta, open(WORK / 'emb_meta.json', 'w'), indent=2)
print('всего:', round((time.time() - T0) / 60, 1), 'мин')