#!/usr/bin/env python3
"""Reference dumps for the shape-SLAT stage ("512" pipeline type, geometry only).

Produces dumps/reference_slat.gguf with:

  coords            [L, 4] f32    active voxels at 32^3 (from the SS stage)
  slat_noise        [L, 32]       fixed sampling noise (seed 4321)
  flow_t500_out     [L, 32]       SLAT flow forward at t=500 (f32)
  slat              [L, 32]       full 12-step sampler output, denormalized
  slat_mean/std     [32]          shape_slat_normalization
  lvl{i}.in_coords  [L_i, 4]      decoder level i active voxels
  lvl{i}.pre_up     [L_i, C_i]    features after level i's ConvNeXt blocks
  lvl{i}.subdiv     [L_i, 8]      subdivision logits of level i's up-block
  out7              [L_4, 7]      decoder output (pre split/sigmoid)
  out_coords        [L_4, 4]

The coords come from decoding tests/ss_sample_ref.bin's reference latent, so
the same scaffold is reproducible from the validated C++ SS stages.

Run inside the container: see scripts/refgen.sh.
"""

import argparse
import json
import os
import struct
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ref_common  # noqa: E402

ref_common.setup()  # sdpa sparse attention + pure-torch sparse conv

import numpy as np  # noqa: E402
import torch  # noqa: E402
import torch.nn.functional as F  # noqa: E402


