"""onw-Image runtime v2: conv-form DiT graphs (zit.cdit) on an OpenVINO device; host does the small parts.

    eng = Engine(conv_dir, device="NPU", size=(1024, 1024))
    lat = eng.denoise("prompt", seed=42)

NPU graphs (DiT blocks, text encoder segments, VAE) are compiled once into exported blobs in <conv>/npu_blobs3 and
imported from there (with NPU_TURBO). With onw's NPU kit (zit.npukit) the blobs are weightless: the weights are rebuilt
on the NPU from the model's .bin files at import, so the blobs hold the graphs only (0.1 kept the weights in every
blob: ~36 GB next to a 10 GB model on Panther Lake).
"""
import json, os, time
import numpy as np
import openvino as ov

from . import host, dit, te, cdit, npukit

# (0.1: npu_blobs2 = full blobs, and an OpenVINO CACHE_DIR for the text encoder and VAE; ONW_IMAGE_BLOB_DIR: tests)
BLOB_DIR = os.environ.get("ONW_IMAGE_BLOB_DIR", "npu_blobs3")

CHAT = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
D = cdit.D


CAP_BUCKETS = (64, 256, 512)   # caption token buckets (each needs its own compiled DiT set; Japanese ~1 token/char)


def cap_bucket(n):
    for b in CAP_BUCKETS:
        if n <= b:
            return b
    raise ValueError("prompt too long (max 512 tokens)")


def f16(a):
    return np.ascontiguousarray(a, np.float16)


class HostData(dict):
    """Host-side arrays: <conv>/host/<key>.npy memory-mapped on first use (the 778 MB token table is only touched
    row by row), else everything from host.npz."""

    def __init__(self, conv):
        super().__init__()
        self.dir = os.path.join(conv, "host")
        if not os.path.isdir(self.dir):
            self.update(dict(np.load(os.path.join(conv, "host.npz"))))
            self.dir = None

    def __missing__(self, k):
        if self.dir is None:
            raise KeyError(k)
        v = np.load(os.path.join(self.dir, k.replace(":", "@") + ".npy"), mmap_mode="r")
        self[k] = v
        return v

    def __contains__(self, k):
        return dict.__contains__(self, k) or (self.dir is not None and
                                              os.path.exists(os.path.join(self.dir, k.replace(":", "@") + ".npy")))


def _arr(x):
    """numpy view of a block output (an ov.Tensor with zero-copy chaining, else already numpy)."""
    return x.data if isinstance(x, ov.Tensor) else x


