# SPDX-License-Identifier: Apache-2.0
"""iris-3b Turbo text-to-image on LiteRT CompiledModel: the Python reference pipeline.

A prompt becomes a 768x768 image through the 20 graphs in ../model:

    tokenize (Qwen3-VL byte-level BPE) -> bf16 embedding lookup
    -> 12 text encoder graphs (one per hidden-state tap) -> window rows x mask
    -> adapter (once per prompt)
    -> 4 steps x (dual0..2 -> [text; image] -> single0..2 -> image rows -> pixel stage -> sampler)
    -> clamp -> patch-major pixels to a CHW image

It reads the same files and runs the same host-side steps as the Android engine in ../android
(IrisPipeline.kt, HostOps.kt, PromptTokens.kt, QwenTokenizer.kt). Every graph runs on the CPU
(XNNPACK); the phone runs the same graphs on its GPU at FP32. A seed's starting noise is
java.util.Random(seed).nextGaussian(), as in the Android app.

    pip install -r requirements.txt
    python iris3b_litert.py --model-dir ../model --prompt "a red apple on a wooden table" --out apple.png
"""

import argparse
import json
import math
import os
import resource
import sys
import time
import unicodedata
from pathlib import Path

import numpy as np
import regex

CONFIG_FILE = "iris3b_config.json"

SEQ = 334
PREFIX = 34
WINDOW = 300
TE_HIDDEN = 2560
TE_HEADS = 32
TAPS = 12
ADAPTER_HEADS = 20
DIT_HIDDEN = 2560
IMAGE_SIZE = 768
PATCH = 16
GRID = IMAGE_SIZE // PATCH
PATCHES = GRID * GRID
PIXELS = PATCH * PATCH
CHANNELS = 3
IMAGE_FLOATS = CHANNELS * IMAGE_SIZE * IMAGE_SIZE
MASK_NEG = np.float32(-30000.0)
PAD_TOKEN_ID = 151643

# The Unicode White_Space set, written out (the Android port does the same so that every regex
# engine agrees on it), and tokenizer.json's Split pattern with `\s` and `(?i:...)` written out.
_WS = "\t\n\x0b\x0c\r \x85\xa0  - 

  　"
_PRE_TOKENIZER = regex.compile(
    r"'(?:[sSſ]|[tT]|[rR][eE]|[vV][eE]|[mM]|[lL][lL]|[dD])"
    r"|[^\r\n\p{L}\p{N}]?\p{L}+"
    r"|\p{N}"
    r"| ?[^" + _WS + r"\p{L}\p{N}]+[\r\n]*"
    r"|[" + _WS + r"]*[\r\n]+"
    r"|[" + _WS + r"]+(?![^" + _WS + r"])"
    r"|[" + _WS + r"]+"
)


