"""
Precompute SigLIP + Gemma embeddings for all training images.

Run precompute_wd_tags.py first to produce WD_TAGS_JSON.

Output files in EMBEDDINGS_DIR:
  index.json          {image_path: row_index}
  siglip_pools.npy    float16 [N, 256, 1152]   SigLIP patch features (pre-projection)
  gemma_embs.npy      float16 [N, 300, 2304]   Gemma text embeddings
  progress.json       {last_completed: int}      restartable

These files are memory-mapped at training time by train.py.
"""
import json
import os
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from tqdm import tqdm
from transformers import AutoImageProcessor, AutoModelForCausalLM, AutoTokenizer, SiglipVisionModel
from transformers import logging as hf_logging
hf_logging.set_verbosity_error()

# Settings --------------------------------------------------------------------

IMAGE_DIR      = "/home/nobus/Raid0/DataSet/Images1"
WD_TAGS_JSON   = "/home/nobus/Raid0/DataSet/embeddings/wd_tags.json"
EMBEDDINGS_DIR = "/home/nobus/Raid0/DataSet/embeddings"

SIGLIP_DEVICE  = "cuda:0"
GEMMA_DEVICE   = "cuda:0"
SIGLIP_BATCH   = 64
GEMMA_BATCH    = 16

# -----------------------------------------------------------------------------

_SIGLIP_ID  = "google/siglip-so400m-patch14-384"
_GEMMA_ID   = "Efficient-Large-Model/gemma-2-2b-it"
_N_IP       = 256
_SIGLIP_DIM = 1152
_N_TEXT     = 300
_GEMMA_DIM  = 2304

_CHI_PROMPT = "\n".join([
    'Given a user prompt, generate an "Enhanced prompt" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:',
    "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.",
    "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.",
    "Here are examples of how to transform or refine prompts:",
    "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.",
    "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.",
    "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:",
    "User Prompt: ",
])

_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp"}


def _find_images(root):
    return sorted(p for p in Path(root).rglob("*") if p.is_file() and p.suffix.lower() in _IMAGE_EXTS)


def _open_or_create_memmap(path, shape, dtype="float16"):
    if os.path.exists(path):
        mm = np.lib.format.open_memmap(path, mode="r+")
        assert mm.shape == shape and mm.dtype == np.dtype(dtype), \
            f"Existing {path} shape/dtype mismatch: {mm.shape} vs {shape}"
        return mm
    return np.lib.format.open_memmap(path, mode="w+", dtype=dtype, shape=shape)