class Engine:
    def __init__(self, conv, device="NPU", size=(1024, 1024), cache_dir=None, log=print, te_device=None, props=None,
                 lazy=None):
        self.conv, self.device, self.log = conv, device, log
        self.lazy = (device == "CPU") if lazy is None else lazy   # CPU checks: compile per use (f32 copies are big)
        self.Hh = HostData(conv)
        self.V = dict(np.load(os.path.join(conv, "dit2_vec.npz")))
        self.core = npukit.core()
        if cache_dir and device != "NPU":               # (NPU graphs: exported blobs, see _npu_request)
            self.core.set_property({"CACHE_DIR": cache_dir})
        self.props = dict(props or {})
        self.export_only = False                        # precompile children: compile + export, no import
        # block outputs stay in the plugin's tensors and feed the next block directly ($ONW_IMAGE_ZEROCOPY=0: copies)
        self.zero_copy = device == "NPU" and os.environ.get("ONW_IMAGE_ZEROCOPY", "1") != "0"
        self._wmodels = {}                              # graph key -> its ov.Model (a weightless import reads it)
        # exported DiT blobs are imported with NPU_TURBO (higher NPU clock: ~15% faster convolutions on Panther Lake);
        # it is a runtime setting at import - passed to compile_model it would change the cache key and recompile
        turbo = os.environ.get("ONW_IMAGE_TURBO", "1") != "0"
        self.import_props = {**self.props, "NPU_TURBO": "YES"} if turbo else dict(self.props)
        self.te_device = te_device or device
        self.H, self.W = int(size[0]), int(size[1])
        self.ht, self.wt = self.H // 16, self.W // 16
        self.n_img = self.ht * self.wt
        self._te, self._blocks, self._bucket, self._models = {}, {}, None, []
        # zit.cdit_u builds (dit2/uniform.json): every block of a kind is the same graph -> one weightless blob per kind
        self.uniform = os.path.exists(os.path.join(conv, "dit2", "uniform.json"))
        # 0.3: uniform blocks are rebuilt in memory in a faster graph form (zit.cdit_u2: same weight constants, no new
        # model files): G4 (default; FlashAttention-2 with the kit's compiler) or G3; $ONW_IMAGE_GRAPH=0.2 keeps the
        # graphs as stored
        gf = os.environ.get("ONW_IMAGE_GRAPH", "g4")
        self.graph_form = gf if self.uniform and gf in ("g3", "g4") else ""
        from tokenizers import Tokenizer
        self.tok = Tokenizer.from_file(os.path.join(conv, "tokenizer.json"))

    # ---------- compile ----------
    def _shapes(self, streams, sfxs=("",)):
        sh = {}
        for s, (h, w) in streams.items():
            n = h * w
            sh[f"x_{s}"] = [1, D, h, w]
            sh[f"mask_{s}"] = [1, 1, 1, n]
            for x in sfxs:
                sh.update({f"{t}_{s}{x}": [1, 1, n, cdit.HD] for t in ("cq", "sq", "ck", "sk")})
        for x in sfxs:
            sh.update({n + x: [1, D, 1, 1] for n in ("m_a", "g_a", "m_f", "g_f")})
        return sh

    def groups(self):
        """[[layer names]] when the main layers are chained in groups (dit2/groups.json), else one per layer."""
        p = os.path.join(self.conv, "dit2", "groups.json")
        return json.load(open(p)) if os.path.exists(p) else [[f"layers.{i}"] for i in range(30)]

    def _graph_name(self, gi, grp):
        return f"group.{gi}" if len(grp) > 1 or grp[0].startswith("group") else grp[0]

    def blob_path(self, b, shapes):
        """Exported NPU blob for graph b at these shapes (<conv>/npu_blobs3/<b>_<HxW>[_<HxW>].blob); uniform builds
        with weightless blobs: one per kind (<kind>.all_...: the 30 main layers, the noise / context refiners)."""
        tag = "_".join(f"{v[2]}x{v[3]}" for k, v in sorted(shapes.items()) if k.startswith("x_"))
        if self.uniform and npukit.weightless() and "." in b:
            b = b.split(".")[0] + ".all"
        form = f"-{self.graph_form}" if self.graph_form else ""
        return os.path.join(self.conv, BLOB_DIR, f"{b}{form}_{tag}.blob")

    def _npu_request(self, path, model, device, cfg=None):
        """Infer request for an NPU graph from its exported blob `path` - compiled and exported first when missing
        (weightless with the kit). Imports carry NPU_TURBO (a runtime setting: given to the compile it would be part
        of the blob). export_only: compile + export, return None."""
        cfg = dict(self.props if cfg is None else cfg)
        if npukit.blob_file(path) is None:
            flash = self.graph_form == "g4" and "_s" not in os.path.basename(path) and not os.path.basename(path).startswith("vae")
            if flash:                                   # (the kit's compiler: FlashAttention-2 for the DiT's SDPA)
                os.environ["ONW_FLASH_SDPA"] = "1"
            try:
                npukit.compile_blob(self.core, model, device, cfg, path, log=self.log)
            finally:
                if flash:
                    os.environ.pop("ONW_FLASH_SDPA", None)
        if self.export_only:
            return None
        self._wmodels[path] = model
        imp = {**cfg, **{k: v for k, v in self.import_props.items() if k == "NPU_TURBO"}}
        return npukit.import_blob(self.core, model, device, imp, path).create_infer_request()

    def _compile_block(self, b, streams, n_sub=1):
        sh = self._shapes(streams, [f"@{j}" for j in range(n_sub)] if n_sub > 1 else ("",))
        xml = os.path.join(self.conv, "dit2", b + ".xml")
        if self.graph_form:
            from . import cdit_u2
            m = cdit_u2.build(xml, cdit_u2.G4 if self.graph_form == "g4" else cdit_u2.G3, self.core)
        else:
            m = self.core.read_model(xml)
        m.reshape(sh)
        if self.device == "NPU":
            return self._npu_request(self.blob_path(b, sh), m, self.device)
        self._models.append(m)       # CPU graphs may reference the model constants
        return self.core.compile_model(m, self.device, self.props).create_infer_request()

    def load_blocks(self, bucket):
        if self._bucket == bucket:
            return
        self._blocks, self._models = {}, []
        t = time.time()
        img, cap = (self.ht, self.wt), (1, bucket)
        self._groups = self.groups()
        self._nsub = {}
        self._streams = {**{f"noise_refiner.{i}": {"img": img} for i in range(2)},
                         **{f"context_refiner.{i}": {"cap": cap} for i in range(2)}}
        for gi, grp in enumerate(self._groups):
            name = self._graph_name(gi, grp)
            self._streams[name] = {"img": img, "cap": cap}
            self._nsub[name] = len(grp) if name.startswith("group") else 1
        if self.lazy:
            self._blocks = {b: None for b in self._streams}
            self._bucket = bucket
            return
        for i in range(2):
            self._blocks[f"noise_refiner.{i}"] = self._compile_block(f"noise_refiner.{i}", {"img": img})
            self._blocks[f"context_refiner.{i}"] = self._compile_block(f"context_refiner.{i}", {"cap": cap})
        for gi, grp in enumerate(self._groups):
            name = self._graph_name(gi, grp)
            self._blocks[name] = self._compile_block(name, {"img": img, "cap": cap}, self._nsub[name])
            self.log(f"  {name} ready ({time.time()-t:.0f}s)")
        self._bucket = bucket

    # ---------- text ----------
    def _te_paths(self):
        d = os.path.join(self.conv, "te_seg")
        if os.path.isdir(d):
            n = len([f for f in os.listdir(d) if f.endswith(".xml")])
            return [os.path.join(d, f"te_{i}.xml") for i in range(n)]
        return [os.path.join(self.conv, "te", "te.xml")]

    def encode_compile(self, S):
        """Compile (or import from the cache) the text encoder segments for prompt bucket S."""
        if S in self._te:
            return
        reqs = []
        for path in self._te_paths():
            m = self.core.read_model(path)
            m.reshape({"emb": [1, S, te.HID], "cos": [1, S, 1, 64], "sin": [1, S, 1, 64], "mask": [1, 1, S, S]})
            if self.te_device == "NPU":
                name = os.path.splitext(os.path.basename(path))[0]
                reqs.append(self._npu_request(os.path.join(self.conv, BLOB_DIR, f"{name}_s{S}.blob"), m, "NPU"))
            else:
                reqs.append(self.core.compile_model(m, self.te_device, self.props).create_infer_request())
            del m
        self._te = {S: reqs}

    def embed(self, ids):
        """Token embeddings [n, 2560] f32: from the INT8 table (te.embed_tokens_q + per-row scale te.embed_tokens_s,
        0.39 GB) when the model has it, else the f16 one (0.78 GB)."""
        if "te.embed_tokens_q" in self.Hh:
            return self.Hh["te.embed_tokens_q"][ids].astype(np.float32) * \
                self.Hh["te.embed_tokens_s"][ids].astype(np.float32)[:, None]
        return self.Hh["te.embed_tokens"][ids].astype(np.float32)

    def encode(self, prompt):
        """Prompt -> caption features [n, 2560] (text encoder segments compiled per prompt bucket, then kept until
        free_text_encoder())."""
        ids = self.tok.encode(CHAT.format(prompt)).ids[:512]
        S = next(b for b in (64, 128, 256, 512) if len(ids) <= b)
        self.encode_compile(S)
        x = np.zeros((1, S, te.HID), np.float32)
        x[0, :len(ids)] = self.embed(ids)
        cos, sin, mask = te.host_inputs(ids, S)
        for rq in self._te[S]:
            rq.infer({"emb": x, "cos": cos, "sin": sin, "mask": mask})
            x = rq.get_output_tensor(0).data.astype(np.float32).copy()
        return x[0, :len(ids)].copy()

    def free_text_encoder(self):
        self._te = {}
        import gc
        gc.collect()

    # ---------- DiT ----------
    def _run(self, b, feed):
        return self._finish(b, self._start(b, feed))

    def _start(self, b, feed):
        """Start block b on its own infer request (async on the NPU: the host prepares the next block meanwhile)."""
        rq = self._blocks[b]
        if rq is None:
            self._models = []
            rq = self._compile_block(b, self._streams[b], self._nsub.get(b, 1))
        if self.zero_copy:
            # NPU: a block's output tensor (the plugin's own device-visible buffer) goes straight into the next block -
            # no copy out to numpy and back in (~60 MB per block at 1024); numpy inputs are shared, not copied
            rq.start_async(feed, share_inputs=True)
        else:
            rq.infer(feed)
        return rq

    def _finish(self, b, rq):
        if self.zero_copy:
            rq.wait()
            out = {o.get_any_name(): rq.get_tensor(o) for o in rq.model_outputs}
        else:
            out = {o.get_any_name(): rq.get_tensor(o).data.copy() for o in rq.model_outputs}
        if self.lazy:
            self._models = []
        return out

    def _vecs(self, b, mods=None):
        V = self.V
        a_in, a_out, f_in, f_out = (V[f"{b}:{k}"] for k in ("a_in", "a_out", "f_in", "f_out"))
        if mods is not None:
            s1, g1, s2, g2 = (m.ravel() for m in mods)
            a_in, a_out, f_in, f_out = a_in * s1, a_out * g1, f_in * s2, f_out * g2
        out = {n: f16(v.reshape(1, D, 1, 1)) for n, v in zip(("m_a", "g_a", "m_f", "g_f"), (a_in, a_out, f_in, f_out))}
        for lin in ("qkv", "out", "w1", "w3", "w2"):    # zit.cdit_u shared weights: the caption stream's guards
            v = V.get(f"{b}:p_{lin}")
            if v is not None:
                out[f"p_{lin}"] = f16(v.reshape(1, -1, 1, 1))
        return out

    def _sublayers(self, b):
        if b.startswith("group"):
            return self._groups[int(b.split(".")[1])]
        return [b]

    def _tables(self, b, s, ang):
        cq, sq, ck, sk = cdit.rope_tables(ang, self.V[f"{b}:q_g"], self.V[f"{b}:k_g"])
        return {f"cq_{s}": cq, f"sq_{s}": sq, f"ck_{s}": ck, f"sk_{s}": sk}

    def prepare(self, cap_feats):
        Hh = self.Hh
        n = len(cap_feats)
        cp = host.cap_padded_len(n)
        bucket = cap_bucket(cp)
        self.load_blocks(bucket)
        c = dit.rms_np(cap_feats, Hh["cap_embedder.0.weight"], 1e-5)
        c = c @ Hh["cap_embedder.1.weight"].T + Hh["cap_embedder.1.bias"]
        cap = np.empty((bucket, D), np.float32)
        cap[:n] = c
        cap[n:] = Hh["cap_pad_token"]                  # pad and bucket rows (bucket rows: masked keys, never zero)
        img_ids, cap_ids = host.position_ids(n, self.ht, self.wt)
        ids_b = np.zeros((bucket, 3), np.int64)
        ids_b[:cp] = cap_ids
        a_img, a_cap = host.rope_angles(img_ids), host.rope_angles(ids_b)
        mask_cap = np.zeros((1, 1, 1, bucket), np.float16)
        mask_cap[..., cp:] = cdit.NEG
        mask_img = np.zeros((1, 1, 1, self.n_img), np.float16)
        self.cp = cp
        self.const = {}
        self._fbc, self.fbcache_skipped, self.fbcache_log = None, 0, []     # (First Block Cache: per image)
        for b in self._blocks:
            f = {}
            subs = self._sublayers(b)
            for j, lb in enumerate(subs):
                sfx = f"@{j}" if len(subs) > 1 or b.startswith("group") else ""
                if not b.startswith("context"):
                    f.update({k + sfx: v for k, v in self._tables(lb, "img", a_img).items()})
                if not b.startswith("noise"):
                    f.update({k + sfx: v for k, v in self._tables(lb, "cap", a_cap).items()})
            if not b.startswith("context"):
                f["mask_img"] = mask_img
            if not b.startswith("noise"):
                f["mask_cap"] = mask_cap
            self.const[b] = f
        k = cdit.RESID_K
        x = f16((cap / k["context_refiner"]).T.reshape(1, D, 1, bucket))
        for i in range(2):
            b = f"context_refiner.{i}"
            x = self._run(b, dict(self.const[b], x_cap=x, **self._vecs(b)))["out_cap"]
        self.cap = f16(_arr(x).astype(np.float32) * (k["context_refiner"] / k["layers"]))

    def dit(self, lat, t, trace=None):
        Hh, k = self.Hh, cdit.RESID_K
        temb = dit.t_embed(Hh, host.timestep_freq(t))
        xi = host.patchify(lat) @ Hh["all_x_embedder.2-1.weight"].T + Hh["all_x_embedder.2-1.bias"]
        xi = f16((xi / k["noise_refiner"]).T.reshape(1, D, self.ht, self.wt))
        for i in range(2):
            b = f"noise_refiner.{i}"
            xi = self._run(b, dict(self.const[b], x_img=xi, **self._vecs(b, dit.modulation(Hh, b + ".", temb))))["out_img"]
            if trace is not None:
                trace[b] = _arr(xi)[0].reshape(D, -1).T.astype(np.float32) * k["noise_refiner"]
        xi = f16(_arr(xi).astype(np.float32) * (k["noise_refiner"] / k["layers"]))
        xc = self.cap
        # First Block Cache ($ONW_IMAGE_FBCACHE = threshold, e.g. 0.1): after the first main block, when its residual
        # changed little since the previous step, the other blocks are skipped and their previous residual reused
        thr = float(os.environ.get("ONW_IMAGE_FBCACHE", "0") or 0)
        fb = getattr(self, "_fbc", None) if thr > 0 else None
        pending = None                 # (block, request) running on the NPU while the host prepares the next block
        x_first_in = xi
        skipped = False
        for gi, grp in enumerate(self._groups + [None]):
            vecs = {}
            if grp is not None:
                b = self._graph_name(gi, grp)
                for j, lb in enumerate(grp):
                    sfx = f"@{j}" if b.startswith("group") else ""
                    vecs.update({n + sfx: v for n, v in self._vecs(lb, dit.modulation(Hh, lb + ".", temb)).items()})
            if pending is not None:
                r = self._finish(*pending)
                xi, xc = r["out_img"], r["out_cap"]
                pending = None
                if trace is not None:
                    trace[self._groups[gi - 1][-1]] = np.concatenate([_arr(xi)[0].reshape(D, -1).T, _arr(xc)[0].reshape(D, -1).T[:self.cp]]).astype(np.float32) * k["layers"]
                if thr > 0 and gi == 1:
                    r1 = _arr(xi).astype(np.float32) - _arr(x_first_in).astype(np.float32)
                    x_first = _arr(xi).astype(np.float32)
                    if fb is not None and fb["r1"].shape == r1.shape:
                        change = float(np.abs(r1 - fb["r1"]).mean() / max(np.abs(fb["r1"]).mean(), 1e-12))
                        self.fbcache_log = getattr(self, "fbcache_log", []) + [change]
                        if change < thr:
                            xi = f16(x_first + fb["rest"])
                            self.fbcache_skipped = getattr(self, "fbcache_skipped", 0) + 1
                            skipped = True     # (the reference residuals stay: errors do not drift)
                            break
            if grp is None:
                break
            pending = (b, self._start(b, dict(self.const[b], x_img=xi, x_cap=xc, **vecs)))
        if thr > 0 and not skipped:
            self._fbc = {"r1": r1, "rest": _arr(xi).astype(np.float32) - x_first}
        out = dit.final_layer(Hh, _arr(xi)[0].reshape(D, -1).T.astype(np.float32) * k["layers"], temb)
        return host.unpatchify(out, self.H // 8, self.W // 8)

    # ---------- VAE ----------
    # On the NPU the VAE decodes in tiles of VAE_TILE x VAE_TILE latents (512x512 pixels) overlapping by >= VAE_OVERLAP,
    # blended linearly: ONE compiled graph for every image size, and a 1024x1024 graph alone was 189 MB of NPU task
    # descriptors and a 14-minute compile (512: 52 MB, 99 s). ONW_IMAGE_VAE_TILE=0: the whole image at once.
    VAE_TILE, VAE_OVERLAP = 64, 16

    @staticmethod
    def _tile_starts(n, t, ov):
        """Evenly spread tile starts covering [0, n) with tiles of t and overlaps >= ov."""
        if n <= t:
            return [0]
        k = -(-(n - t) // (t - ov)) + 1
        return [round(i * (n - t) / (k - 1)) for i in range(k)]

    def _vae_graph(self, dev, h, w):
        key = (dev, h, w)
        if getattr(self, "_vae_key", None) != key:
            m = self.core.read_model(os.path.join(self.conv, "vae", "vae_decoder.xml"))
            m.reshape({0: [1, 16, h, w]})
            self._vae_m = m
            if dev == "NPU":
                bp = os.path.join(self.conv, BLOB_DIR, f"vae_{h}x{w}.blob")
                self._vae = self._npu_request(bp, m, "NPU", {"INFERENCE_PRECISION_HINT": "f16"})
            else:
                self._vae = self.core.compile_model(m, dev).create_infer_request()
            self._vae_key = key if self._vae is not None else None
        return self._vae

    def decode(self, lat, device=None):
        """latent [16,H/8,W/8] -> PIL image (Flux VAE decoder IR in <conv>/vae)."""
        from PIL import Image
        dev = device or getattr(self, "vae_device", None) or self.device
        _, lh, lw = lat.shape
        t = self.VAE_TILE
        tiled = os.environ.get("ONW_IMAGE_VAE_TILE", "1")          # (force: tiles on any device - checks)
        if (dev != "NPU" and tiled != "force") or tiled == "0" or (lh <= t and lw <= t):
            rq = self._vae_graph(dev, lh, lw)
            if rq is None:                              # (export_only)
                return None
            rq.infer({0: lat[None].astype(np.float32)})
            img = rq.get_output_tensor(0).data[0]
        else:
            rq = self._vae_graph(dev, t, t)
            if rq is None:
                return None
            ys, xs = self._tile_starts(lh, t, self.VAE_OVERLAP), self._tile_starts(lw, t, self.VAE_OVERLAP)
            acc = np.zeros((3, lh * 8, lw * 8), np.float32)
            wsum = np.zeros((1, lh * 8, lw * 8), np.float32)

            def ramp(starts, i, n):
                """Weights along one axis for tile i: linear ramps over the overlaps with its neighbours."""
                w = np.ones(t * 8, np.float32)
                s = starts[i]
                if i > 0:
                    ov = (starts[i - 1] + t - s) * 8
                    w[:ov] = np.minimum(w[:ov], (np.arange(ov) + 0.5) / ov)
                if i < len(starts) - 1:
                    ov = (s + t - starts[i + 1]) * 8
                    w[t * 8 - ov:] = np.minimum(w[t * 8 - ov:], (ov - np.arange(ov) - 0.5) / ov)
                return w

            for iy, y0 in enumerate(ys):
                wy = ramp(ys, iy, lh)
                for ix, x0 in enumerate(xs):
                    wx = ramp(xs, ix, lw)
                    rq.infer({0: np.ascontiguousarray(lat[None, :, y0:y0 + t, x0:x0 + t], np.float32)})
                    tile = rq.get_output_tensor(0).data[0].astype(np.float32)
                    w2 = wy[:, None] * wx[None, :]
                    acc[:, y0 * 8:(y0 + t) * 8, x0 * 8:(x0 + t) * 8] += tile * w2
                    wsum[:, y0 * 8:(y0 + t) * 8, x0 * 8:(x0 + t) * 8] += w2
            img = acc / np.maximum(wsum, 1e-6)
        img = ((np.clip(img, -1, 1) + 1) * 127.5).round().astype(np.uint8).transpose(1, 2, 0)
        return Image.fromarray(img)

    def denoise(self, prompt, seed=42, steps=9, latents=None, callback=None):
        t0 = time.time()
        cap = self.encode(prompt)
        self.free_text_encoder()          # 16 GB machines: the text encoder and the DiT do not fit together
        t1 = time.time()
        self.prepare(cap)
        if latents is None:
            latents = np.random.default_rng(seed).standard_normal((16, self.H // 8, self.W // 8)).astype(np.float32)
        lat = latents.astype(np.float32)
        sig = host.sigmas(steps)
        for i in range(steps):
            lat = host.euler(lat, self.dit(lat, host.model_t(sig[i])), sig[i], sig[i + 1])
            if callback:
                callback(i, steps, lat)
        self.timing = {"text": t1 - t0, "denoise": time.time() - t1}
        return lat