class QwenTokenizer:
    """Byte-level BPE of the Qwen3-VL-4B-Instruct text encoder (a port of QwenTokenizer.kt).

    Added tokens are split out first (leftmost, then longest), the text between them is NFC
    normalized, pre-tokenized with the Split pattern, byte-level encoded (GPT-2 bytes_to_unicode)
    and merged pair by pair in rank order.
    """

    def __init__(self, files):
        vocab = files["qwen_vocab.txt"].split("\n")
        if vocab[-1] == "":
            vocab.pop()
        self._tokens = vocab
        self._token_to_id = {token: i for i, token in enumerate(vocab)}
        self._ranks = {}
        for rank, line in enumerate(files["qwen_merges.txt"].split("\n")):
            space = line.find(" ")
            if space < 0:
                continue
            left = self._token_to_id.get(line[:space])
            right = self._token_to_id.get(line[space + 1:])
            if left is not None and right is not None:
                self._ranks[(left, right)] = rank
        added = []
        for line in files["qwen_added.txt"].split("\n"):
            if line:
                tab = line.rindex("\t")
                added.append((line[:tab], int(line[tab + 1:])))
        added.sort(key=lambda t: -len(t[0]))
        self._added = {}
        for token, token_id in added:
            self._added.setdefault(token[0], []).append((token, token_id))
        printable = set(range(ord("!"), ord("~") + 1)) | set(range(ord("¡"), ord("¬") + 1)) | set(
            range(ord("®"), ord("ÿ") + 1))
        byte_to_char, extra = [], 0
        for byte in range(256):
            if byte in printable:
                byte_to_char.append(chr(byte))
            else:
                byte_to_char.append(chr(256 + extra))
                extra += 1
        self._byte_ids = [self._token_to_id[c] for c in byte_to_char]
        self._cache = {}

    def encode(self, text):
        """Token ids of raw text, no template and no special tokens added."""
        ids, start, i = [], 0, 0
        while i < len(text):
            hit = next((a for a in self._added.get(text[i], ()) if text.startswith(a[0], i)), None)
            if hit is None:
                i += 1
                continue
            if start < i:
                self._encode_segment(text[start:i], ids)
            ids.append(hit[1])
            i += len(hit[0])
            start = i
        if start < len(text):
            self._encode_segment(text[start:], ids)
        return ids

    def _encode_segment(self, segment, out):
        text = unicodedata.normalize("NFC", segment)
        position = 0
        while position < len(text):
            match = _PRE_TOKENIZER.search(text, position)
            if match is None:
                out.extend(self._merge(text[position:]))
                break
            if match.start() > position:
                out.extend(self._merge(text[position:match.start()]))
            out.extend(self._merge(match.group()))
            position = match.end()

    def _merge(self, piece):
        cached = self._cache.get(piece)
        if cached is not None:
            return cached
        symbols = [self._byte_ids[b] for b in piece.encode("utf-8")]
        while len(symbols) > 1:
            best_rank, best_index = None, -1
            for i in range(len(symbols) - 1):
                rank = self._ranks.get((symbols[i], symbols[i + 1]))
                if rank is not None and (best_rank is None or rank < best_rank):
                    best_rank, best_index = rank, i
            if best_index < 0:
                break
            merged = self._tokens[symbols[best_index]] + self._tokens[symbols[best_index + 1]]
            symbols[best_index:best_index + 2] = [self._token_to_id[merged]]
        self._cache[piece] = symbols
        return symbols


class PromptTokens:
    """The prompt layout and the bf16 embedding lookup (a port of PromptTokens.kt).

    System-prompt prefix (34 ids) + caption cut to 295 ids + assistant-turn suffix (5 ids),
    right-padded with <|endoftext|> to 334; the real tokens are prefix + caption + suffix.
    """

    def __init__(self, files, embed_path):
        self.tokenizer = QwenTokenizer(files)
        self.prefix = [int(v) for v in files["prompt_prefix_ids.txt"].strip().split(",")]
        self.suffix = [int(v) for v in files["prompt_suffix_ids.txt"].strip().split(",")]
        self.table = np.memmap(embed_path, dtype="<u2", mode="r").reshape(-1, TE_HIDDEN)

    def encode(self, caption):
        """(ids [334], real token count, caption tokens before the cut)."""
        caption_ids = self.tokenizer.encode(caption)
        n = min(len(caption_ids), WINDOW - len(self.suffix))
        ids = self.prefix + caption_ids[:n] + self.suffix
        real = len(ids)
        return np.array(ids + [PAD_TOKEN_ID] * (SEQ - real), np.int64), real, len(caption_ids)

    def embed(self, ids):
        """inputs_embeds [334, 2560] fp32 from the bf16 table (exact: bf16 -> fp32 is a 16-bit shift)."""
        return (self.table[ids].astype(np.uint32) << 16).view(np.float32)


def te_mask(real):
    """Additive causal mask [32, 334, 334]: 0 where key <= query and the key is real."""
    query = np.arange(SEQ)[:, None]
    key = np.arange(SEQ)[None, :]
    plane = np.where((key <= query) & (key < real), np.float32(0), MASK_NEG).astype(np.float32)
    return np.ascontiguousarray(np.broadcast_to(plane, (TE_HEADS, SEQ, SEQ)))


def window_mask(real):
    """1 on the caption + suffix rows of the 300-row window."""
    return (PREFIX + np.arange(WINDOW) < real).astype(np.float32)