def load_ss_sample_latent(path):
    """Read the z_s reference produced by tests/ref_ss_sample.py (SSSAMP01)."""
    with open(path, "rb") as f:
        assert f.read(8) == b"SSSAMP01", "bad magic"
        R, cin, lkv, cctx, steps = struct.unpack("<5i", f.read(20))
        f.read(4 * 3)  # gs, rescale, rescale_t
        f.read(8)      # seed
        n = cin * R * R * R
        f.seek(-(n * 4), os.SEEK_END)
        z = np.frombuffer(f.read(n * 4), dtype="<f4").reshape(1, cin, R, R, R)
    return torch.from_numpy(z.copy())


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--models", default=os.path.join(ref_common.MODELS, "TRELLIS.2-4B", "ckpts"))
    ap.add_argument("--ss-dec", default=os.path.join(ref_common.MODELS, "TRELLIS-image-large",
                                                     "ckpts", "ss_dec_conv3d_16l8_fp16"))
    ap.add_argument("--dinodata", default=os.path.join(ref_common.DUMPS, "fixture.dinodata"))
    ap.add_argument("--ss-latent", default=os.path.join(ref_common.REPO, "tests", "ss_sample_ref.bin"))
    ap.add_argument("--pipeline-json", default=os.path.join(ref_common.MODELS, "TRELLIS.2-4B", "pipeline.json"))
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    ap.add_argument("--t", type=float, default=500.0)
    ap.add_argument("--seed", type=int, default=4321)
    ap.add_argument("--resolution", type=int, default=512)
    ap.add_argument("--out", default=os.path.join(ref_common.DUMPS, "reference_slat.gguf"))
    args = ap.parse_args()

    from safetensors.torch import load_file
    from trellis2.models.sparse_structure_vae import SparseStructureDecoder
    from trellis2.models.structured_latent_flow import SLatFlowModel
    from trellis2.models.sc_vaes.fdg_vae import FlexiDualGridVaeDecoder
    from trellis2.pipelines.samplers import FlowEulerGuidanceIntervalSampler
    from trellis2.modules import sparse as sp

    dev = torch.device(args.device)
    caps = {}

    # ── coords from the SS stage reference latent ────────────────────────────
    with open(args.ss_dec + ".json") as f:
        ss_cfg = json.load(f)["args"]
    ss_cfg["use_fp16"] = False
    ss_dec = SparseStructureDecoder(**ss_cfg)
    ss_dec.load_state_dict({k: v.float() for k, v in load_file(args.ss_dec + ".safetensors").items()})
    ss_dec.dtype = torch.float32
    ss_dec.eval().float().to(dev)

    z_s = load_ss_sample_latent(args.ss_latent).to(dev)
    with torch.no_grad():
        occ = ss_dec(z_s) > 0                                   # [1,1,64,64,64]
    occ = (F.max_pool3d(occ.float(), 2, 2, 0) > 0.5)            # ss_res 32
    coords = torch.argwhere(occ)[:, [0, 2, 3, 4]].int().contiguous()
    L = coords.shape[0]
    print(f"coords: {L} active voxels at 32^3 "
          f"({100.0 * L / 32**3:.2f}% occupancy)")
    caps["coords"] = coords.float()
    del ss_dec

    # ── conditioning ─────────────────────────────────────────────────────────
    with open(args.dinodata, "rb") as f:
        assert f.read(8) == b"DINOCOND"
        _, _, ndim = struct.unpack("<III", f.read(12))
        shape = struct.unpack("<%dI" % ndim, f.read(4 * ndim))
        cond_np = np.frombuffer(f.read(), dtype="<f4").reshape(shape)
    cond = torch.from_numpy(cond_np.copy()).float().to(dev)
    neg_cond = torch.zeros_like(cond)

    # ── SLAT flow: forward parity point + full sampler ───────────────────────
    stem = os.path.join(args.models, "slat_flow_img2shape_dit_1_3B_512_bf16")
    with open(stem + ".json") as f:
        cfg = json.load(f)["args"]
    cfg.pop("initialization", None)
    cfg.pop("dtype", None)
    flow = SLatFlowModel(**cfg, dtype="float32")
    sd = {k: v.float() for k, v in load_file(stem + ".safetensors").items()}
    missing, unexpected = flow.load_state_dict(sd, strict=False)
    assert not unexpected, unexpected
    flow.convert_to(torch.float32)
    flow.eval().to(dev)

    g = torch.Generator().manual_seed(args.seed)
    noise = torch.randn(L, cfg["in_channels"], generator=g).to(dev)
    caps["slat_noise"] = noise

    x = sp.SparseTensor(feats=noise.clone(), coords=coords.to(dev))
    with torch.no_grad():
        out = flow(x, torch.tensor([args.t], device=dev), cond)
    caps["flow_t500_out"] = out.feats
    print(f"flow t={args.t}: out mean={out.feats.mean().item():.6f} "
          f"l2={out.feats.norm().item():.4f}")

    with open(args.pipeline_json) as f:
        pj = json.load(f)["args"]
    norm = pj["shape_slat_normalization"]
    sampler_params = pj["shape_slat_sampler"]["params"]
    print("sampler params:", sampler_params)

    sampler = FlowEulerGuidanceIntervalSampler(sigma_min=1e-5)
    x0 = sp.SparseTensor(feats=noise.clone(), coords=coords.to(dev))
    with torch.no_grad():
        slat = sampler.sample(flow, x0, cond=cond, neg_cond=neg_cond,
                              **sampler_params, verbose=True).samples
    mean = torch.tensor(norm["mean"], device=dev)[None]
    std = torch.tensor(norm["std"], device=dev)[None]
    slat = slat * std + mean
    caps["slat"] = slat.feats
    caps["slat_mean"] = mean[0]
    caps["slat_std"] = std[0]
    print(f"slat: mean={slat.feats.mean().item():.5f} std={slat.feats.std().item():.5f}")
    del flow
    if dev.type == "cuda":
        torch.cuda.empty_cache()

    # ── FDG decoder, level by level (mirrors SparseUnetVaeDecoder.forward) ──
    dstem = os.path.join(args.models, "shape_dec_next_dc_f16c32_fp16")
    with open(dstem + ".json") as f:
        dcfg = json.load(f)["args"]
    dcfg.pop("use_fp16", None)
    dcfg.pop("resolution", None)
    dec = FlexiDualGridVaeDecoder(resolution=args.resolution, use_fp16=False, **dcfg)
    dec.load_state_dict({k: v.float() for k, v in load_file(dstem + ".safetensors").items()})
    dec.eval().float().to(dev)

    with torch.no_grad():
        h = dec.from_latent(slat.float())
        for i, res in enumerate(dec.blocks):
            caps[f"lvl{i}.in_coords"] = h.coords.float()
            for j, block in enumerate(res):
                if i < len(dec.blocks) - 1 and j == len(res) - 1:
                    caps[f"lvl{i}.pre_up"] = h.feats
                    h, sub = block(h)
                    caps[f"lvl{i}.subdiv"] = sub.feats
                else:
                    h = block(h)
            print(f"level {i}: {h.feats.shape[0]} voxels x {h.feats.shape[1]} ch")
        hn = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
        out7 = dec.output_layer(hn)
    caps["out7"] = out7.feats
    caps["out_coords"] = out7.coords.float()
    print(f"out7: {out7.feats.shape}, offsets mean={torch.sigmoid(out7.feats[:, 0:3]).mean().item():.4f}, "
          f"intersected frac={(out7.feats[:, 3:6] > 0).float().mean().item():.4f}")

    import gguf
    writer = gguf.GGUFWriter(args.out, "reference")
    manifest = {"shapes": {}, "atol": 2e-3, "rtol": 2e-3}
    for name, t in caps.items():
        a = t.detach().cpu().float().numpy()
        manifest["shapes"][name] = list(a.shape)
        writer.add_tensor(name, np.ascontiguousarray(a.reshape(-1), dtype=np.float32))
    writer.write_header_to_file()
    writer.write_kv_data_to_file()
    writer.write_tensors_to_file()
    writer.close()
    with open(os.path.join(ref_common.DUMPS, "manifest_slat.json"), "w") as f:
        json.dump(manifest, f, indent=1)
    print(f"wrote {args.out} ({os.path.getsize(args.out):,} bytes), {len(caps)} tensors")


if __name__ == "__main__":
    main()