"""Run upstream Pixal3D SLatFlowModel on CPU torch to produce a reference output.

The SLAT flows contain no sparse convolution, so upstream is runnable here once the two
flash-attn kernels are swapped for torch SDPA. That gives a real numerical oracle.
"""
import json
import sys

import numpy as np
import torch
import torch.nn.functional as F

sys.path.insert(0, "/Users/m3ultra/Documents/pixal3d_mrp_mlx/upstream/Pixal3D")

import pixal3d.modules.sparse as sp  # noqa: E402
from pixal3d.modules.sparse.basic import VarLenTensor  # noqa: E402
import pixal3d.modules.attention.modules as AM  # noqa: E402
import pixal3d.modules.sparse.attention.modules as SAM  # noqa: E402


def _dense_sdpa(q, k, v, **kw):
    return F.scaled_dot_product_attention(
        q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
    ).permute(0, 2, 1, 3)


AM.scaled_dot_product_attention = _dense_sdpa


def _flatten_kv(t):
    """-> (feats [T,H,C], list of (start,stop)). Handles varlen and dense [N,L,H,C]."""
    if isinstance(t, VarLenTensor):
        return t.feats, [(s.start, s.stop) for s in t.layout]
    if t.dim() == 4:  # [N, L, H, C] dense -> flatten, one segment per batch item
        n, l = t.shape[0], t.shape[1]
        return t.reshape(n * l, *t.shape[2:]), [(i * l, (i + 1) * l) for i in range(n)]
    return t, [(0, t.shape[0])]


def _sparse_sdpa(*args):
    if len(args) == 1:
        ref = args[0]
        f = ref.feats
        q, k, v = f[:, 0], f[:, 1], f[:, 2]
        lq = lkv = [(s.start, s.stop) for s in ref.layout]
    elif len(args) == 2:
        ref, kv = args
        q, lq = _flatten_kv(ref)
        kvf, lkv = _flatten_kv(kv)
        k, v = kvf[:, 0], kvf[:, 1]
    else:
        ref, k_, v_ = args
        q, lq = _flatten_kv(ref)
        k, lkv = _flatten_kv(k_)
        v, _ = _flatten_kv(v_)
    outs = []
    for (a, b), (c, d) in zip(lq, lkv):
        o = F.scaled_dot_product_attention(
            q[a:b].permute(1, 0, 2)[None],
            k[c:d].permute(1, 0, 2)[None],
            v[c:d].permute(1, 0, 2)[None],
        )
        outs.append(o[0].permute(1, 0, 2))
    out = torch.cat(outs, 0)
    return ref.replace(out) if isinstance(ref, VarLenTensor) else out


SAM.sparse_scaled_dot_product_attention = _sparse_sdpa

from pixal3d.models.structured_latent_flow import SLatFlowModel  # noqa: E402
from safetensors.torch import load_file  # noqa: E402

CK = (
    "/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts/"
    "slat_flow_img2shape_dit_1_3B_512_bf16"
)
cfg = json.load(open(CK + ".json"))
a = dict(cfg["args"])
a["dtype"] = "float32"
m = SLatFlowModel(**a).eval()
m.load_state_dict(
    {k: v.float() for k, v in load_file(CK + ".safetensors").items()}, strict=False
)

rng = np.random.default_rng(0)
N = 64
co = np.concatenate(
    [np.zeros((N, 1), dtype=np.int32), rng.integers(0, 32, (N, 3)).astype(np.int32)], 1
)
co = co[np.lexsort((co[:, 3], co[:, 2], co[:, 1], co[:, 0]))]
fe = rng.standard_normal((N, 32)).astype(np.float32)
g = (rng.standard_normal((1, 37, 1024)) * 0.1).astype(np.float32)
pr = (rng.standard_normal((N, 2048)) * 0.1).astype(np.float32)
t = torch.tensor(np.array([0.5], dtype=np.float32))

with torch.no_grad():
    out = m(
        sp.SparseTensor(torch.tensor(fe), torch.tensor(co)),
        t,
        {
            "global": torch.tensor(g),
            "proj": sp.SparseTensor(torch.tensor(pr), torch.tensor(co)),
        },
    )
o = out.feats.numpy()
print(f"  UPSTREAM slat: {o.shape} mean {o.mean():+.5f} std {o.std():.5f}")
for n, v in [("ref", o), ("co", co), ("fe", fe), ("g", g), ("pr", pr)]:
    np.save(f"/tmp/px_slat_{n}.npy", v)