def adapter_mask_bias(wmask):
    """Adapter key mask [20, 300, 300]: a query sees the real keys and itself."""
    keep = (wmask[None, :] != 0) | np.eye(WINDOW, dtype=bool)
    plane = np.where(keep, np.float32(0), MASK_NEG).astype(np.float32)
    return np.ascontiguousarray(np.broadcast_to(plane, (ADAPTER_HEADS, WINDOW, WINDOW)))


def sampler_step(x, v, sigma, next_sigma):
    """Turbo few-step sampler in fp32: clean = x - sigma v, then x at the next sigma (if any)."""
    clean = x - sigma * v
    if next_sigma is None:
        return x, clean
    implied = (x - (np.float32(1) - sigma) * clean) / sigma
    return (np.float32(1) - next_sigma) * clean + next_sigma * implied, clean


def to_image(x_pix):
    """Patch-major pixels [2304, 256, 3] -> CHW image [3, 768, 768]."""
    return x_pix.reshape(GRID, GRID, PATCH, PATCH, CHANNELS).transpose(4, 0, 2, 1, 3).reshape(
        CHANNELS, IMAGE_SIZE, IMAGE_SIZE)


def seed_noise(seed):
    """Standard normal noise [IMAGE_FLOATS] from java.util.Random(seed).nextGaussian(), as the app."""
    mult, add, mask = 0x5DEECE66D, 0xB, (1 << 48) - 1
    unit = 2.0 ** -53
    state = (seed ^ mult) & mask
    out = []
    log, sqrt = math.log, math.sqrt
    while len(out) < IMAGE_FLOATS:
        while True:
            state = (state * mult + add) & mask
            hi = state >> 22
            state = (state * mult + add) & mask
            v1 = 2 * (((hi << 27) + (state >> 21)) * unit) - 1
            state = (state * mult + add) & mask
            hi = state >> 22
            state = (state * mult + add) & mask
            v2 = 2 * (((hi << 27) + (state >> 21)) * unit) - 1
            s = v1 * v1 + v2 * v2
            if 0 < s < 1:
                break
        multiplier = sqrt(-2 * log(s) / s)
        out.append(v1 * multiplier)
        out.append(v2 * multiplier)
    return np.array(out[:IMAGE_FLOATS], np.float64).astype(np.float32)


def save_png(image, path):
    """Clamped CHW image in [-1, 1] -> 8-bit RGB PNG (the app's rounding)."""
    from PIL import Image

    level = (image.astype(np.float32) + np.float32(1)) * np.float32(127.5) + np.float32(0.5)
    rgb = np.clip(np.nan_to_num(level, nan=0.0), 0, 255).astype(np.uint8).transpose(1, 2, 0)
    Image.fromarray(np.ascontiguousarray(rgb)).save(path)


class _Graph:
    """One graph file on LiteRT CompiledModel (CPU), driven as its config entry says."""

    def __init__(self, path, spec, threads, environment):
        from ai_edge_litert.compiled_model import CompiledModel, CpuOptions, HardwareAccelerator, Options

        options = Options(hardware_accelerators=HardwareAccelerator.CPU,
                          cpu_options=CpuOptions(num_threads=threads))
        self.model = CompiledModel.from_file(str(path), options=options, environment=environment)
        key = spec["signature"]
        self.index = self.model.get_signature_index(key)
        names = self.model.get_signature_list()[key]
        self.inputs = self.model.create_input_buffers(self.index)
        self.outputs = self.model.create_output_buffers(self.index)
        self.feed = [names["inputs"].index(name) for name in spec["inputs"]]
        self.read = [(names["outputs"].index(name), count) for name, count in spec["outputs"]]

    def run(self, feeds):
        for position, values in zip(self.feed, feeds):
            self.inputs[position].write(np.ascontiguousarray(values, np.float32))
        self.model.run_by_index(self.index, self.inputs, self.outputs)
        return [self.outputs[position].read(count, np.float32) for position, count in self.read]

    def close(self):
        self.inputs = self.outputs = None
        self.model.close()
        self.model = None


