"""PathGen-CLIP-L (jamessyx/PathGen-CLIP-L, CLIP trained on PathGen-1.6M) text + image towers.
Usage: python 16_extract_pathgenclip.py {text|image} [--pathgen|--prompts]
  text            -> cache/text_pathgenclip.npy          (N, 4, 768)  ORION templated texts (cache/text.json)
  text --pathgen  -> cache/pathgen_text_pathgenclip.npy  (N, 1, 768)  PathGen captions (cache/pathgen_meta.csv)
  text --prompts  -> cache/prompts_text_pathgenclip.npy  (P, 1, 768)  concept prompts (prompts.json)
  image           -> cache/he_pathgenclip.npy            (N, 8, 768)  ORION tiles, 8 dihedral views
  image --pathgen -> cache/pathgen_he_pathgenclip.npy    (N, 8, 768)  PathGen images, 8 dihedral views
All float16, projected CLIP embeddings (encode_text / encode_image), NOT L2-normalised -- same as 04b / 08.

Architecture: the README's 'ViT-B-16' is wrong for -L. The checkpoint has visual.conv1 (1024,3,14,14), 24 visual
blocks, visual.positional_embedding (577,1024) = 24x24+1 -> 336px, text width 768 / 12 blocks, proj -> 768.
That is OpenAI ViT-L/14@336; the authors' HF conversion (jamessyx/pathgenclip-vit-large-patch14-hf) is
initialised from openai/clip-vit-large-patch14-336 with hidden_act=quick_gelu -> open_clip 'ViT-L-14-336-quickgelu',
loaded strictly. Preprocess: 336px bicubic, OpenAI-CLIP mean/std (all ORION / PathGen images are square, so the
open_clip Resize+CenterCrop equals the plain square resize used in 04b). Tokenizer: OpenAI CLIP BPE, 77 tokens.
"""
import json
import os
import sys

import numpy as np
import torch
from PIL import Image, PngImagePlugin

PngImagePlugin.MAX_TEXT_CHUNK = 64 * 1024 * 1024  # some PathGen PNGs embed very large ICC profiles

import open_clip
from huggingface_hub import hf_hub_download

from common import CACHE, ORION, ROOT, dihedral, load_all_tiles

ARCH = "ViT-L-14-336-quickgelu"
NAME = "pathgenclip"
mode = sys.argv[1]
SET = sys.argv[sys.argv.index("--set") + 1] if "--set" in sys.argv else "pathgen" if "--pathgen" in sys.argv else None
PATHGEN = SET is not None  # encode an image-text pair set (cache/<SET>_meta.csv, e.g. PathGen-1.6M) instead of ORION tiles
PROMPTS = "--prompts" in sys.argv
dev = "cuda"
CLIPN = ([0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711])


def load_model():
    ckpt = hf_hub_download("jamessyx/PathGen-CLIP-L", "pathgen-clip-l.pt")
    model = open_clip.create_model(ARCH, pretrained=None)
    sd = torch.load(ckpt, map_location="cpu")
    sd = sd.get("state_dict", sd)
    sd = {k[len("module."):] if k.startswith("module.") else k: v for k, v in sd.items()}
    model.load_state_dict(sd, strict=True)
    return model.eval().to(dev)


model = load_model()

if mode == "text":
    tok = open_clip.get_tokenizer(ARCH)
    if PROMPTS:
        pr = json.load(open(os.path.join(ROOT, "prompts.json")))
        flat = [v["text"] for v in pr.values()]
        ids, NV = list(range(len(flat))), 1
    elif PATHGEN:
        import pandas as pd
        flat = pd.read_csv(os.path.join(CACHE, f"{SET}_meta.csv")).caption.tolist()
        ids, NV = list(range(len(flat))), 1
    else:
        texts = json.load(open(os.path.join(CACHE, "text.json")))
        ids, NV = sorted(int(k) for k in texts), 4
        flat = [texts[str(i)][v] for i in ids for v in range(NV)]
    out = []
    with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
        for s in range(0, len(flat), 256):
            out.append(model.encode_text(tok(flat[s:s + 256]).to(dev)).float().cpu().numpy())
    out = np.concatenate(out).reshape(len(ids), NV, -1).astype(np.float16)
    assert ids == list(range(len(ids)))
    pre = "prompts_" if PROMPTS else f"{SET}_" if PATHGEN else ""
    fn = os.path.join(CACHE, pre + f"text_{NAME}.npy")

elif mode == "image":
    size = model.visual.image_size[0] if isinstance(model.visual.image_size, (tuple, list)) else model.visual.image_size
    assert size == 336, size
    mean = torch.tensor(CLIPN[0]).view(1, 3, 1, 1)
    std = torch.tensor(CLIPN[1]).view(1, 3, 1, 1)
    if PATHGEN:
        import pandas as pd
        paths = pd.read_csv(os.path.join(CACHE, f"{SET}_meta.csv")).image_path.values
    else:
        paths = np.array([os.path.join(ORION, p) for p in load_all_tiles().image_path])
    out = None
    bs = 32
    with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
        for s in range(0, len(paths), bs):
            ims = [np.asarray(Image.open(p).convert("RGB").resize((size, size), Image.BICUBIC))
                   for p in paths[s:s + bs]]
            x = (torch.from_numpy(np.stack(ims)).permute(0, 3, 1, 2).float() / 255 - mean) / std
            for d in range(8):
                e = model.encode_image(dihedral(x, d).to(dev)).float().cpu().numpy()
                if out is None:
                    out = np.zeros((len(paths), 8, e.shape[1]), np.float16)
                out[s:s + bs, d] = e
            if s % 640 == 0:
                print(NAME, s, len(paths), flush=True)
    fn = os.path.join(CACHE, (f"{SET}_" if PATHGEN else "") + f"he_{NAME}.npy")
else:
    raise ValueError(mode)

np.save(fn, out)
print("saved", fn, out.shape, out.dtype, "finite:", np.isfinite(out).all(),
      "peak GPU MB:", torch.cuda.max_memory_allocated() // 2**20)