def main():
    os.makedirs(EMBEDDINGS_DIR, exist_ok=True)

    paths = _find_images(IMAGE_DIR)
    N     = len(paths)
    print(f"Found {N} images in {IMAGE_DIR}")

    with open(WD_TAGS_JSON, encoding="utf-8") as fh:
        wd_tags = json.load(fh)
    print(f"Loaded WD tags: {len(wd_tags)} entries")

    # ── index ─────────────────────────────────────────────────────────────────
    index_path = os.path.join(EMBEDDINGS_DIR, "index.json")
    if os.path.exists(index_path):
        with open(index_path) as fh:
            index = json.load(fh)
        assert len(index) == N, f"Index has {len(index)} entries but found {N} images. Delete index.json to rebuild."
        print(f"Loaded existing index ({len(index)} entries)")
    else:
        index = {str(p): i for i, p in enumerate(paths)}
        with open(index_path, "w") as fh:
            json.dump(index, fh, indent=2)
        print(f"Built index → {index_path}")

    # ── progress ───────────────────────────────────────────────────────────────
    progress_path = os.path.join(EMBEDDINGS_DIR, "progress.json")
    start_from    = 0
    if os.path.exists(progress_path):
        with open(progress_path) as fh:
            start_from = json.load(fh).get("last_completed", -1) + 1
        print(f"Resuming from image {start_from}/{N}")

    if start_from >= N:
        print("All images already computed.")
        return

    # ── memmap files ───────────────────────────────────────────────────────────
    siglip_mm = _open_or_create_memmap(
        os.path.join(EMBEDDINGS_DIR, "siglip_pools.npy"), (N, _N_IP, _SIGLIP_DIM)
    )
    gemma_mm  = _open_or_create_memmap(
        os.path.join(EMBEDDINGS_DIR, "gemma_embs.npy"),   (N, _N_TEXT, _GEMMA_DIM)
    )
    print(f"siglip_pools.npy  {siglip_mm.nbytes / 1e9:.1f} GB")
    print(f"gemma_embs.npy    {gemma_mm.nbytes  / 1e9:.1f} GB")

    remaining = paths[start_from:]

    # ── SigLIP ─────────────────────────────────────────────────────────────────
    print(f"\nLoading SigLIP on {SIGLIP_DEVICE}...")
    proc   = AutoImageProcessor.from_pretrained(_SIGLIP_ID)
    siglip = SiglipVisionModel.from_pretrained(_SIGLIP_ID, torch_dtype=torch.float16).eval().to(SIGLIP_DEVICE)

    for bi in tqdm(range(0, len(remaining), SIGLIP_BATCH), desc="SigLIP"):
        batch_paths = remaining[bi: bi + SIGLIP_BATCH]
        imgs, valid = [], []
        for j, p in enumerate(batch_paths):
            try:
                imgs.append(Image.open(p).convert("RGB"))
                valid.append(j)
            except Exception as e:
                print(f"  skip {p}: {e}")
        if not imgs:
            continue
        inputs  = proc(images=imgs, return_tensors="pt").to(SIGLIP_DEVICE)
        with torch.no_grad():
            patches = siglip(**inputs).last_hidden_state        # [B, 729, 1152]
        pooled = (
            F.adaptive_avg_pool1d(patches.float().permute(0, 2, 1), _N_IP)
            .permute(0, 2, 1).to(torch.float16).cpu().numpy()   # [B, 256, 1152]
        )
        for arr, j in zip(pooled, valid):
            siglip_mm[start_from + bi + j] = arr
        siglip_mm.flush()

    del siglip
    torch.cuda.empty_cache()

    # ── Gemma ──────────────────────────────────────────────────────────────────
    print(f"\nLoading Gemma on {GEMMA_DEVICE}...")
    tokenizer = AutoTokenizer.from_pretrained(_GEMMA_ID)
    tokenizer.padding_side = "right"
    gemma = (
        AutoModelForCausalLM.from_pretrained(_GEMMA_ID, torch_dtype=torch.float16)
        .get_decoder().eval().to(GEMMA_DEVICE)
    )

    num_chi = len(tokenizer.encode(_CHI_PROMPT))
    max_len = num_chi + _N_TEXT - 2
    select  = [0] + list(range(-(_N_TEXT - 1), 0))

    for bi in tqdm(range(0, len(remaining), GEMMA_BATCH), desc="Gemma"):
        batch_paths = remaining[bi: bi + GEMMA_BATCH]
        texts = [_CHI_PROMPT + wd_tags.get(str(p.resolve()), "") for p in batch_paths]
        tok   = tokenizer(texts, max_length=max_len, padding="max_length",
                          truncation=True, return_tensors="pt").to(GEMMA_DEVICE)
        with torch.no_grad():
            emb = gemma(input_ids=tok.input_ids, attention_mask=tok.attention_mask).last_hidden_state
            emb = emb[:, select, :].to(torch.float16).cpu().numpy()   # [B, 300, 2304]
        for j in range(len(batch_paths)):
            gemma_mm[start_from + bi + j] = emb[j]
        gemma_mm.flush()
        with open(progress_path, "w") as fh:
            json.dump({"last_completed": start_from + bi + len(batch_paths) - 1}, fh)

    del gemma
    print(f"\nDone!")
    print(f"  {EMBEDDINGS_DIR}/siglip_pools.npy")
    print(f"  {EMBEDDINGS_DIR}/gemma_embs.npy")
    print(f"  {EMBEDDINGS_DIR}/index.json")


if __name__ == "__main__":
    main()