"""Write src/clip-keyword-vectors.ts and test/clip-keywords-sunset.json.

bin/clip-keyword-vectors runs this in a pinned Python image. It downloads
Ente's Discover list, Ente's CLIP text and image models, the CLIP tokenizer
and a public-domain photo of a sunset, checks each against its sha256, and
writes:

- src/clip-keyword-vectors.ts: each Discover category of Ente's photos app,
  with the text embedding Ente's CLIP text model gives its query;
- test/clip-keywords-sunset.json: the CLIP embedding Ente's image model gives
  the photo, prepared as Ente's apps prepare a photo.

Both embeddings are scaled to length 1, as Ente's apps scale them.
"""

import hashlib
import io
import json
import urllib.request

import numpy as np
import onnxruntime
from PIL import Image, ImageOps
from tokenizers import Tokenizer

ENTE_COMMIT = "f5ad70a860ed25d0c8057b0efb4bd6eaefe4b0c5"
DISCOVER_PATH = "mobile/apps/photos/assets/discover.json"
DISCOVER_URL = (
    f"https://raw.githubusercontent.com/ente-io/ente/{ENTE_COMMIT}/{DISCOVER_PATH}"
)
DISCOVER_SHA256 = "cf70ac7116b592f7921ca4cf67cdbb421a8e6b287fbe44b2be678b8d198e5c24"

TEXT_MODEL = "mobileclip_s2_text_opset18_quant.onnx"
TEXT_MODEL_SHA256 = "d92f33dfcff83077fc2e0d3414250710efbb51795dfd89767bdbefb5fdc47322"
IMAGE_MODEL = "mobileclip_s2_image_gelu_opset20.onnx"
IMAGE_MODEL_SHA256 = "205a430af825e501c5138e5bb9abea942482a7a4fd4a680e98e47cf0830dce7e"

# The CLIP tokenizer Hugging Face publishes for OpenAI's CLIP. Its vocabulary
# and merges are those of bpe_simple_vocab_16e6.txt, the vocabulary Ente's
# apps tokenize with.
TOKENIZER_URL = (
    "https://huggingface.co/openai/clip-vit-base-patch32/resolve/"
    "3d74acf9a28c67741b2f4f2ea7635f0aaf6f0268/tokenizer.json"
)
TOKENIZER_SHA256 = "b556ac8c99757ffb677208af34bc8c6721572114111a6e0aaf5fa69ff0b8d842"

# "Sunset off the Massachusetts coast", a NOAA photo in the public domain:
# https://commons.wikimedia.org/wiki/File:Sunset_off_the_Massachusetts_coast_-_NOAA.jpg
PHOTO_URL = (
    "https://upload.wikimedia.org/wikipedia/commons/d/db/"
    "Sunset_off_the_Massachusetts_coast_-_NOAA.jpg"
)
PHOTO_SHA256 = "be62f3289c9bf7b75af68cd6978021d54e8a531d5d4cc493f6e9fbf5e1ab74d3"

# Ente's text model takes 77 token ids, padded with 0; its image model takes
# a 256x256 RGB image.
TOKEN_COUNT = 77
IMAGE_SIZE = 256


def download(url, sha256):
    # Wikimedia refuses requests without a User-Agent.
    request = urllib.request.Request(url, headers={"User-Agent": "quak"})
    with urllib.request.urlopen(request) as response:
        data = response.read()
    digest = hashlib.sha256(data).hexdigest()
    if digest != sha256:
        raise SystemExit(f"{url} has sha256 {digest}, not {sha256}")
    return data


def ente_model(name, sha256):
    return onnxruntime.InferenceSession(
        download(f"https://models.ente.com/{name}", sha256),
        providers=["CPUExecutionProvider"],
    )


def run(session, tensor):
    (output,) = session.run(None, {session.get_inputs()[0].name: tensor})
    embedding = output[0]
    return embedding / np.linalg.norm(embedding)


def text_embedding(session, tokenizer, query):
    ids = tokenizer.encode(query).ids
    ids += [0] * (TOKEN_COUNT - len(ids))
    return run(session, np.array([ids], dtype=np.int32))


def image_embedding(session, photo):
    image = ImageOps.exif_transpose(Image.open(io.BytesIO(photo))).convert("RGB")
    # Scale to cover IMAGE_SIZE x IMAGE_SIZE and crop the centre.
    square = ImageOps.fit(image, (IMAGE_SIZE, IMAGE_SIZE), Image.Resampling.BILINEAR)
    pixels = np.asarray(square, dtype=np.float32) / 255
    return run(session, pixels.transpose(2, 0, 1)[np.newaxis])


def numbers(embedding):
    return ", ".join(str(value) for value in embedding)


def write_categories(prompts, embeddings):
    lines = [
        "// Written by bin/clip-keyword-vectors; do not edit by hand.",
        "//",
        "// The Discover categories of Ente's photos app, in their order, from",
        f"// {DISCOVER_PATH} in https://github.com/ente-io/ente",
        f"// at commit {ENTE_COMMIT}. Each has the text",
        f"// embedding Ente's CLIP text model ({TEXT_MODEL})",
        "// gives its query, scaled to length 1. `showVideo` is false for a category",
        "// that leaves out videos.",
        "",
        "export const DISCOVER_CATEGORIES = [",
    ]
    for prompt, embedding in zip(prompts, embeddings):
        lines += [
            "    {",
            f"        title: {json.dumps(prompt['title'])},",
            f"        query: {json.dumps(prompt['query'])},",
            f"        minScore: {json.dumps(prompt['minScore'])},",
            f"        showVideo: {json.dumps(prompt.get('showVideo', True))},",
            f"        textEmbedding: [{numbers(embedding)}],",
            "    },",
        ]
    lines.append("];")
    with open("src/clip-keyword-vectors.ts", "w", encoding="utf-8") as file:
        file.write("\n".join(lines) + "\n")


def write_fixture(embedding):
    with open("test/clip-keywords-sunset.json", "w", encoding="utf-8") as file:
        file.write(
            "{\n"
            f'    "photo": {json.dumps(PHOTO_URL)},\n'
            f'    "sha256": {json.dumps(PHOTO_SHA256)},\n'
            f'    "embedding": [{numbers(embedding)}]\n'
            "}\n"
        )


def main():
    prompts = json.loads(download(DISCOVER_URL, DISCOVER_SHA256))["prompts"]
    tokenizer = Tokenizer.from_str(
        download(TOKENIZER_URL, TOKENIZER_SHA256).decode("utf-8")
    )
    text_model = ente_model(TEXT_MODEL, TEXT_MODEL_SHA256)
    embeddings = [
        text_embedding(text_model, tokenizer, prompt["query"]) for prompt in prompts
    ]
    write_categories(prompts, embeddings)

    image_model = ente_model(IMAGE_MODEL, IMAGE_MODEL_SHA256)
    write_fixture(image_embedding(image_model, download(PHOTO_URL, PHOTO_SHA256)))


if __name__ == "__main__":
    main()