# -*- coding: utf-8 -*-
"""mine_layer3_step1_extract.py - 用老下半区模型对全库下半区图提特征（挖矿原料）。

输入: lower_all_data/images/*.jpg
输出: lower_all_feats.npy (n,1024) + lower_all_ids.json
"""
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

import json
import time

import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from transformers import Dinov2Model

WZJ = "/home/user/顾工交接/wzj"
IMG_DIR = os.path.join(WZJ, "lower_all_data/images")
MODEL_PATH = os.path.join(WZJ, "layer3_bg_model_output/best_layer3_bottom_model.pth")
OUT_NPY = os.path.join(WZJ, "lower_all_feats.npy")
OUT_IDS = os.path.join(WZJ, "lower_all_ids.json")

HF_MODEL_ID = "facebook/dinov2-large"
IMG_HEIGHT, IMG_WIDTH = 196, 392
FREEZE_BLOCKS = 18
BATCH = 128
DEVICE = torch.device("cuda:0")


class Dinov2Layer3Model(nn.Module):
    def __init__(self, freeze_blocks=FREEZE_BLOCKS):
        super().__init__()
        self.backbone = Dinov2Model.from_pretrained(HF_MODEL_ID)
        for p in self.backbone.parameters():
            p.requires_grad = False
        for i in range(freeze_blocks, len(self.backbone.encoder.layer)):
            for p in self.backbone.encoder.layer[i].parameters():
                p.requires_grad = True
        for p in self.backbone.layernorm.parameters():
            p.requires_grad = True

    def forward(self, x):
        out = self.backbone(x, interpolate_pos_encoding=True)
        return F.normalize(out.last_hidden_state[:, 0, :], p=2, dim=1)


class ImgDataset(Dataset):
    def __init__(self, files):
        self.files = files

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

    def __getitem__(self, i):
        img = cv2.imread(os.path.join(IMG_DIR, self.files[i]))
        if img is None:
            img = np.zeros((IMG_HEIGHT, IMG_WIDTH, 3), dtype=np.uint8)
        else:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, (IMG_WIDTH, IMG_HEIGHT), interpolation=cv2.INTER_CUBIC)
        img = (img.astype(np.float32) / 255.0 - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
        return torch.from_numpy(img.astype(np.float32).transpose(2, 0, 1))


def main():
    files = sorted(f for f in os.listdir(IMG_DIR) if f.endswith(".jpg"))
    print(f"[files] {len(files)}", flush=True)
    ids = [f[:-4] for f in files]

    print("[model] load old layer3 weights ...", flush=True)
    model = Dinov2Layer3Model()
    model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
    model.to(DEVICE).eval()

    loader = DataLoader(ImgDataset(files), batch_size=BATCH, shuffle=False,
                        num_workers=16, pin_memory=True)
    feats = np.zeros((len(files), 1024), dtype=np.float32)
    t0 = time.time()
    i0 = 0
    with torch.no_grad():
        for bi, batch in enumerate(loader):
            batch = batch.to(DEVICE, non_blocking=True)
            with torch.amp.autocast("cuda"):
                f = model(batch)
            feats[i0:i0 + batch.shape[0]] = f.float().cpu().numpy()
            i0 += batch.shape[0]
            if bi % 50 == 0:
                el = (time.time() - t0) / 60
                print(f"  {i0}/{len(files)} elapsed={el:.1f}m", flush=True)

    np.save(OUT_NPY, feats)
    with open(OUT_IDS, "w", encoding="utf-8") as f:
        json.dump(ids, f, ensure_ascii=False)
    print(f"[DONE] {i0} feats -> {OUT_NPY} elapsed={(time.time()-t0)/60:.1f}m", flush=True)


if __name__ == "__main__":
    main()