class Iris3B:
    """The pipeline over a model folder (the 24 files of ../model)."""

    def __init__(self, model_dir, threads=None, keep_dit=False):
        self.dir = Path(model_dir)
        self.config = json.loads((self.dir / CONFIG_FILE).read_text(encoding="utf-8"))
        graphs = self.config["graphs"]
        self.te_specs = graphs["te"]
        self.adapter_spec = graphs["adapter"]
        self.dual = list(zip(graphs["dual"], graphs["dual_mod_rows"]))
        self.single = list(zip(graphs["single"], graphs["single_mod_rows"]))
        self.pixel_spec = graphs["pixel"]
        self._host_file = self.dir / self.config["host"]["file"]
        self._host_index = self.config["host"]["tensors"]
        self._host = {}
        files = json.loads((self.dir / self.config["tokenizer"]).read_text(encoding="utf-8"))["files"]
        self.tokens = PromptTokens(files, self.dir / self.config["embed"])
        self.threads = threads or os.cpu_count()
        self.keep_dit = keep_dit
        self._environment = None
        self._resident = {}
        self.timeline = []

    def host(self, name):
        """A host tensor (rotary tables, modulation rows, sigmas, reference noise ...) by name."""
        if name not in self._host:
            t = self._host_index[name]
            self._host[name] = np.fromfile(self._host_file, dtype="<f4", count=t["count"], offset=t["offset"])
        return self._host[name]

    def reference_noise(self):
        """The upstream reference run's starting noise (torch.randn, seed 0), patch-major."""
        return self.host("noise_pix").copy()

    def generate(self, prompt, noise):
        """Prompt + noise [IMAGE_FLOATS] -> clamped CHW image [3, 768, 768] in [-1, 1]."""
        self.timeline = []
        t0 = time.perf_counter()
        ids, real, caption_tokens = self.tokens.encode(prompt)
        self.real_len, self.caption_tokens = real, caption_tokens
        hidden = self.tokens.embed(ids)
        mask, wmask = te_mask(real), window_mask(real)
        cos, sin = self.host("te_rope_cos"), self.host("te_rope_sin")
        te = np.zeros((WINDOW, TAPS, TE_HIDDEN), np.float32)
        for tap, spec in enumerate(self.te_specs):
            hidden = self._once(spec, [hidden, mask, cos, sin])[0]
            te[:, tap, :] = hidden.reshape(SEQ, TE_HIDDEN)[PREFIX:] * wmask[:, None]
        y0 = self._dit(self.adapter_spec, [te, adapter_mask_bias(wmask), self.host("y_pos")])[0]

        sigmas = self.host("sigmas")
        cos_j, sin_j = self.host("cos_j"), self.host("sin_j")
        cos_i, sin_i, pix_pos = self.host("cos_i"), self.host("sin_i"), self.host("pix_pos")
        x = np.asarray(noise, np.float32).reshape(-1).copy()
        assert x.size == IMAGE_FLOATS, x.size
        clean = None
        for step in range(sigmas.size):
            mods = self.host(f"mods_s{step}").reshape(-1, DIT_HIDDEN)
            image, text = x, y0
            for spec, (first, end) in self.dual:
                outs = self._dit(spec, [image, text, cos_j, sin_j, mods[first:end]])
                image = _by_size(outs, PATCHES * DIT_HIDDEN)
                text = _by_size(outs, WINDOW * DIT_HIDDEN)
            joint = np.concatenate([text, image])
            for spec, (first, end) in self.single:
                joint = self._dit(spec, [joint, cos_j, sin_j, mods[first:end]])[0]
            v = self._dit(self.pixel_spec, [joint[WINDOW * DIT_HIDDEN:], self.host(f"t_emb_s{step}"), x, pix_pos,
                                            cos_i, sin_i])[0]
            next_sigma = sigmas[step + 1] if step + 1 < sigmas.size else None
            x, clean = sampler_step(x, v, sigmas[step], next_sigma)
            print(f"step {step + 1}/{sigmas.size} done at {time.perf_counter() - t0:.1f} s", file=sys.stderr)
        for graph in self._resident.values():
            graph.close()
        self._resident = {}
        self.wall_s = time.perf_counter() - t0
        return to_image(np.clip(clean, -1, 1))

    def _env(self):
        if self._environment is None:
            import ai_edge_litert.compiled_model as cm

            self._environment = cm.Environment.create(
                options=cm.EnvironmentOptions(runtime_path=os.path.dirname(os.path.abspath(cm.__file__))))
        return self._environment

    def _load(self, spec):
        t0 = time.perf_counter()
        graph = _Graph(self.dir / spec["file"], spec, self.threads, self._env())
        self.timeline.append({"graph": spec["role"], "event": "load", "s": time.perf_counter() - t0})
        return graph

    def _run(self, graph, spec, feeds):
        t0 = time.perf_counter()
        outs = graph.run(feeds)
        self.timeline.append({"graph": spec["role"], "event": "run", "s": time.perf_counter() - t0})
        return outs

    def _once(self, spec, feeds):
        graph = self._load(spec)
        try:
            return self._run(graph, spec, feeds)
        finally:
            graph.close()

    def _dit(self, spec, feeds):
        if not self.keep_dit:
            return self._once(spec, feeds)
        graph = self._resident.get(spec["role"])
        if graph is None:
            graph = self._resident[spec["role"]] = self._load(spec)
        return self._run(graph, spec, feeds)


def _by_size(outputs, count):
    hits = [o for o in outputs if o.size == count]
    assert len(hits) == 1, [o.size for o in outputs]
    return hits[0]


def main():
    here = Path(__file__).resolve().parent
    ap = argparse.ArgumentParser(description="iris-3b Turbo text-to-image on LiteRT CompiledModel (CPU)")
    ap.add_argument("--model-dir", default=str(here.parent / "model"), help="folder with the 24 model files")
    ap.add_argument("--prompt", required=True)
    ap.add_argument("--seed", type=int, default=0, help="java.util.Random seed of the starting noise (as the app)")
    ap.add_argument("--reference-noise", action="store_true",
                    help="use the upstream reference run's noise from the host pack instead of --seed")
    ap.add_argument("--out", default="iris3b.png", help="PNG to write")
    ap.add_argument("--npy", help="also write the clamped CHW image [1, 3, 768, 768] float32 here")
    ap.add_argument("--threads", type=int, default=None, help="XNNPACK threads (default: all cores)")
    ap.add_argument("--keep-dit", action="store_true",
                    help="keep the 8 DiT graphs loaded across the 4 steps instead of loading each graph for "
                         "every use (uses more memory)")
    ap.add_argument("--report", help="write timings and memory as JSON here")
    args = ap.parse_args()

    pipe = Iris3B(args.model_dir, threads=args.threads, keep_dit=args.keep_dit)
    noise = pipe.reference_noise() if args.reference_noise else seed_noise(args.seed)
    image = pipe.generate(args.prompt, noise)
    save_png(image, args.out)
    if args.npy:
        np.save(args.npy, image[None].astype(np.float32))
    nonfinite = int(np.size(image) - np.count_nonzero(np.isfinite(image)))
    rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    peak_bytes = rss if sys.platform == "darwin" else rss * 1024
    print(f"{args.out}: {pipe.wall_s:.1f} s, real tokens {pipe.real_len}, caption tokens {pipe.caption_tokens}, "
          f"non-finite {nonfinite}, peak RSS {peak_bytes / 2**30:.2f} GiB", file=sys.stderr)
    if args.report:
        loads = sum(e["s"] for e in pipe.timeline if e["event"] == "load")
        runs = sum(e["s"] for e in pipe.timeline if e["event"] == "run")
        Path(args.report).write_text(json.dumps({
            "prompt": args.prompt, "seed": None if args.reference_noise else args.seed,
            "noise": "reference" if args.reference_noise else "java.util.Random", "threads": pipe.threads,
            "keep_dit": args.keep_dit, "wall_s": pipe.wall_s, "load_s": loads, "run_s": runs,
            "real_len": pipe.real_len, "caption_tokens": pipe.caption_tokens, "nonfinite": nonfinite,
            "peak_rss_bytes": peak_bytes, "timeline": pipe.timeline}, indent=1))
    return 0 if nonfinite == 0 else 1


if __name__ == "__main__":
    sys.exit(main())